API Reference

Ordered the way the library is used: the two objects everything is built on (a causal graph and an SCM), then the agents and environments that act in them, then the causal machinery those agents call — intervention sets, identification, transport, discovery — and finally the decision, bounds and certificate layer that scores a policy. The — Complete Reference sections at the end carry an entry for every name in causalrl.__all__, grouped by causalrl.API_TIERS.

The full-lifecycle agent front doors are CausalMBRLAgent (fit → act on confounded logs) and OnlineCausalMBRL (learn the SCM while acting); both are in the complete-reference sections below.

Graphs

causalrl.scm.graph.CausalGraph

A causal graph: a DAG over observed variables plus bidirected edges that denote unobserved confounders (an ADMG).

Source code in src/causalrl/scm/graph.py

class CausalGraph:
    """A causal graph: a DAG over observed variables plus bidirected edges that
    denote unobserved confounders (an ADMG)."""

    def __init__(
        self,
        directed_edges: list[tuple[str, str]],
        bidirected_edges: list[tuple[str, str]] | None = None,
        nodes: list[str] | None = None,
        complete_parents: Mapping[str, str] | None = None,
    ) -> None:
        self._dag: nx.DiGraph[str] = nx.DiGraph()
        self._bi: nx.Graph[str] = nx.Graph()
        if nodes:
            self._dag.add_nodes_from(nodes)
            self._bi.add_nodes_from(nodes)
        self._dag.add_edges_from(directed_edges)
        for a, b in bidirected_edges or []:
            self._bi.add_node(a)
            self._bi.add_node(b)
            self._bi.add_edge(a, b)
        # keep node sets aligned
        self._bi.add_nodes_from(self._dag.nodes)
        self._dag.add_nodes_from(self._bi.nodes)
        if not nx.is_directed_acyclic_graph(self._dag):
            raise CausalGraphError("directed edges must form a DAG (cycle detected)")
        self._complete_parents: dict[str, str] = dict(complete_parents or {})
        for node in self._complete_parents:
            if node not in self._dag:
                raise CausalGraphError(f"unknown node in complete_parents: {node!r}")
            if self._bi.degree(node) > 0:
                partners = sorted(self._bi.neighbors(node))
                raise CausalGraphError(
                    f"{node!r} is declared to have complete parents but shares a bidirected edge "
                    f"with {partners}: a bidirected edge asserts an unobserved common cause, which "
                    f"IS an unobserved parent of {node!r}. The two statements cannot both hold. "
                    "Use CausalGraph.assert_complete_parents to drop those edges deliberately, "
                    "which records why."
                )

    @property
    def nodes(self) -> list[str]:
        return list(self._dag.nodes)

    @property
    def directed_edges(self) -> list[tuple[str, str]]:
        """The directed edges as ``(parent, child)`` pairs."""
        return [(u, v) for u, v in self._dag.edges]

    @property
    def bidirected_edges(self) -> list[tuple[str, str]]:
        """The bidirected (latent-confounding) edges."""
        return [(u, v) for u, v in self._bi.edges]

    def _check(self, node: str) -> None:
        if node not in self._dag:
            raise CausalGraphError(f"unknown node: {node!r}")

    def parents(self, node: str) -> list[str]:
        self._check(node)
        return list(self._dag.predecessors(node))

    def children(self, node: str) -> list[str]:
        self._check(node)
        return list(self._dag.successors(node))

    def topological_order(self) -> list[str]:
        return list(nx.topological_sort(self._dag))

    def is_confounded(self, a: str, b: str) -> bool:
        self._check(a)
        self._check(b)
        return self._bi.has_edge(a, b)

    def has_bidirected_edges(self) -> bool:
        """Whether this graph contains latent-confounding (bidirected) edges."""
        return self._bi.number_of_edges() > 0

    @property
    def complete_parents(self) -> dict[str, str]:
        """Nodes asserted to have NO unobserved parent, mapped to why.

        See :meth:`has_complete_parents`.
        """
        return dict(self._complete_parents)

    def has_complete_parents(self, node: str) -> bool:
        """Whether ``node``'s observed parents are asserted to be all the parents there are.

        This is the assertion that converts a bound into a point estimate, and it is the one thing
        a graph cannot learn from data: no observational test distinguishes "no unobserved parent"
        from "an unobserved parent I did not measure". It is licensed by *design* -- a randomised
        assignment, a feature flag, an experimenter's schedule, a rule-based pricing policy -- where
        the mechanism that sets the node is written down and contains nothing else.

        The library already had :class:`~causalrl.scm.fitters.PinnedMechanism`, which asserts the
        node's *equation*. That is a different and weaker claim: an equation can be supplied while
        the world still feeds the node something you did not model. This asserts the absence.
        """
        self._check(node)
        return node in self._complete_parents

    def assert_complete_parents(self, *nodes: str, reason: str) -> CausalGraph:
        """A copy of this graph in which ``nodes`` have no unobserved parents, and why.

        Bidirected edges incident to ``nodes`` are DROPPED, because they contradict the assertion:
        a latent common cause of ``A`` and ``Y`` is an unobserved parent of ``A``. This is what
        upgrades an effect from bounded to identified, so it is deliberately a distinct call with a
        mandatory ``reason`` rather than a constructor flag -- the assertion is a trust primitive,
        and an incorrect one silently converts an honest interval into a false point estimate.

        ``reason`` travels with the graph into every certificate built from it, so a reader can see
        which claims rest on it.
        """
        if not nodes:
            raise ValueError("assert_complete_parents needs at least one node")
        if not reason.strip():
            raise ValueError(
                "reason must be non-empty: this assertion is not checkable from data, so the "
                "record of why it is licensed is the only thing standing behind every claim it "
                "strengthens."
            )
        for node in nodes:
            self._check(node)
        targets = set(nodes)
        kept = [(a, b) for a, b in self.bidirected_edges if a not in targets and b not in targets]
        merged = dict(self._complete_parents)
        merged.update({node: reason for node in nodes})
        return CausalGraph(
            directed_edges=self.directed_edges,
            bidirected_edges=kept,
            nodes=self.nodes,
            complete_parents=merged,
        )

    def has_incident_bidirected_edges(self, node: str) -> bool:
        """Whether `node` is incident to any latent-confounding edge."""
        self._check(node)
        return self._bi.degree(node) > 0

    def c_components(self) -> list[set[str]]:
        """Connected components of the bidirected graph (isolated nodes are singletons)."""
        return [set(c) for c in nx.connected_components(self._bi)]

    def remove_incoming_edges(self, node: str) -> CausalGraph:
        """Return a copy with all directed edges into `node` removed (graph mutilation)."""
        self._check(node)
        directed = [(u, v) for u, v in self._dag.edges if v != node]
        bidirected = list(self._bi.edges)
        return CausalGraph(directed, bidirected, nodes=self.nodes)

    def _as_node_set(self, nodes: str | Iterable[str]) -> set[str]:
        ns = {nodes} if isinstance(nodes, str) else set(nodes)
        for n in ns:
            self._check(n)
        return ns

    def ancestors(self, nodes: str | Iterable[str]) -> set[str]:
        """Ancestors of `nodes` via directed edges, INCLUDING the inputs (the inclusive An(·)
        convention used by the identification/POMIS literature)."""
        ns = self._as_node_set(nodes)
        result = set(ns)
        for n in ns:
            result |= set(nx.ancestors(self._dag, n))  # type: ignore[reportUnknownMemberType]
        return result

    def descendants(self, nodes: str | Iterable[str]) -> set[str]:
        """Strict descendants of `nodes` via directed edges (excludes the inputs)."""
        ns = self._as_node_set(nodes)
        result: set[str] = set()
        for n in ns:
            result |= set(nx.descendants(self._dag, n))  # type: ignore[reportUnknownMemberType]
        return result

    def induced_subgraph(self, nodes: str | Iterable[str]) -> CausalGraph:
        """Subgraph on `nodes`: keep directed/bidirected edges with both endpoints in `nodes`."""
        keep = self._as_node_set(nodes)
        directed = [(u, v) for u, v in self._dag.edges if u in keep and v in keep]
        bidirected = [(u, v) for u, v in self._bi.edges if u in keep and v in keep]
        return CausalGraph(directed, bidirected, nodes=list(keep))

    def do_mutilate(self, intervened: str | Iterable[str]) -> CausalGraph:
        """ADMG mutilation for do(intervened): drop incoming directed edges to each
        intervened node AND every bidirected edge incident to an intervened node
        (intervention severs latent confounding into the set). Distinct from
        ``remove_incoming_edges``, which keeps bidirected edges."""
        x = self._as_node_set(intervened)
        directed = [(u, v) for u, v in self._dag.edges if v not in x]
        bidirected = [(u, v) for u, v in self._bi.edges if u not in x and v not in x]
        return CausalGraph(directed, bidirected, nodes=self.nodes)

    def _kept_reach(self, children: Iterable[str], keep: set[str]) -> set[str]:
        """Kept nodes reachable from `children` via directed edges whose interior nodes are
        all outside `keep` (a kept node is a boundary: it is reached but not expanded)."""
        reached: set[str] = set()
        seen: set[str] = set()
        frontier = list(children)
        while frontier:
            w = frontier.pop()
            if w in keep:
                reached.add(w)
            elif w not in seen:
                seen.add(w)
                frontier.extend(list(self._dag.successors(w)))
        return reached

    def latent_projection(self, keep: str | Iterable[str]) -> CausalGraph:
        """Latent projection onto `keep`: marginalize out every node not in `keep`, adding a
        directed edge for each directed path through removed nodes and a bidirected edge for
        each confounding path through removed nodes (the Tian-Pearl / Verma projection).

        Directed ``Vi -> Vj`` when a directed path from ``Vi`` to ``Vj`` has all interior
        nodes removed; bidirected ``Vi <-> Vj`` when a removed common cause (a marginalized
        node, or the latent behind a bidirected edge) reaches both through removed interiors.
        Removing a collider induces no confounding (its parents are never in its reached set)."""
        keep_set = self._as_node_set(keep)
        directed: list[tuple[str, str]] = []
        for vi in keep_set:
            for w in self._kept_reach(list(self._dag.successors(vi)), keep_set):
                directed.append((vi, w))

        bidirected: set[tuple[str, str]] = set()
        sources = [list(self._dag.successors(w)) for w in self._dag.nodes if w not in keep_set]
        sources += [[a, b] for a, b in self._bi.edges]
        for children in sources:
            reached = sorted(self._kept_reach(children, keep_set))
            for i in range(len(reached)):
                for j in range(i + 1, len(reached)):
                    bidirected.add((reached[i], reached[j]))

        return CausalGraph(directed, list(bidirected), nodes=list(keep_set))

directed_edges property

The directed edges as (parent, child) pairs.

bidirected_edges property

The bidirected (latent-confounding) edges.

complete_parents property

Nodes asserted to have NO unobserved parent, mapped to why.

See :meth:has_complete_parents.

has_bidirected_edges()

Whether this graph contains latent-confounding (bidirected) edges.

Source code in src/causalrl/scm/graph.py

def has_bidirected_edges(self) -> bool:
    """Whether this graph contains latent-confounding (bidirected) edges."""
    return self._bi.number_of_edges() > 0

has_complete_parents(node)

Whether node's observed parents are asserted to be all the parents there are.

This is the assertion that converts a bound into a point estimate, and it is the one thing a graph cannot learn from data: no observational test distinguishes "no unobserved parent" from "an unobserved parent I did not measure". It is licensed by design -- a randomised assignment, a feature flag, an experimenter's schedule, a rule-based pricing policy -- where the mechanism that sets the node is written down and contains nothing else.

The library already had :class:~causalrl.scm.fitters.PinnedMechanism, which asserts the node's equation. That is a different and weaker claim: an equation can be supplied while the world still feeds the node something you did not model. This asserts the absence.

Source code in src/causalrl/scm/graph.py

def has_complete_parents(self, node: str) -> bool:
    """Whether ``node``'s observed parents are asserted to be all the parents there are.

    This is the assertion that converts a bound into a point estimate, and it is the one thing
    a graph cannot learn from data: no observational test distinguishes "no unobserved parent"
    from "an unobserved parent I did not measure". It is licensed by *design* -- a randomised
    assignment, a feature flag, an experimenter's schedule, a rule-based pricing policy -- where
    the mechanism that sets the node is written down and contains nothing else.

    The library already had :class:`~causalrl.scm.fitters.PinnedMechanism`, which asserts the
    node's *equation*. That is a different and weaker claim: an equation can be supplied while
    the world still feeds the node something you did not model. This asserts the absence.
    """
    self._check(node)
    return node in self._complete_parents

assert_complete_parents(*nodes, reason)

A copy of this graph in which nodes have no unobserved parents, and why.

Bidirected edges incident to nodes are DROPPED, because they contradict the assertion: a latent common cause of A and Y is an unobserved parent of A. This is what upgrades an effect from bounded to identified, so it is deliberately a distinct call with a mandatory reason rather than a constructor flag -- the assertion is a trust primitive, and an incorrect one silently converts an honest interval into a false point estimate.

reason travels with the graph into every certificate built from it, so a reader can see which claims rest on it.

Source code in src/causalrl/scm/graph.py

def assert_complete_parents(self, *nodes: str, reason: str) -> CausalGraph:
    """A copy of this graph in which ``nodes`` have no unobserved parents, and why.

    Bidirected edges incident to ``nodes`` are DROPPED, because they contradict the assertion:
    a latent common cause of ``A`` and ``Y`` is an unobserved parent of ``A``. This is what
    upgrades an effect from bounded to identified, so it is deliberately a distinct call with a
    mandatory ``reason`` rather than a constructor flag -- the assertion is a trust primitive,
    and an incorrect one silently converts an honest interval into a false point estimate.

    ``reason`` travels with the graph into every certificate built from it, so a reader can see
    which claims rest on it.
    """
    if not nodes:
        raise ValueError("assert_complete_parents needs at least one node")
    if not reason.strip():
        raise ValueError(
            "reason must be non-empty: this assertion is not checkable from data, so the "
            "record of why it is licensed is the only thing standing behind every claim it "
            "strengthens."
        )
    for node in nodes:
        self._check(node)
    targets = set(nodes)
    kept = [(a, b) for a, b in self.bidirected_edges if a not in targets and b not in targets]
    merged = dict(self._complete_parents)
    merged.update({node: reason for node in nodes})
    return CausalGraph(
        directed_edges=self.directed_edges,
        bidirected_edges=kept,
        nodes=self.nodes,
        complete_parents=merged,
    )

has_incident_bidirected_edges(node)

Whether node is incident to any latent-confounding edge.

Source code in src/causalrl/scm/graph.py

def has_incident_bidirected_edges(self, node: str) -> bool:
    """Whether `node` is incident to any latent-confounding edge."""
    self._check(node)
    return self._bi.degree(node) > 0

c_components()

Connected components of the bidirected graph (isolated nodes are singletons).

Source code in src/causalrl/scm/graph.py

def c_components(self) -> list[set[str]]:
    """Connected components of the bidirected graph (isolated nodes are singletons)."""
    return [set(c) for c in nx.connected_components(self._bi)]

remove_incoming_edges(node)

Return a copy with all directed edges into node removed (graph mutilation).

Source code in src/causalrl/scm/graph.py

def remove_incoming_edges(self, node: str) -> CausalGraph:
    """Return a copy with all directed edges into `node` removed (graph mutilation)."""
    self._check(node)
    directed = [(u, v) for u, v in self._dag.edges if v != node]
    bidirected = list(self._bi.edges)
    return CausalGraph(directed, bidirected, nodes=self.nodes)

ancestors(nodes)

Ancestors of nodes via directed edges, INCLUDING the inputs (the inclusive An(·) convention used by the identification/POMIS literature).

Source code in src/causalrl/scm/graph.py

def ancestors(self, nodes: str | Iterable[str]) -> set[str]:
    """Ancestors of `nodes` via directed edges, INCLUDING the inputs (the inclusive An(·)
    convention used by the identification/POMIS literature)."""
    ns = self._as_node_set(nodes)
    result = set(ns)
    for n in ns:
        result |= set(nx.ancestors(self._dag, n))  # type: ignore[reportUnknownMemberType]
    return result

descendants(nodes)

Strict descendants of nodes via directed edges (excludes the inputs).

Source code in src/causalrl/scm/graph.py

def descendants(self, nodes: str | Iterable[str]) -> set[str]:
    """Strict descendants of `nodes` via directed edges (excludes the inputs)."""
    ns = self._as_node_set(nodes)
    result: set[str] = set()
    for n in ns:
        result |= set(nx.descendants(self._dag, n))  # type: ignore[reportUnknownMemberType]
    return result

induced_subgraph(nodes)

Subgraph on nodes: keep directed/bidirected edges with both endpoints in nodes.

Source code in src/causalrl/scm/graph.py

def induced_subgraph(self, nodes: str | Iterable[str]) -> CausalGraph:
    """Subgraph on `nodes`: keep directed/bidirected edges with both endpoints in `nodes`."""
    keep = self._as_node_set(nodes)
    directed = [(u, v) for u, v in self._dag.edges if u in keep and v in keep]
    bidirected = [(u, v) for u, v in self._bi.edges if u in keep and v in keep]
    return CausalGraph(directed, bidirected, nodes=list(keep))

do_mutilate(intervened)

ADMG mutilation for do(intervened): drop incoming directed edges to each intervened node AND every bidirected edge incident to an intervened node (intervention severs latent confounding into the set). Distinct from remove_incoming_edges, which keeps bidirected edges.

Source code in src/causalrl/scm/graph.py

def do_mutilate(self, intervened: str | Iterable[str]) -> CausalGraph:
    """ADMG mutilation for do(intervened): drop incoming directed edges to each
    intervened node AND every bidirected edge incident to an intervened node
    (intervention severs latent confounding into the set). Distinct from
    ``remove_incoming_edges``, which keeps bidirected edges."""
    x = self._as_node_set(intervened)
    directed = [(u, v) for u, v in self._dag.edges if v not in x]
    bidirected = [(u, v) for u, v in self._bi.edges if u not in x and v not in x]
    return CausalGraph(directed, bidirected, nodes=self.nodes)

latent_projection(keep)

Latent projection onto keep: marginalize out every node not in keep, adding a directed edge for each directed path through removed nodes and a bidirected edge for each confounding path through removed nodes (the Tian-Pearl / Verma projection).

Directed Vi -> Vj when a directed path from Vi to Vj has all interior nodes removed; bidirected Vi <-> Vj when a removed common cause (a marginalized node, or the latent behind a bidirected edge) reaches both through removed interiors. Removing a collider induces no confounding (its parents are never in its reached set).

Source code in src/causalrl/scm/graph.py

def latent_projection(self, keep: str | Iterable[str]) -> CausalGraph:
    """Latent projection onto `keep`: marginalize out every node not in `keep`, adding a
    directed edge for each directed path through removed nodes and a bidirected edge for
    each confounding path through removed nodes (the Tian-Pearl / Verma projection).

    Directed ``Vi -> Vj`` when a directed path from ``Vi`` to ``Vj`` has all interior
    nodes removed; bidirected ``Vi <-> Vj`` when a removed common cause (a marginalized
    node, or the latent behind a bidirected edge) reaches both through removed interiors.
    Removing a collider induces no confounding (its parents are never in its reached set)."""
    keep_set = self._as_node_set(keep)
    directed: list[tuple[str, str]] = []
    for vi in keep_set:
        for w in self._kept_reach(list(self._dag.successors(vi)), keep_set):
            directed.append((vi, w))

    bidirected: set[tuple[str, str]] = set()
    sources = [list(self._dag.successors(w)) for w in self._dag.nodes if w not in keep_set]
    sources += [[a, b] for a, b in self._bi.edges]
    for children in sources:
        reached = sorted(self._kept_reach(children, keep_set))
        for i in range(len(reached)):
            for j in range(i + 1, len(reached)):
                bidirected.add((reached[i], reached[j]))

    return CausalGraph(directed, list(bidirected), nodes=list(keep_set))

Structural Causal Models

causalrl.scm.scm.StructuralCausalModel

Executable explicit-latent DAG SCM supporting L1/L2/L3 queries.

Bidirected-edge ADMGs are accepted by :class:CausalGraph for analytical graph algorithms, but they are not executable SCMs: shared latent causes must be represented as explicit parent nodes with their own mechanism and exogenous distribution.

Source code in src/causalrl/scm/scm.py

class StructuralCausalModel:
    """Executable explicit-latent DAG SCM supporting L1/L2/L3 queries.

    Bidirected-edge ADMGs are accepted by :class:`CausalGraph` for analytical graph
    algorithms, but they are not executable SCMs: shared latent causes must be represented
    as explicit parent nodes with their own mechanism and exogenous distribution.
    """

    def __init__(
        self,
        graph: CausalGraph,
        mechanisms: dict[str, Mechanism],
        exogenous: dict[str, Distribution],
        *,
        provenance: Literal["specified", "fitted", "mixed"] = "specified",
        fit_report: FitReport | None = None,
    ) -> None:
        if graph.has_bidirected_edges():
            raise CausalGraphError(
                "StructuralCausalModel requires an explicit latent-variable DAG; "
                "bidirected ADMG edges are analytical only. Represent shared latent causes "
                "as explicit latent nodes."
            )
        nodes = set(graph.nodes)
        for name, entries in (("mechanisms", mechanisms), ("exogenous distributions", exogenous)):
            missing = nodes - set(entries)
            extra = set(entries) - nodes
            if missing or extra:
                raise CausalGraphError(
                    f"{name} must exactly cover graph nodes; "
                    f"missing={sorted(missing)}, extra={sorted(extra)}"
                )
        for node in graph.nodes:
            expected = set(graph.parents(node))
            declared = set(mechanisms[node].parents)
            if declared != expected:
                raise CausalGraphError(
                    f"mechanism parents for {node!r} do not match graph parents; "
                    f"expected={sorted(expected)}, declared={sorted(declared)}"
                )
        self.graph = graph
        self.mechanisms = mechanisms
        self.exogenous = exogenous
        # Explicit annotation: without it pyright infers the attribute's declared type as the
        # widened `str` (not the parameter's narrower Literal), so reading it back and passing
        # it into another StructuralCausalModel(...) call -- e.g. from do() -- fails strict
        # checking even though the value only ever holds one of the three literal members.
        self.provenance: Literal["specified", "fitted", "mixed"] = provenance
        self.fit_report = fit_report
        self._generator = torch.Generator()  # type: ignore[reportPrivateImportUsage]
        self._generator.seed()

    def _sample_exogenous(self, n: int, seed: int | None) -> dict[str, Tensor]:
        """Sample using a private CPU RNG stream without mutating process-global Torch state."""
        generator = torch.Generator()  # type: ignore[reportPrivateImportUsage]
        if seed is None:
            generator.set_state(self._generator.get_state())
        else:
            generator.manual_seed(seed)
        with torch.random.fork_rng():  # type: ignore[reportUnknownMemberType]
            torch.random.set_rng_state(generator.get_state())  # type: ignore[reportUnknownMemberType]
            out = {
                name: dist.sample((n,)).reshape(n).float() for name, dist in self.exogenous.items()
            }
            if seed is None:
                self._generator.set_state(torch.random.get_rng_state())  # type: ignore[reportUnknownMemberType]
        return out

    def _evaluate(self, noise: dict[str, Tensor]) -> dict[str, Tensor]:
        values: dict[str, Tensor] = {}
        for node in self.graph.topological_order():
            mech = self.mechanisms[node]
            parent_values = {p: values[p] for p in self.graph.parents(node)}
            values[node] = mech(parent_values, noise[node])
        return values

    def do(self, interventions: Mapping[str, Value]) -> StructuralCausalModel:
        """Layer 2: return the mutilated SCM under do(interventions). Original is unchanged.

        Preserves ``provenance``/``fit_report``: the un-intervened mechanisms are the same
        objects (shallow-copied dict), so a fitted model's L3 guard must still see them as
        fitted after mutilation -- otherwise ``abduct`` on the result would silently stop
        refusing point counterfactuals for the non-invertible nodes ``do`` never touched.
        """
        graph = self.graph
        mechanisms = dict(self.mechanisms)
        for node, value in interventions.items():
            if node not in self.mechanisms:
                raise CausalGraphError(f"cannot intervene on unknown node: {node!r}")
            graph = graph.remove_incoming_edges(node)
            if isinstance(value, (int, float, bool)):
                const = float(value)
                mechanisms[node] = FunctionalMechanism(
                    [],
                    lambda pa, u, _c=const: torch.full_like(u, _c),  # type: ignore[reportPrivateImportUsage]
                )
            else:
                # Per-sample vector: pin elementwise, broadcast-checked against the unit count.
                mechanisms[node] = FunctionalMechanism(
                    [],
                    lambda pa, u, _v=value: _broadcast_value(_v, u.shape[0]).to(u.dtype),
                )
        return StructuralCausalModel(
            graph,
            mechanisms,
            self.exogenous,
            provenance=self.provenance,
            fit_report=self.fit_report,
        )

    def non_invertible_nodes(self) -> list[str]:
        """Nodes whose fitted mechanism cannot recover its noise from (parents, value).

        Mechanisms without an ``invertible`` attribute are treated as invertible: a hand-written
        mechanism is an assertion by its author, so its coupling is given rather than inferred.
        """
        return [
            node for node, mech in self.mechanisms.items() if not getattr(mech, "invertible", True)
        ]

    def abduct(
        self,
        evidence: dict[str, float] | None = None,
        *,
        known: Mapping[str, Value] | None = None,
        n: int = 20_000,
        seed: int | None = None,
        atol: float = 1e-6,
    ) -> ExogenousPosterior:
        """Layer 3, step 1 — infer the exogenous posterior given evidence/known noise.

        ``known`` pins supplied exogenous values *exactly* (the exact, continuous path: no
        rejection). Remaining exogenous are sampled; if ``evidence`` is given they are
        rejection-filtered so the factual evaluation matches ``evidence`` within ``atol``.
        Returns an :class:`ExogenousPosterior`; call its ``predict(do=...)``.

        On a **fitted** SCM containing a non-invertible node this raises: L1 data identifies the
        mechanisms but not the noise-to-value coupling, so a point counterfactual would report a
        modelling choice as a result. Use :func:`causalrl.counterfactual_interval` instead. A
        ``"mixed"`` model (some nodes pinned, some learned) is gated the same way: pinning one
        equation does not identify the coupling at a node that was still learned from data.
        """
        if self.provenance in ("fitted", "mixed"):
            ambiguous = sorted(self.non_invertible_nodes())
            if ambiguous:
                raise NotIdentifiableError(
                    f"counterfactuals on a {self.provenance} SCM are not identified: node(s) "
                    f"{ambiguous} have a non-invertible mechanism, so their "
                    "noise-to-value coupling is a modelling choice, not a fitted quantity. "
                    "Use counterfactual_interval(...) for the identified interval.",
                    witness=ambiguous,
                )
        return self._abduct(evidence, known=known, n=n, seed=seed, atol=atol)

    def _abduct(
        self,
        evidence: dict[str, float] | None = None,
        *,
        known: Mapping[str, Value] | None = None,
        n: int = 20_000,
        seed: int | None = None,
        atol: float = 1e-6,
    ) -> ExogenousPosterior:
        """Unguarded abduction, for a caller that has already established what its query licenses
        — sub-project 4's counterfactual data augmentation replays *invertible* mechanisms on a
        fitted model and must not be blocked by :meth:`abduct`'s provenance guard."""
        known = known or {}
        bad = set(known) - set(self.exogenous)
        if bad:
            raise CausalGraphError(f"unknown exogenous node(s): {sorted(bad)}")
        # Exact path: all-known exogenous, no evidence to reject against.
        if known and not evidence:
            noise = {
                name: (
                    _broadcast_value(known[name], n)
                    if name in known
                    else dist.sample((n,)).reshape(n).float()
                )
                for name, dist in self.exogenous.items()
            }
            return ExogenousPosterior(self, noise)
        # Evidence path: sample, pin any known, reject to match evidence.
        sampled = self._sample_exogenous(n, seed)
        for name, val in known.items():
            sampled[name] = _broadcast_value(val, n)
        factual = self._evaluate(sampled)
        mask = torch.ones(n, dtype=torch.bool)  # type: ignore[reportPrivateImportUsage]
        for node, val in (evidence or {}).items():
            if node not in self.mechanisms:
                raise CausalGraphError(f"unknown evidence node: {node!r}")
            mask &= (factual[node] - float(val)).abs() <= atol
        if int(mask.sum()) == 0:
            raise RealizabilityError(
                f"no exogenous draws match evidence {evidence!r}; "
                "increase n or check that the evidence has nonzero probability"
            )
        return ExogenousPosterior(self, {name: u[mask] for name, u in sampled.items()})

    def counterfactual(
        self,
        evidence: dict[str, float],
        interventions: Mapping[str, Value],
        n: int,
        *,
        seed: int | None = None,
        atol: float = 1e-6,
    ) -> dict[str, Tensor]:
        """Layer 3: abduction-action-prediction. Sugar over :meth:`abduct` + predict."""
        post = self.abduct(evidence, n=n, seed=seed, atol=atol)
        return post.predict(do=interventions or None)

    def see(self, n: int, *, seed: int | None = None) -> dict[str, Tensor]:
        """Layer 1: draw n observational samples P(V)."""
        noise = self._sample_exogenous(n, seed)
        return self._evaluate(noise)

do(interventions)

Layer 2: return the mutilated SCM under do(interventions). Original is unchanged.

Preserves provenance/fit_report: the un-intervened mechanisms are the same objects (shallow-copied dict), so a fitted model's L3 guard must still see them as fitted after mutilation -- otherwise abduct on the result would silently stop refusing point counterfactuals for the non-invertible nodes do never touched.

Source code in src/causalrl/scm/scm.py

def do(self, interventions: Mapping[str, Value]) -> StructuralCausalModel:
    """Layer 2: return the mutilated SCM under do(interventions). Original is unchanged.

    Preserves ``provenance``/``fit_report``: the un-intervened mechanisms are the same
    objects (shallow-copied dict), so a fitted model's L3 guard must still see them as
    fitted after mutilation -- otherwise ``abduct`` on the result would silently stop
    refusing point counterfactuals for the non-invertible nodes ``do`` never touched.
    """
    graph = self.graph
    mechanisms = dict(self.mechanisms)
    for node, value in interventions.items():
        if node not in self.mechanisms:
            raise CausalGraphError(f"cannot intervene on unknown node: {node!r}")
        graph = graph.remove_incoming_edges(node)
        if isinstance(value, (int, float, bool)):
            const = float(value)
            mechanisms[node] = FunctionalMechanism(
                [],
                lambda pa, u, _c=const: torch.full_like(u, _c),  # type: ignore[reportPrivateImportUsage]
            )
        else:
            # Per-sample vector: pin elementwise, broadcast-checked against the unit count.
            mechanisms[node] = FunctionalMechanism(
                [],
                lambda pa, u, _v=value: _broadcast_value(_v, u.shape[0]).to(u.dtype),
            )
    return StructuralCausalModel(
        graph,
        mechanisms,
        self.exogenous,
        provenance=self.provenance,
        fit_report=self.fit_report,
    )

non_invertible_nodes()

Nodes whose fitted mechanism cannot recover its noise from (parents, value).

Mechanisms without an invertible attribute are treated as invertible: a hand-written mechanism is an assertion by its author, so its coupling is given rather than inferred.

Source code in src/causalrl/scm/scm.py

def non_invertible_nodes(self) -> list[str]:
    """Nodes whose fitted mechanism cannot recover its noise from (parents, value).

    Mechanisms without an ``invertible`` attribute are treated as invertible: a hand-written
    mechanism is an assertion by its author, so its coupling is given rather than inferred.
    """
    return [
        node for node, mech in self.mechanisms.items() if not getattr(mech, "invertible", True)
    ]

abduct(evidence=None, *, known=None, n=20000, seed=None, atol=1e-06)

Layer 3, step 1 — infer the exogenous posterior given evidence/known noise.

known pins supplied exogenous values exactly (the exact, continuous path: no rejection). Remaining exogenous are sampled; if evidence is given they are rejection-filtered so the factual evaluation matches evidence within atol. Returns an :class:ExogenousPosterior; call its predict(do=...).

On a fitted SCM containing a non-invertible node this raises: L1 data identifies the mechanisms but not the noise-to-value coupling, so a point counterfactual would report a modelling choice as a result. Use :func:causalrl.counterfactual_interval instead. A "mixed" model (some nodes pinned, some learned) is gated the same way: pinning one equation does not identify the coupling at a node that was still learned from data.

Source code in src/causalrl/scm/scm.py

def abduct(
    self,
    evidence: dict[str, float] | None = None,
    *,
    known: Mapping[str, Value] | None = None,
    n: int = 20_000,
    seed: int | None = None,
    atol: float = 1e-6,
) -> ExogenousPosterior:
    """Layer 3, step 1 — infer the exogenous posterior given evidence/known noise.

    ``known`` pins supplied exogenous values *exactly* (the exact, continuous path: no
    rejection). Remaining exogenous are sampled; if ``evidence`` is given they are
    rejection-filtered so the factual evaluation matches ``evidence`` within ``atol``.
    Returns an :class:`ExogenousPosterior`; call its ``predict(do=...)``.

    On a **fitted** SCM containing a non-invertible node this raises: L1 data identifies the
    mechanisms but not the noise-to-value coupling, so a point counterfactual would report a
    modelling choice as a result. Use :func:`causalrl.counterfactual_interval` instead. A
    ``"mixed"`` model (some nodes pinned, some learned) is gated the same way: pinning one
    equation does not identify the coupling at a node that was still learned from data.
    """
    if self.provenance in ("fitted", "mixed"):
        ambiguous = sorted(self.non_invertible_nodes())
        if ambiguous:
            raise NotIdentifiableError(
                f"counterfactuals on a {self.provenance} SCM are not identified: node(s) "
                f"{ambiguous} have a non-invertible mechanism, so their "
                "noise-to-value coupling is a modelling choice, not a fitted quantity. "
                "Use counterfactual_interval(...) for the identified interval.",
                witness=ambiguous,
            )
    return self._abduct(evidence, known=known, n=n, seed=seed, atol=atol)

counterfactual(evidence, interventions, n, *, seed=None, atol=1e-06)

Layer 3: abduction-action-prediction. Sugar over :meth:abduct + predict.

Source code in src/causalrl/scm/scm.py

def counterfactual(
    self,
    evidence: dict[str, float],
    interventions: Mapping[str, Value],
    n: int,
    *,
    seed: int | None = None,
    atol: float = 1e-6,
) -> dict[str, Tensor]:
    """Layer 3: abduction-action-prediction. Sugar over :meth:`abduct` + predict."""
    post = self.abduct(evidence, n=n, seed=seed, atol=atol)
    return post.predict(do=interventions or None)

see(n, *, seed=None)

Layer 1: draw n observational samples P(V).

Source code in src/causalrl/scm/scm.py

def see(self, n: int, *, seed: int | None = None) -> dict[str, Tensor]:
    """Layer 1: draw n observational samples P(V)."""
    noise = self._sample_exogenous(n, seed)
    return self._evaluate(noise)

Assumption-Aware Agents

causalrl.agents.dovi.DOVI

Bases: Agent

Deconfounded Optimistic Value Iteration (Wang, Yang, Wang, Bareinboim 2021), tabular.

Finite-horizon backward induction (LSVI form) whose optimism is capped by the Manski upper causal bound on each (state, action)'s immediate reward, deconfounding the value estimate. With horizon H::

r̃(s,a)   = min( mean_online(s,a) + ucb_bonus(s,a),  hi_Manski(s,a) )
Q_h(s,a) = r̃(s,a) + Σ_s' P̂(s'|s,a) · V_{h+1}(s')
V_h(s)   = max_a Q_h(s,a),    V_{H+1} ≡ 0,    h = H, H-1, …, 1

P̂ is the empirical next-state distribution pooled from offline logs and online steps; transitions that ended an episode (done) carry zero future value (the absorbing terminal). At H = 1 the backup reduces exactly to v0.2's immediate Manski ceiling, so the horizon-1 DTR result is preserved.

Bound validity: the backup is a certified optimistic bound on the return only when the dynamics do not depend on the hidden confounder (true of SequentialDTREnv, whose transitions are U-independent). On ConfoundedGridworld the behavior policy steers around the hidden hazard, so P̂ is confounded and the backup is heuristic value propagation, not a certified bound.

Source code in src/causalrl/agents/dovi.py

class DOVI(Agent):
    """Deconfounded Optimistic Value Iteration (Wang, Yang, Wang, Bareinboim 2021), tabular.

    Finite-horizon backward induction (LSVI form) whose optimism is capped by the Manski
    upper causal bound on each (state, action)'s immediate reward, deconfounding the value
    estimate. With horizon ``H``::

        r̃(s,a)   = min( mean_online(s,a) + ucb_bonus(s,a),  hi_Manski(s,a) )
        Q_h(s,a) = r̃(s,a) + Σ_s' P̂(s'|s,a) · V_{h+1}(s')
        V_h(s)   = max_a Q_h(s,a),    V_{H+1} ≡ 0,    h = H, H-1, …, 1

    ``P̂`` is the empirical next-state distribution pooled from offline logs and online steps;
    transitions that ended an episode (``done``) carry zero future value (the absorbing
    terminal). At ``H = 1`` the backup reduces exactly to v0.2's immediate Manski ceiling, so
    the horizon-1 DTR result is preserved.

    Bound validity: the backup is a certified optimistic bound on the return only when the
    dynamics do not depend on the hidden confounder (true of ``SequentialDTREnv``, whose
    transitions are U-independent). On ``ConfoundedGridworld`` the behavior policy steers
    around the hidden hazard, so ``P̂`` is confounded and the backup is heuristic value
    propagation, not a certified bound.
    """

    def __init__(
        self,
        n_states: int,
        n_actions: int,
        horizon: int,
        seed: int | None = None,
        *,
        reward_max: float = 1.0,
        transition_assumption: TransitionAssumption = "unknown",
        allow_heuristic: bool = False,
    ) -> None:
        if transition_assumption not in ("unknown", "unconfounded"):
            raise ValueError(
                "transition_assumption must be 'unknown' or 'unconfounded', "
                f"got {transition_assumption!r}"
            )
        self.n_states = n_states
        self.n_actions = n_actions
        self.horizon = horizon
        self.reward_max = reward_max
        self.transition_assumption = transition_assumption
        self.allow_heuristic = allow_heuristic
        self._counts = np.zeros((n_states, n_actions))
        self._sums = np.zeros((n_states, n_actions))
        # Transition model: non-terminal next-state counts, and total transitions seen.
        self._trans = np.zeros((n_states, n_actions, n_states))
        self._trans_n = np.zeros((n_states, n_actions))
        self._t = 1
        self._rng = np.random.default_rng(seed)
        self._ceiling = np.ones((n_states, n_actions))  # Manski upper bound on immediate reward
        self._q: np.ndarray | None = None  # cached plan, shape (H+1, S, A), indexed h = 1..H

    @property
    def is_certified(self) -> bool:
        """Whether value propagation is within the documented causal guarantee."""
        return self.horizon <= 1 or self.transition_assumption == "unconfounded"

    def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
        if not self.is_certified and not self.allow_heuristic:
            raise UnverifiedAssumptionError(
                "multi-stage DOVI requires transition_assumption='unconfounded' for "
                "certified value propagation; set allow_heuristic=True to run without "
                "that transition guarantee"
            )
        for (s, a), (_lo, hi) in bounds_table(dataset).items():
            self._ceiling[s, a] = hi
        for tr in dataset.transitions:
            self._trans_n[tr.state, tr.action] += 1.0
            if not tr.done:
                self._trans[tr.state, tr.action, tr.next_state] += 1.0
        self._q = None

    def _r_tilde_table(self) -> np.ndarray:
        n = self._counts
        safe_n = np.maximum(n, 1.0)
        mean = np.where(n > 0, self._sums / safe_n, 0.0)
        bonus = np.where(
            n > 0,
            self.reward_max * np.sqrt(2.0 * math.log(self._t) / safe_n),
            self.reward_max,
        )
        return np.minimum(mean + bonus, self._ceiling)

    def _plan(self) -> np.ndarray:
        h_max, n_states, n_actions = self.horizon, self.n_states, self.n_actions
        rtil = self._r_tilde_table()
        denom = np.maximum(self._trans_n, 1.0)  # (S, A)
        q = np.zeros((h_max + 1, n_states, n_actions))
        v_next = np.zeros(n_states)  # V_{H+1}
        for h in range(h_max, 0, -1):
            # boot[s,a] = (Σ_s' trans[s,a,s'] · V_next[s']) / total transitions seen at (s,a)
            boot = (self._trans @ v_next) / denom
            q[h] = rtil + boot
            v_next = q[h].max(axis=1)
        return q

    def optimistic_q(self, state: int, action: int, h: int = 1) -> float:
        if self._q is None:
            self._q = self._plan()
        return float(self._q[h, state, action])

    def act(self, observation: dict[str, Any]) -> int:
        s = int(observation["state"])
        h = min(int(observation.get("t", 0)) + 1, self.horizon)
        if self._q is None:
            self._q = self._plan()
        scores = self._q[h, s]
        best = float(scores.max())
        winners = [a for a in range(self.n_actions) if scores[a] >= best - 1e-12]
        return int(self._rng.choice(winners))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        s = int(observation["state"])
        self._counts[s, action] += 1.0
        self._sums[s, action] += reward
        self._t += 1
        self._q = None

    def observe_transition(self, state: int, action: int, next_state: int, done: bool) -> None:
        self._trans_n[state, action] += 1.0
        if not done:
            self._trans[state, action, next_state] += 1.0
        self._q = None

is_certified property

Whether value propagation is within the documented causal guarantee.

causalrl.agents.scbandit.POMISThompsonSampling

Bases: _ArmSubsetThompsonSampling

Thompson sampling over only the arms whose intervened-variable set is a POMIS.

manipulable (the variables you may intervene on) must be given explicitly; non-manipulable variables are handled by the POMIS engine via latent projection (r40). When every non-reward variable is manipulable this matches the unconstrained POMIS.

Source code in src/causalrl/agents/scbandit.py

class POMISThompsonSampling(_ArmSubsetThompsonSampling):
    """Thompson sampling over only the arms whose intervened-variable set is a POMIS.

    ``manipulable`` (the variables you may intervene on) must be given explicitly; non-manipulable
    variables are handled by the POMIS engine via latent projection (r40). When every non-reward
    variable is manipulable this matches the unconstrained POMIS.
    """

    def __init__(
        self,
        graph: CausalGraph,
        reward: str,
        arms: list[dict[str, int]],
        seed: int | None = None,
        *,
        manipulable: Iterable[str] | None = None,
    ) -> None:
        arm_variables = {v for arm in arms for v in arm}
        if manipulable is None:
            raise ValueError(
                "POMISThompsonSampling requires manipulable=<intervenable variables>; "
                "inferring it from the arms is no longer supported"
            )
        permitted = set(manipulable)
        unexpected = arm_variables - permitted
        if unexpected:
            raise ValueError(f"arms intervene on non-manipulable variables: {sorted(unexpected)}")
        pomis_sets = set(pomis(graph, reward, manipulable=permitted))
        allowed = [i for i, arm in enumerate(arms) if frozenset(arm.keys()) in pomis_sets]
        super().__init__(allowed, seed)

Causal Gymnasium Wrapper

causalrl.envs.wrapper.CausalEnvWrapper

Bases: Wrapper[Any, Any, Any, Any]

A Gymnasium wrapper that exposes the wrapped environment's causal structure.

Parameters

env: Any gymnasium.Env. When the env carries a non-None .scm attribute and reward_node is provided (and valid), the full causal interface is enabled. When env.scm is None or reward_node is None, construction still succeeds but the causal interface is disabled (pass-through mode). reward_node: The SCM variable name that represents the reward / return signal. Optional — None disables the causal interface. If the env has a live SCM and reward_node is supplied but not present in the graph, a ValueError is raised.

Attributes

reward_node: The SCM node treated as the reward/return, or None when not set. has_causal_interface: True iff the wrapped env exposes a non-None SCM and reward_node is set and present in the graph. All causal methods require this to be True. reward_parents: Direct SCM parents of reward_node in topological order. Requires has_causal_interface. scm: The underlying :class:~causalrl.scm.scm.StructuralCausalModel (live reference to the wrapped env's .scm), or None when not available. active_interventions: The currently stored intervention mapping, or None if no persistent intervention is active.

Source code in src/causalrl/envs/wrapper.py

class CausalEnvWrapper(gym.Wrapper[Any, Any, Any, Any]):
    """A Gymnasium wrapper that exposes the wrapped environment's causal structure.

    Parameters
    ----------
    env:
        Any ``gymnasium.Env``.  When the env carries a non-``None`` ``.scm`` attribute and
        ``reward_node`` is provided (and valid), the full causal interface is enabled.
        When ``env.scm is None`` **or** ``reward_node`` is ``None``, construction still
        succeeds but the causal interface is disabled (pass-through mode).
    reward_node:
        The SCM variable name that represents the reward / return signal.  Optional —
        ``None`` disables the causal interface.  If the env has a live SCM and
        ``reward_node`` is supplied but not present in the graph, a ``ValueError`` is raised.

    Attributes
    ----------
    reward_node:
        The SCM node treated as the reward/return, or ``None`` when not set.
    has_causal_interface:
        ``True`` iff the wrapped env exposes a non-``None`` SCM *and* ``reward_node`` is set
        and present in the graph.  All causal methods require this to be ``True``.
    reward_parents:
        Direct SCM parents of ``reward_node`` in topological order.  Requires
        ``has_causal_interface``.
    scm:
        The underlying :class:`~causalrl.scm.scm.StructuralCausalModel` (live reference to
        the wrapped env's ``.scm``), or ``None`` when not available.
    active_interventions:
        The currently stored intervention mapping, or ``None`` if no persistent
        intervention is active.
    """

    def __init__(
        self,
        env: gym.Env[Any, Any],
        *,
        reward_node: str | None = None,
    ) -> None:
        super().__init__(env)
        # Validate reward_node only when both the SCM and reward_node are provided.
        raw_scm = getattr(env, "scm", None)
        if reward_node is not None and raw_scm is not None:
            scm = cast("StructuralCausalModel", raw_scm)
            if reward_node not in scm.graph.nodes:
                raise ValueError(
                    f"reward_node {reward_node!r} is not a node in the SCM graph "
                    f"(nodes: {sorted(scm.graph.nodes)})"
                )
        self.reward_node: str | None = reward_node
        # Persistent-intervention state.
        self._active_interventions: Mapping[str, Value] | None = None
        self._mutilated_scm: StructuralCausalModel | None = None

    # ------------------------------------------------------------------
    # Causal interface availability
    # ------------------------------------------------------------------

    @property
    def has_causal_interface(self) -> bool:
        """``True`` iff the causal interface is fully operational.

        Requires the wrapped env to expose a non-``None`` ``.scm`` **and** ``reward_node``
        to be set and present in the graph.
        """
        raw_scm = getattr(self.env, "scm", None)
        return raw_scm is not None and self.reward_node is not None

    def _require_causal_interface(self, method: str) -> None:
        """Raise :class:`~causalrl.exceptions.CausalInterfaceUnavailableError` if disabled."""
        if not self.has_causal_interface:
            raw_scm = getattr(self.env, "scm", None)
            if raw_scm is None:
                reason = (
                    f"the wrapped env {type(self.env).__name__!r} has scm=None "
                    "(confounded MDPs and envs without an explicit SCM do not support "
                    "the causal interface)"
                )
            else:
                reason = "reward_node was not provided at construction time"
            raise CausalInterfaceUnavailableError(
                f"{method} requires has_causal_interface=True, but {reason}. "
                "Construct CausalEnvWrapper with a SCM-backed env and a valid reward_node "
                "to enable the causal interface."
            )

    # ------------------------------------------------------------------
    # Causal interface — pure graph queries
    # ------------------------------------------------------------------

    @property
    def scm(self) -> StructuralCausalModel | None:
        """The underlying SCM (live reference to the wrapped env's ``.scm``), or ``None``."""
        return cast("StructuralCausalModel | None", getattr(self.env, "scm", None))

    @property
    def reward_parents(self) -> list[str]:
        """Direct SCM parents of ``reward_node`` in graph-topological order.

        These are the variables that causally determine the immediate reward signal.  Pass
        their names as the ``factor_nodes`` argument of
        :func:`~causalrl.agents.factored_advantage.factored_advantage`.

        Raises
        ------
        CausalInterfaceUnavailableError
            When ``has_causal_interface`` is ``False``.
        """
        self._require_causal_interface("reward_parents")
        scm = cast("StructuralCausalModel", self.scm)
        reward_node = cast(str, self.reward_node)
        return scm.graph.parents(reward_node)

    def do(self, interventions: Mapping[str, Value]) -> StructuralCausalModel:
        """Return a *new* SCM mutilated by ``do(interventions)``.

        The running environment's SCM is NOT modified.  This is a pure causal-graph query
        suitable for off-policy reasoning or shaping.

        Parameters
        ----------
        interventions:
            Mapping ``{node_name: value}`` passed to
            :meth:`~causalrl.scm.scm.StructuralCausalModel.do`.

        Returns
        -------
        StructuralCausalModel
            The mutilated SCM under the specified do-intervention.

        Raises
        ------
        CausalInterfaceUnavailableError
            When ``has_causal_interface`` is ``False``.
        """
        self._require_causal_interface("do")
        return cast("StructuralCausalModel", self.scm).do(interventions)

    def intervene(self, node: str, value: Value) -> StructuralCausalModel:
        """Convenience wrapper: ``do({node: value})``.

        Parameters
        ----------
        node:
            The SCM variable to intervene on.
        value:
            The value to pin it to (scalar, sequence, or Tensor).

        Returns
        -------
        StructuralCausalModel
            The mutilated SCM under ``do({node: value})``.

        Raises
        ------
        CausalInterfaceUnavailableError
            When ``has_causal_interface`` is ``False``.
        """
        return self.do({node: value})

    # ------------------------------------------------------------------
    # Persistent interventional rollouts
    # ------------------------------------------------------------------

    @property
    def active_interventions(self) -> Mapping[str, Value] | None:
        """The currently stored intervention mapping, or ``None`` if none is active."""
        return self._active_interventions

    def set_intervention(self, interventions: Mapping[str, Value]) -> None:
        """Store a persistent intervention that affects subsequent ``reset`` and ``step``.

        After this call, every ``reset()`` and ``step()`` swaps the wrapped env's ``.scm``
        to the pre-computed mutilated SCM for the duration of the call, then restores the
        original SCM in a ``finally`` block.  Precomputed baselines stored on the env
        (e.g. ``arm_values``) are **not** recomputed.

        Parameters
        ----------
        interventions:
            Mapping ``{node_name: value}`` to pin persistently.

        Raises
        ------
        CausalInterfaceUnavailableError
            When ``has_causal_interface`` is ``False``.
        """
        self._require_causal_interface("set_intervention")
        self._active_interventions = dict(interventions)
        # Pre-compute the mutilated SCM once so set_intervention is idempotent-fast.
        self._mutilated_scm = cast("StructuralCausalModel", self.scm).do(interventions)

    def clear_intervention(self) -> None:
        """Remove the persistent intervention; subsequent calls use the unintervened SCM."""
        self._active_interventions = None
        self._mutilated_scm = None

    # ------------------------------------------------------------------
    # Gymnasium pass-through with optional SCM swap
    # ------------------------------------------------------------------

    def reset(
        self,
        *,
        seed: int | None = None,
        options: dict[str, Any] | None = None,
    ) -> tuple[Any, dict[str, Any]]:
        """Reset the wrapped environment, forwarding ``seed`` and ``options``.

        When a persistent intervention is active (``set_intervention`` was called), the
        wrapped env's ``.scm`` is temporarily replaced with the mutilated SCM for the
        duration of this call, then restored unconditionally.
        """
        if self._mutilated_scm is not None:
            orig_scm = cast("StructuralCausalModel", self.scm)
            self.env.scm = self._mutilated_scm  # type: ignore[union-attr]
            try:
                return self.env.reset(seed=seed, options=options)  # type: ignore[return-value]
            finally:
                self.env.scm = orig_scm  # type: ignore[union-attr]
        return self.env.reset(seed=seed, options=options)  # type: ignore[return-value]

    def step(self, action: Any) -> tuple[Any, float, bool, bool, dict[str, Any]]:
        """Step the wrapped environment.

        When a persistent intervention is active (``set_intervention`` was called), the
        wrapped env's ``.scm`` is temporarily replaced with the mutilated SCM for the
        duration of this call, then restored unconditionally.
        """
        if self._mutilated_scm is not None:
            orig_scm = cast("StructuralCausalModel", self.scm)
            self.env.scm = self._mutilated_scm  # type: ignore[union-attr]
            try:
                return self.env.step(action)  # type: ignore[return-value]
            finally:
                self.env.scm = orig_scm  # type: ignore[union-attr]
        return self.env.step(action)  # type: ignore[return-value]

has_causal_interface property

True iff the causal interface is fully operational.

Requires the wrapped env to expose a non-None .scm and reward_node to be set and present in the graph.

scm property

The underlying SCM (live reference to the wrapped env's .scm), or None.

reward_parents property

Direct SCM parents of reward_node in graph-topological order.

These are the variables that causally determine the immediate reward signal. Pass their names as the factor_nodes argument of :func:~causalrl.agents.factored_advantage.factored_advantage.

Raises

CausalInterfaceUnavailableError When has_causal_interface is False.

active_interventions property

The currently stored intervention mapping, or None if none is active.

do(interventions)

Return a new SCM mutilated by do(interventions).

The running environment's SCM is NOT modified. This is a pure causal-graph query suitable for off-policy reasoning or shaping.

Parameters

interventions: Mapping {node_name: value} passed to :meth:~causalrl.scm.scm.StructuralCausalModel.do.

Returns

StructuralCausalModel The mutilated SCM under the specified do-intervention.

Raises

CausalInterfaceUnavailableError When has_causal_interface is False.

Source code in src/causalrl/envs/wrapper.py

def do(self, interventions: Mapping[str, Value]) -> StructuralCausalModel:
    """Return a *new* SCM mutilated by ``do(interventions)``.

    The running environment's SCM is NOT modified.  This is a pure causal-graph query
    suitable for off-policy reasoning or shaping.

    Parameters
    ----------
    interventions:
        Mapping ``{node_name: value}`` passed to
        :meth:`~causalrl.scm.scm.StructuralCausalModel.do`.

    Returns
    -------
    StructuralCausalModel
        The mutilated SCM under the specified do-intervention.

    Raises
    ------
    CausalInterfaceUnavailableError
        When ``has_causal_interface`` is ``False``.
    """
    self._require_causal_interface("do")
    return cast("StructuralCausalModel", self.scm).do(interventions)

intervene(node, value)

Convenience wrapper: do({node: value}).

Parameters

node: The SCM variable to intervene on. value: The value to pin it to (scalar, sequence, or Tensor).

Returns

StructuralCausalModel The mutilated SCM under do({node: value}).

Raises

CausalInterfaceUnavailableError When has_causal_interface is False.

Source code in src/causalrl/envs/wrapper.py

def intervene(self, node: str, value: Value) -> StructuralCausalModel:
    """Convenience wrapper: ``do({node: value})``.

    Parameters
    ----------
    node:
        The SCM variable to intervene on.
    value:
        The value to pin it to (scalar, sequence, or Tensor).

    Returns
    -------
    StructuralCausalModel
        The mutilated SCM under ``do({node: value})``.

    Raises
    ------
    CausalInterfaceUnavailableError
        When ``has_causal_interface`` is ``False``.
    """
    return self.do({node: value})

set_intervention(interventions)

Store a persistent intervention that affects subsequent reset and step.

After this call, every reset() and step() swaps the wrapped env's .scm to the pre-computed mutilated SCM for the duration of the call, then restores the original SCM in a finally block. Precomputed baselines stored on the env (e.g. arm_values) are not recomputed.

Parameters

interventions: Mapping {node_name: value} to pin persistently.

Raises

CausalInterfaceUnavailableError When has_causal_interface is False.

Source code in src/causalrl/envs/wrapper.py

def set_intervention(self, interventions: Mapping[str, Value]) -> None:
    """Store a persistent intervention that affects subsequent ``reset`` and ``step``.

    After this call, every ``reset()`` and ``step()`` swaps the wrapped env's ``.scm``
    to the pre-computed mutilated SCM for the duration of the call, then restores the
    original SCM in a ``finally`` block.  Precomputed baselines stored on the env
    (e.g. ``arm_values``) are **not** recomputed.

    Parameters
    ----------
    interventions:
        Mapping ``{node_name: value}`` to pin persistently.

    Raises
    ------
    CausalInterfaceUnavailableError
        When ``has_causal_interface`` is ``False``.
    """
    self._require_causal_interface("set_intervention")
    self._active_interventions = dict(interventions)
    # Pre-compute the mutilated SCM once so set_intervention is idempotent-fast.
    self._mutilated_scm = cast("StructuralCausalModel", self.scm).do(interventions)

clear_intervention()

Remove the persistent intervention; subsequent calls use the unintervened SCM.

Source code in src/causalrl/envs/wrapper.py

def clear_intervention(self) -> None:
    """Remove the persistent intervention; subsequent calls use the unintervened SCM."""
    self._active_interventions = None
    self._mutilated_scm = None

reset(*, seed=None, options=None)

Reset the wrapped environment, forwarding seed and options.

When a persistent intervention is active (set_intervention was called), the wrapped env's .scm is temporarily replaced with the mutilated SCM for the duration of this call, then restored unconditionally.

Source code in src/causalrl/envs/wrapper.py

def reset(
    self,
    *,
    seed: int | None = None,
    options: dict[str, Any] | None = None,
) -> tuple[Any, dict[str, Any]]:
    """Reset the wrapped environment, forwarding ``seed`` and ``options``.

    When a persistent intervention is active (``set_intervention`` was called), the
    wrapped env's ``.scm`` is temporarily replaced with the mutilated SCM for the
    duration of this call, then restored unconditionally.
    """
    if self._mutilated_scm is not None:
        orig_scm = cast("StructuralCausalModel", self.scm)
        self.env.scm = self._mutilated_scm  # type: ignore[union-attr]
        try:
            return self.env.reset(seed=seed, options=options)  # type: ignore[return-value]
        finally:
            self.env.scm = orig_scm  # type: ignore[union-attr]
    return self.env.reset(seed=seed, options=options)  # type: ignore[return-value]

step(action)

Step the wrapped environment.

When a persistent intervention is active (set_intervention was called), the wrapped env's .scm is temporarily replaced with the mutilated SCM for the duration of this call, then restored unconditionally.

Source code in src/causalrl/envs/wrapper.py

def step(self, action: Any) -> tuple[Any, float, bool, bool, dict[str, Any]]:
    """Step the wrapped environment.

    When a persistent intervention is active (``set_intervention`` was called), the
    wrapped env's ``.scm`` is temporarily replaced with the mutilated SCM for the
    duration of this call, then restored unconditionally.
    """
    if self._mutilated_scm is not None:
        orig_scm = cast("StructuralCausalModel", self.scm)
        self.env.scm = self._mutilated_scm  # type: ignore[union-attr]
        try:
            return self.env.step(action)  # type: ignore[return-value]
        finally:
            self.env.scm = orig_scm  # type: ignore[union-attr]
    return self.env.step(action)  # type: ignore[return-value]

Gymnasium Env Registration

causalrl.envs.registration.register_envs()

Register causalrl demo environments in the Gymnasium registry.

Calling this function more than once in the same process is safe (idempotent).

After calling this function (or importing causalrl), you can use::

import gymnasium
import causalrl  # triggers register_envs()

env = gymnasium.make("causalrl/StructuralCausalBandit-v0")
vec = gymnasium.make_vec("causalrl/StructuralCausalBandit-v0", num_envs=2)

Source code in src/causalrl/envs/registration.py

def register_envs() -> None:
    """Register causalrl demo environments in the Gymnasium registry.

    Calling this function more than once in the same process is safe (idempotent).

    After calling this function (or importing ``causalrl``), you can use::

        import gymnasium
        import causalrl  # triggers register_envs()

        env = gymnasium.make("causalrl/StructuralCausalBandit-v0")
        vec = gymnasium.make_vec("causalrl/StructuralCausalBandit-v0", num_envs=2)
    """
    global _registered
    if _registered:
        return

    # gymnasium.register has an incompletely typed `kwargs` dict parameter in the stubs;
    # suppress the resulting pyright reportUnknownMemberType at each call site.
    gymnasium.register(  # type: ignore[reportUnknownMemberType]
        id="causalrl/StructuralCausalBandit-v0",
        entry_point="causalrl.envs.suite.scbandit:make_confounded_chain_env",
        # disable_env_checker keeps CI fast; callers can re-enable via check_env.
        disable_env_checker=True,
    )

    gymnasium.register(  # type: ignore[reportUnknownMemberType]
        id="causalrl/FrontdoorBandit-v0",
        entry_point="causalrl.envs.suite.scbandit:make_frontdoor_env",
        disable_env_checker=True,
    )

    _registered = True

Causal Graph-Factored Advantage (CGFA-PPO)

The pure-NumPy rollout arithmetic (no framework dependency):

causalrl.agents.factored_advantage.factor_rewards(factor_trace)

Per-factor rewards r^factor_{k,t} = phi_k(s_{t+1}) - phi_k(s_t).

This is the per-step quantity the CGFA environment wrapper publishes in arXiv:2605.06066 §E.1: the change in each SCM factor across a transition. The factors phi(s) are the SCM parents of the return node — for a :class:~causalrl.envs.wrapper.CausalEnvWrapper those are wrapper.reward_parents.

Parameters

factor_trace: Array of shape (T + 1, K): the factor values phi(s_0), …, phi(s_T) observed along a rollout of T transitions.

Returns

NDArray[np.float64] Shape (T, K) — the first difference along the time axis.

Raises

ValueError If factor_trace is not 2-D or has fewer than two rows.

Examples

import numpy as np phi = np.array([[0.0, 1.0], [1.0, 1.0], [1.0, 4.0]]) # (T+1=3, K=2) factor_rewards(phi) array([[1., 0.], [0., 3.]])

References

• Cunha, Mian, French and Liu (2026), arXiv:2605.06066, §E.1.

Source code in src/causalrl/agents/factored_advantage.py

def factor_rewards(factor_trace: NDArray[np.float64]) -> NDArray[np.float64]:
    """Per-factor rewards ``r^factor_{k,t} = phi_k(s_{t+1}) - phi_k(s_t)``.

    This is the per-step quantity the CGFA environment wrapper publishes in
    arXiv:2605.06066 §E.1: the *change* in each SCM factor across a transition.  The factors
    ``phi(s)`` are the SCM parents of the return node — for a
    :class:`~causalrl.envs.wrapper.CausalEnvWrapper` those are ``wrapper.reward_parents``.

    Parameters
    ----------
    factor_trace:
        Array of shape ``(T + 1, K)``: the factor values ``phi(s_0), …, phi(s_T)`` observed
        along a rollout of ``T`` transitions.

    Returns
    -------
    NDArray[np.float64]
        Shape ``(T, K)`` — the first difference along the time axis.

    Raises
    ------
    ValueError
        If ``factor_trace`` is not 2-D or has fewer than two rows.

    Examples
    --------
    >>> import numpy as np
    >>> phi = np.array([[0.0, 1.0], [1.0, 1.0], [1.0, 4.0]])  # (T+1=3, K=2)
    >>> factor_rewards(phi)
    array([[1., 0.],
           [0., 3.]])

    References
    ----------
    * Cunha, Mian, French and Liu (2026), arXiv:2605.06066, §E.1.
    """
    trace = np.asarray(factor_trace, dtype=np.float64)
    if trace.ndim != 2:
        raise ValueError(f"factor_trace must be 2-D (T+1, K); got shape {trace.shape}")
    if trace.shape[0] < 2:
        raise ValueError(
            f"factor_trace needs at least 2 rows (phi(s_0) and phi(s_1)); got {trace.shape[0]}"
        )
    return np.diff(trace, axis=0)

causalrl.agents.factored_advantage.factor_gae(rewards, values, *, gamma, lam=1.0, bootstrap_values=None, dones=None)

Per-factor advantages and returns for the CGFA-PPO factor critic.

Runs the standard generalised-advantage recursion independently per factor::

delta_{k,t} = r^factor_{k,t} + gamma * V_k(s_{t+1}) * (1 - done_t) - V_k(s_t)
A_{k,t}     = delta_{k,t} + gamma * lam * (1 - done_t) * A_{k,t+1}
G_{k,t}     = A_{k,t} + V_k(s_t)

At lam=1.0 with no bootstrap this collapses exactly to the paper's written equations: G_{k,t} = sum_i gamma^i r^factor_{k,t+i} (Eq. 8) and A_{k,t} = G_{k,t} - V_k(s_t) (Eq. 10).

.. note:: Paper ambiguity. arXiv:2605.06066 writes Eq. 8 / Eq. 10 as a Monte-Carlo return and its residual, but the surrounding prose (§5, §E.2) says the per-factor advantage is computed "using the same generalised-advantage truncation as the scalar critic", and Table 6 lists a single GAE lambda = 0.95 for the shared backbone. Those two readings coincide only at lambda = 1. lam therefore defaults to 1.0, which reproduces the paper's explicit equations; pass lam=0.95 for the prose/Table-6 reading.

Parameters

rewards: Per-factor rewards, shape (T, K) — typically :func:factor_rewards output. values: Per-factor critic estimates V_k(s_t), shape (T, K) — typically :meth:~causalrl.agents.cgfa_critic.FactoredCritic.values. gamma: Discount factor, shared with the scalar return (paper default 0.995). lam: GAE lambda. 1.0 (default) reproduces Eq. 8 / Eq. 10 exactly. bootstrap_values: V_k(s_T) for the state after the last stored transition, shape (K,). None (default) means zero — an episode that ended at T. dones: Boolean episode-termination flags, shape (T,). dones[t] true cuts the bootstrap and the recursion at step t. None means no interior terminations.

Returns

tuple[NDArray[np.float64], NDArray[np.float64]] (advantages, returns), both of shape (T, K).

Raises

ValueError On any shape mismatch between rewards, values, bootstrap_values and dones, or if rewards is not 2-D.

Examples

Two factors that accumulate at different rates, with a zero critic so the advantage is the raw discounted return (Eq. 8):

import numpy as np r = np.array([[1.0, 0.0], [0.0, 1.0]]) V = np.zeros((2, 2)) adv, ret = factor_gae(r, V, gamma=0.5) ret array([[1. , 0.5], [0. , 1. ]])

References

• Cunha, Mian, French and Liu (2026), arXiv:2605.06066, Eq. 8, Eq. 10, §E.2.

Source code in src/causalrl/agents/factored_advantage.py

def factor_gae(
    rewards: NDArray[np.float64],
    values: NDArray[np.float64],
    *,
    gamma: float,
    lam: float = 1.0,
    bootstrap_values: NDArray[np.float64] | None = None,
    dones: NDArray[np.bool_] | None = None,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Per-factor advantages and returns for the CGFA-PPO factor critic.

    Runs the standard generalised-advantage recursion **independently per factor**::

        delta_{k,t} = r^factor_{k,t} + gamma * V_k(s_{t+1}) * (1 - done_t) - V_k(s_t)
        A_{k,t}     = delta_{k,t} + gamma * lam * (1 - done_t) * A_{k,t+1}
        G_{k,t}     = A_{k,t} + V_k(s_t)

    At ``lam=1.0`` with no bootstrap this collapses exactly to the paper's written
    equations: ``G_{k,t} = sum_i gamma^i r^factor_{k,t+i}`` (Eq. 8) and
    ``A_{k,t} = G_{k,t} - V_k(s_t)`` (Eq. 10).

    .. note:: **Paper ambiguity.** arXiv:2605.06066 writes Eq. 8 / Eq. 10 as a
       Monte-Carlo return and its residual, but the surrounding prose (§5, §E.2) says the
       per-factor advantage is computed "using the same generalised-advantage truncation as
       the scalar critic", and Table 6 lists a single GAE ``lambda = 0.95`` for the shared
       backbone.  Those two readings coincide only at ``lambda = 1``.  ``lam`` therefore
       defaults to ``1.0``, which reproduces the paper's *explicit equations*; pass
       ``lam=0.95`` for the prose/Table-6 reading.

    Parameters
    ----------
    rewards:
        Per-factor rewards, shape ``(T, K)`` — typically :func:`factor_rewards` output.
    values:
        Per-factor critic estimates ``V_k(s_t)``, shape ``(T, K)`` — typically
        :meth:`~causalrl.agents.cgfa_critic.FactoredCritic.values`.
    gamma:
        Discount factor, shared with the scalar return (paper default ``0.995``).
    lam:
        GAE ``lambda``.  ``1.0`` (default) reproduces Eq. 8 / Eq. 10 exactly.
    bootstrap_values:
        ``V_k(s_T)`` for the state after the last stored transition, shape ``(K,)``.
        ``None`` (default) means zero — an episode that ended at ``T``.
    dones:
        Boolean episode-termination flags, shape ``(T,)``.  ``dones[t]`` true cuts the
        bootstrap and the recursion at step ``t``.  ``None`` means no interior terminations.

    Returns
    -------
    tuple[NDArray[np.float64], NDArray[np.float64]]
        ``(advantages, returns)``, both of shape ``(T, K)``.

    Raises
    ------
    ValueError
        On any shape mismatch between ``rewards``, ``values``, ``bootstrap_values`` and
        ``dones``, or if ``rewards`` is not 2-D.

    Examples
    --------
    Two factors that accumulate at different rates, with a zero critic so the advantage is
    the raw discounted return (Eq. 8):

    >>> import numpy as np
    >>> r = np.array([[1.0, 0.0], [0.0, 1.0]])
    >>> V = np.zeros((2, 2))
    >>> adv, ret = factor_gae(r, V, gamma=0.5)
    >>> ret
    array([[1. , 0.5],
           [0. , 1. ]])

    References
    ----------
    * Cunha, Mian, French and Liu (2026), arXiv:2605.06066, Eq. 8, Eq. 10, §E.2.
    """
    r = np.asarray(rewards, dtype=np.float64)
    v = np.asarray(values, dtype=np.float64)
    if r.ndim != 2:
        raise ValueError(f"rewards must be 2-D (T, K); got shape {r.shape}")
    if v.shape != r.shape:
        raise ValueError(f"values must have the same shape as rewards {r.shape}; got {v.shape}")
    t_steps, k = r.shape

    if bootstrap_values is None:
        boot = np.zeros(k, dtype=np.float64)
    else:
        boot = np.asarray(bootstrap_values, dtype=np.float64)
        if boot.shape != (k,):
            raise ValueError(f"bootstrap_values must have shape ({k},); got {boot.shape}")

    if dones is None:
        not_done = np.ones(t_steps, dtype=np.float64)
    else:
        d = np.asarray(dones)
        if d.shape != (t_steps,):
            raise ValueError(f"dones must have shape ({t_steps},); got {d.shape}")
        not_done = 1.0 - d.astype(np.float64)

    advantages = np.zeros_like(r)
    carry = np.zeros(k, dtype=np.float64)
    next_values = boot
    for t in range(t_steps - 1, -1, -1):
        mask = not_done[t]
        delta = r[t] + gamma * next_values * mask - v[t]
        carry = delta + gamma * lam * mask * carry
        advantages[t] = carry
        next_values = v[t]
    return advantages, advantages + v

causalrl.agents.factored_advantage.blend_advantages(scalar_advantages, factor_advantages, *, gate, weights=None)

The CGFA-PPO state-conditional residual blend (arXiv:2605.06066, Eq. 11).

A_used(s_t, a_t) = (1 - g(s_t)) A^scalar(s_t, a_t) + g(s_t) * sum_k w_k A_k(s_t, a_t)

g -> 0 recovers vanilla PPO exactly; g -> 1 hands the policy update entirely to the factor-aligned signal. In the paper g is a learned state-conditional MLP and w = softmax(beta) are learned mixture logits initialised from the SCM's logistic-regression coefficients on the return node; both are supplied by :class:~causalrl.agents.cgfa_critic.FactoredCritic. Passing constants here is a valid ablation (the paper's "CGFA without the gate" row).

Parameters

scalar_advantages: The ordinary (scalar-critic) advantage, shape (T,). factor_advantages: Per-factor advantages A_{k,t}, shape (T, K) — :func:factor_gae output. gate: g(s_t) in (0, 1): shape (T,) or a scalar broadcast over the rollout. weights: Mixture weights w_k, shape (K,). None means uniform 1/K. The paper uses softmax(beta), which sums to 1; this function does not renormalise, so an unnormalised vector is passed through as given.

Returns

NDArray[np.float64] Shape (T,) — the advantage to feed to the policy-gradient surrogate.

Raises

ValueError On any shape mismatch, or if gate falls outside [0, 1].

Examples

import numpy as np a_scalar = np.array([1.0, 1.0]) a_factor = np.array([[4.0, 0.0], [4.0, 0.0]]) blend_advantages(a_scalar, a_factor, gate=0.0) # pure PPO array([1., 1.]) blend_advantages(a_scalar, a_factor, gate=1.0) # pure factored, uniform w array([2., 2.]) blend_advantages(a_scalar, a_factor, gate=0.5) # half and half array([1.5, 1.5])

References

• Cunha, Mian, French and Liu (2026), arXiv:2605.06066, Eq. 11, §E.3.

Source code in src/causalrl/agents/factored_advantage.py

def blend_advantages(
    scalar_advantages: NDArray[np.float64],
    factor_advantages: NDArray[np.float64],
    *,
    gate: NDArray[np.float64] | float,
    weights: NDArray[np.float64] | None = None,
) -> NDArray[np.float64]:
    """The CGFA-PPO state-conditional residual blend (arXiv:2605.06066, Eq. 11).

    ``A_used(s_t, a_t) = (1 - g(s_t)) A^scalar(s_t, a_t) + g(s_t) * sum_k w_k A_k(s_t, a_t)``

    ``g -> 0`` recovers vanilla PPO exactly; ``g -> 1`` hands the policy update entirely to
    the factor-aligned signal.  In the paper ``g`` is a learned state-conditional MLP and
    ``w = softmax(beta)`` are learned mixture logits initialised from the SCM's
    logistic-regression coefficients on the return node; both are supplied by
    :class:`~causalrl.agents.cgfa_critic.FactoredCritic`.  Passing constants here is a valid
    ablation (the paper's "CGFA without the gate" row).

    Parameters
    ----------
    scalar_advantages:
        The ordinary (scalar-critic) advantage, shape ``(T,)``.
    factor_advantages:
        Per-factor advantages ``A_{k,t}``, shape ``(T, K)`` — :func:`factor_gae` output.
    gate:
        ``g(s_t)`` in ``(0, 1)``: shape ``(T,)`` or a scalar broadcast over the rollout.
    weights:
        Mixture weights ``w_k``, shape ``(K,)``.  ``None`` means uniform ``1/K``.  The paper
        uses ``softmax(beta)``, which sums to 1; this function does not renormalise, so an
        unnormalised vector is passed through as given.

    Returns
    -------
    NDArray[np.float64]
        Shape ``(T,)`` — the advantage to feed to the policy-gradient surrogate.

    Raises
    ------
    ValueError
        On any shape mismatch, or if ``gate`` falls outside ``[0, 1]``.

    Examples
    --------
    >>> import numpy as np
    >>> a_scalar = np.array([1.0, 1.0])
    >>> a_factor = np.array([[4.0, 0.0], [4.0, 0.0]])
    >>> blend_advantages(a_scalar, a_factor, gate=0.0)          # pure PPO
    array([1., 1.])
    >>> blend_advantages(a_scalar, a_factor, gate=1.0)          # pure factored, uniform w
    array([2., 2.])
    >>> blend_advantages(a_scalar, a_factor, gate=0.5)          # half and half
    array([1.5, 1.5])

    References
    ----------
    * Cunha, Mian, French and Liu (2026), arXiv:2605.06066, Eq. 11, §E.3.
    """
    a_scalar = np.asarray(scalar_advantages, dtype=np.float64)
    a_factor = np.asarray(factor_advantages, dtype=np.float64)
    if a_factor.ndim != 2:
        raise ValueError(f"factor_advantages must be 2-D (T, K); got shape {a_factor.shape}")
    t_steps, k = a_factor.shape
    if a_scalar.shape != (t_steps,):
        raise ValueError(
            f"scalar_advantages must have shape ({t_steps},) to match factor_advantages rows; "
            f"got {a_scalar.shape}"
        )

    g = np.asarray(gate, dtype=np.float64)
    if g.ndim == 0:
        g = np.full(t_steps, float(g))
    elif g.shape != (t_steps,):
        raise ValueError(f"gate must be a scalar or have shape ({t_steps},); got {g.shape}")
    if bool(np.any(g < 0.0)) or bool(np.any(g > 1.0)):
        raise ValueError("gate must lie in [0, 1] — it is a residual mixing coefficient")

    if weights is None:
        w = np.full(k, 1.0 / k, dtype=np.float64)
    else:
        w = np.asarray(weights, dtype=np.float64)
        if w.shape != (k,):
            raise ValueError(f"weights must have shape ({k},) to match K; got {w.shape}")

    return (1.0 - g) * a_scalar + g * (a_factor @ w)

causalrl.agents.factored_advantage.factored_advantage(factor_values, baselines, *, config=None, aggregation='sum', weights=None)

Combine per-factor value estimates against a common baseline into one advantage.

Given K causal parent factors of the return, and for each rollout step a vector of per-factor value estimates V_1, …, V_K and a shared scalar baseline b, the per-factor advantage is A_i = V_i - b and the combined advantage is their (weighted) sum or mean.

When K = 1 (single factor) the output reduces exactly to the standard advantage A = V - b, so this is a strict generalisation of the scalar advantage.

Relation to CGFA-PPO

This is a decomposition primitive, not Eq. 10 of arXiv:2605.06066. Eq. 10 is A_{k,t} = G_{k,t} - V_k(s_t): the per-factor return minus the per-factor value, i.e. a per-factor baseline of shape (T, K), which this signature cannot express (its baselines is (T,)). Use :func:factor_gae for Eq. 10 and :func:blend_advantages for the Eq. 11 residual blend; both are consumed by :class:~causalrl.agents.cgfa_critic.FactoredCritic.

Parameters

factor_values: Array of shape (T, K) where T is the number of rollout steps and K is the number of causal factors (SCM parents of the return). Each column [:,i] is the critic's value estimate for factor i. baselines: Array of shape (T,) — the shared scalar baseline for each step (typically the current value-function estimate V(s_t)). config: A :class:FactoredAdvantageConfig that carries factor_nodes, aggregation, and optional weights. When provided, aggregation and weights keyword arguments are ignored (config takes precedence). aggregation: Used when config is None. "sum" (default) or "mean". weights: Used when config is None. Per-factor weights, shape (K,). None means uniform unit weights.

Returns

NDArray[np.float64] Shape (T,) — the combined causal-graph-factored advantage for each step.

Raises

ValueError If factor_values is not 2-D, if baselines length does not match T, or if weights shape does not match K.

Examples

Single-factor (reduces to standard advantage):

import numpy as np V = np.array([[2.0], [3.0], [1.0]]) # (T=3, K=1) b = np.array([1.5, 2.5, 0.5]) factored_advantage(V, b) array([0.5, 0.5, 0.5])

Two-factor sum (CGFA-PPO with two SCM parents of the return):

V2 = np.array([[2.0, 1.0], [3.0, 0.5]]) # (T=2, K=2) b2 = np.array([1.5, 2.0]) factored_advantage(V2, b2) # A_i = V_i - b; sum over i array([ 0. , -0.5])

References

• Cristiano da Costa Cunha, Ajmal Mian, Tim French, and Wei Liu (2026). "Causal Reinforcement Learning for Complex Card Games: A Magic: The Gathering Benchmark." arXiv:2605.06066.

Source code in src/causalrl/agents/factored_advantage.py

def factored_advantage(
    factor_values: NDArray[np.float64],
    baselines: NDArray[np.float64],
    *,
    config: FactoredAdvantageConfig | None = None,
    aggregation: Literal["sum", "mean"] = "sum",
    weights: NDArray[np.float64] | None = None,
) -> NDArray[np.float64]:
    """Combine per-factor value estimates against a common baseline into one advantage.

    Given ``K`` causal parent factors of the return, and for each rollout step a vector of
    per-factor value estimates ``V_1, …, V_K`` and a **shared scalar** baseline ``b``, the
    per-factor advantage is ``A_i = V_i - b`` and the combined advantage is their (weighted)
    sum or mean.

    When ``K = 1`` (single factor) the output reduces exactly to the standard advantage
    ``A = V - b``, so this is a strict generalisation of the scalar advantage.

    Relation to CGFA-PPO
    --------------------
    This is a decomposition primitive, **not** Eq. 10 of arXiv:2605.06066.  Eq. 10 is
    ``A_{k,t} = G_{k,t} - V_k(s_t)``: the per-factor *return* minus the *per-factor* value,
    i.e. a per-factor baseline of shape ``(T, K)``, which this signature cannot express (its
    ``baselines`` is ``(T,)``).  Use :func:`factor_gae` for Eq. 10 and
    :func:`blend_advantages` for the Eq. 11 residual blend; both are consumed by
    :class:`~causalrl.agents.cgfa_critic.FactoredCritic`.

    Parameters
    ----------
    factor_values:
        Array of shape ``(T, K)`` where ``T`` is the number of rollout steps and ``K`` is
        the number of causal factors (SCM parents of the return).  Each column ``[:,i]``
        is the critic's value estimate for factor ``i``.
    baselines:
        Array of shape ``(T,)`` — the shared scalar baseline for each step (typically the
        current value-function estimate ``V(s_t)``).
    config:
        A :class:`FactoredAdvantageConfig` that carries ``factor_nodes``, ``aggregation``,
        and optional ``weights``.  When provided, ``aggregation`` and ``weights`` keyword
        arguments are ignored (config takes precedence).
    aggregation:
        Used when ``config`` is ``None``.  ``"sum"`` (default) or ``"mean"``.
    weights:
        Used when ``config`` is ``None``.  Per-factor weights, shape ``(K,)``.  ``None``
        means uniform unit weights.

    Returns
    -------
    NDArray[np.float64]
        Shape ``(T,)`` — the combined causal-graph-factored advantage for each step.

    Raises
    ------
    ValueError
        If ``factor_values`` is not 2-D, if ``baselines`` length does not match ``T``, or
        if ``weights`` shape does not match ``K``.

    Examples
    --------
    Single-factor (reduces to standard advantage):

    >>> import numpy as np
    >>> V = np.array([[2.0], [3.0], [1.0]])   # (T=3, K=1)
    >>> b = np.array([1.5, 2.5, 0.5])
    >>> factored_advantage(V, b)
    array([0.5, 0.5, 0.5])

    Two-factor sum (CGFA-PPO with two SCM parents of the return):

    >>> V2 = np.array([[2.0, 1.0], [3.0, 0.5]])  # (T=2, K=2)
    >>> b2 = np.array([1.5, 2.0])
    >>> factored_advantage(V2, b2)   # A_i = V_i - b; sum over i
    array([ 0. , -0.5])

    References
    ----------
    * Cristiano da Costa Cunha, Ajmal Mian, Tim French, and Wei Liu (2026). "Causal
      Reinforcement Learning for Complex Card Games: A Magic: The Gathering Benchmark."
      arXiv:2605.06066.
    """
    fv = np.asarray(factor_values, dtype=np.float64)
    bl = np.asarray(baselines, dtype=np.float64)

    if fv.ndim != 2:
        raise ValueError(f"factor_values must be 2-D (T, K); got shape {fv.shape}")
    t, k = fv.shape
    if bl.shape != (t,):
        raise ValueError(
            f"baselines must have shape ({t},) to match factor_values rows; got {bl.shape}"
        )

    # Resolve aggregation and weights.
    if config is not None:
        agg = config.aggregation
        w = config.weights_array
        if len(w) != k:
            raise ValueError(
                f"FactoredAdvantageConfig has {len(w)} factors but factor_values has K={k} columns"
            )
    else:
        agg = aggregation
        if weights is not None:
            w = np.asarray(weights, dtype=np.float64)
            if w.shape != (k,):
                raise ValueError(f"weights must have shape ({k},) to match K; got {w.shape}")
        else:
            w = np.ones(k, dtype=np.float64)

    # Per-factor advantages: A_i(t) = V_i(t) - b(t).  Shape (T, K).
    per_factor: NDArray[np.float64] = fv - bl[:, np.newaxis]  # broadcast baseline

    # Combine: weighted sum or weighted mean over the K factors.
    combined = per_factor @ w  # (T,)
    if agg == "mean":
        combined = combined / float(w.sum()) if w.sum() != 0.0 else combined

    return combined

causalrl.agents.factored_advantage.FactoredAdvantageConfig dataclass

Configuration bundle for :func:factored_advantage.

Parameters

factor_nodes: Ordered list of SCM parent-node names whose per-factor value estimates are passed to :func:factored_advantage. The order must match the columns of the factor_values array. aggregation: How to combine per-factor advantages into the final scalar advantage.

* ``"sum"`` (default): ``A = Σ_i A_i`` — the formulation in arXiv:2605.06066.
* ``"mean"``: ``A = mean_i A_i`` — normalises when the number of factors varies.

weights: Optional per-factor weights w_i (length must match factor_nodes when provided). The combined advantage becomes A = Σ_i w_i * A_i for aggregation="sum" or the weighted mean for aggregation="mean". None means uniform unit weights.

Source code in src/causalrl/agents/factored_advantage.py

@dataclass
class FactoredAdvantageConfig:
    """Configuration bundle for :func:`factored_advantage`.

    Parameters
    ----------
    factor_nodes:
        Ordered list of SCM parent-node names whose per-factor value estimates are passed to
        :func:`factored_advantage`.  The order must match the columns of the
        ``factor_values`` array.
    aggregation:
        How to combine per-factor advantages into the final scalar advantage.

        * ``"sum"`` (default): ``A = Σ_i A_i`` — the formulation in arXiv:2605.06066.
        * ``"mean"``: ``A = mean_i A_i`` — normalises when the number of factors varies.
    weights:
        Optional per-factor weights ``w_i`` (length must match ``factor_nodes`` when
        provided).  The combined advantage becomes ``A = Σ_i w_i * A_i`` for
        ``aggregation="sum"`` or the weighted mean for ``aggregation="mean"``.  ``None``
        means uniform unit weights.
    """

    factor_nodes: list[str]
    aggregation: Literal["sum", "mean"] = "sum"
    weights: list[float] | None = None
    # derived / validated at post-init
    _weights_arr: NDArray[np.float64] = field(init=False, repr=False)

    def __post_init__(self) -> None:
        n = len(self.factor_nodes)
        if self.weights is not None:
            if len(self.weights) != n:
                raise ValueError(
                    f"weights length ({len(self.weights)}) must match factor_nodes length ({n})"
                )
            self._weights_arr = np.asarray(self.weights, dtype=np.float64)
        else:
            self._weights_arr = np.ones(n, dtype=np.float64)

    @property
    def weights_array(self) -> NDArray[np.float64]:
        """The validated per-factor weight vector (uniform unit weights when unset)."""
        return self._weights_arr

weights_array property

The validated per-factor weight vector (uniform unit weights when unset).

The K-head critic that makes it an algorithm (needs the causalrl[torch] extra):

causalrl.agents.cgfa_critic.FactoredCritic

The K-head causal-factor critic of CGFA-PPO (arXiv:2605.06066, §E.1).

One value head per SCM parent of the return, on a trunk shared with the scalar critic, plus the learnable mixture logits and the state-conditional residual gate that turn those K heads into a single advantage for the policy-gradient surrogate.

The heads only differentiate because they are trained against different targets: head k regresses the per-factor return G_{k,t} accumulated from r^factor_{k,t} = phi_k(s_{t+1}) - phi_k(s_t). Give every head the same target and they collapse to the same function — which is exactly the degenerate "share one value head" configuration this class exists to replace.

Parameters

obs_dim: Dimensionality of the (flattened) observation fed to the trunk. factor_nodes: Ordered SCM node names of the return's causal parents — for a :class:~causalrl.envs.wrapper.CausalEnvWrapper, wrapper.reward_parents. Their order fixes the column order of every (T, K) array this class consumes or emits. config: A :class:CGFACriticConfig; None uses the paper's defaults. mixture_init: Initial beta, shape (K,). None gives zeros (a uniform mixture). The paper seeds beta with the SCM's logistic-regression coefficients on the return node, so that training starts from the structural prior. seed: Seed for torch's global RNG before parameter initialisation, for reproducibility.

Raises

ImportError At construction, when PyTorch is not installed — naming the causalrl[torch] extra. ValueError On a non-positive obs_dim, an empty factor_nodes, duplicate factor names, or a mixture_init whose length does not match factor_nodes.

Examples

import numpy as np critic = FactoredCritic(obs_dim=3, factor_nodes=["X3", "U"], seed=0) obs = np.zeros((5, 3)) scalar, factors = critic.values(obs) factors.shape (5, 2) float(critic.mixture_weights().sum()) 1.0

References

• Cunha, Mian, French and Liu (2026), arXiv:2605.06066 — §5, Appendix E, Algorithm 1.

Source code in src/causalrl/agents/cgfa_critic.py

class FactoredCritic:
    """The ``K``-head causal-factor critic of CGFA-PPO (arXiv:2605.06066, §E.1).

    One value head per SCM parent of the return, on a trunk shared with the scalar critic,
    plus the learnable mixture logits and the state-conditional residual gate that turn those
    ``K`` heads into a single advantage for the policy-gradient surrogate.

    The heads only differentiate because they are trained against *different* targets: head
    ``k`` regresses the per-factor return ``G_{k,t}`` accumulated from
    ``r^factor_{k,t} = phi_k(s_{t+1}) - phi_k(s_t)``.  Give every head the same target and
    they collapse to the same function — which is exactly the degenerate "share one value
    head" configuration this class exists to replace.

    Parameters
    ----------
    obs_dim:
        Dimensionality of the (flattened) observation fed to the trunk.
    factor_nodes:
        Ordered SCM node names of the return's causal parents — for a
        :class:`~causalrl.envs.wrapper.CausalEnvWrapper`, ``wrapper.reward_parents``.  Their
        order fixes the column order of every ``(T, K)`` array this class consumes or emits.
    config:
        A :class:`CGFACriticConfig`; ``None`` uses the paper's defaults.
    mixture_init:
        Initial ``beta``, shape ``(K,)``.  ``None`` gives zeros (a uniform mixture).  The
        paper seeds ``beta`` with the SCM's logistic-regression coefficients on the return
        node, so that training starts from the structural prior.
    seed:
        Seed for torch's global RNG before parameter initialisation, for reproducibility.

    Raises
    ------
    ImportError
        At construction, when PyTorch is not installed — naming the ``causalrl[torch]``
        extra.
    ValueError
        On a non-positive ``obs_dim``, an empty ``factor_nodes``, duplicate factor names, or
        a ``mixture_init`` whose length does not match ``factor_nodes``.

    Examples
    --------
    >>> import numpy as np
    >>> critic = FactoredCritic(obs_dim=3, factor_nodes=["X3", "U"], seed=0)
    >>> obs = np.zeros((5, 3))
    >>> scalar, factors = critic.values(obs)
    >>> factors.shape
    (5, 2)
    >>> float(critic.mixture_weights().sum())
    1.0

    References
    ----------
    * Cunha, Mian, French and Liu (2026), arXiv:2605.06066 — §5, Appendix E, Algorithm 1.
    """

    def __init__(
        self,
        obs_dim: int,
        factor_nodes: Sequence[str],
        *,
        config: CGFACriticConfig | None = None,
        mixture_init: NDArray[np.float64] | Sequence[float] | None = None,
        seed: int | None = None,
    ) -> None:
        _require_torch()
        import torch

        if obs_dim <= 0:
            raise ValueError(f"obs_dim must be positive; got {obs_dim}")
        nodes = list(factor_nodes)
        if not nodes:
            raise ValueError("factor_nodes must name at least one SCM parent of the return")
        if len(set(nodes)) != len(nodes):
            raise ValueError(f"factor_nodes must be unique; got {nodes}")

        self._factor_nodes = nodes
        self._config = config if config is not None else CGFACriticConfig()
        cfg = self._config
        k = len(nodes)

        if seed is not None:
            torch.manual_seed(seed)  # type: ignore[reportPrivateImportUsage]

        if mixture_init is None:
            beta0 = torch.zeros(k)  # type: ignore[reportPrivateImportUsage]
        else:
            beta_arr = np.asarray(mixture_init, dtype=np.float64)
            if beta_arr.shape != (k,):
                raise ValueError(
                    f"mixture_init must have shape ({k},) to match factor_nodes; "
                    f"got {beta_arr.shape}"
                )
            beta0 = torch.as_tensor(beta_arr, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]

        trunk_layers: list[torch.nn.Module] = []
        width = obs_dim
        for out in cfg.hidden:
            trunk_layers.append(torch.nn.Linear(width, out))  # type: ignore[reportPrivateImportUsage]
            trunk_layers.append(torch.nn.ReLU())  # type: ignore[reportPrivateImportUsage]
            width = out
        self._trunk = torch.nn.Sequential(*trunk_layers)  # type: ignore[reportUnknownMemberType]
        self._scalar_head = torch.nn.Linear(width, 1)  # type: ignore[reportPrivateImportUsage]
        self._factor_head = torch.nn.Linear(width, k)  # type: ignore[reportPrivateImportUsage]

        gate_out = torch.nn.Linear(cfg.gate_hidden, 1)  # type: ignore[reportPrivateImportUsage]
        # Zero weights + logit(gate_init) bias => g(s) == gate_init everywhere at step 0.
        with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
            gate_out.weight.zero_()
            gate_out.bias.fill_(float(np.log(cfg.gate_init / (1.0 - cfg.gate_init))))
        self._gate_head = torch.nn.Sequential(  # type: ignore[reportUnknownMemberType]
            torch.nn.Linear(width, cfg.gate_hidden),  # type: ignore[reportPrivateImportUsage]
            torch.nn.Tanh(),  # type: ignore[reportPrivateImportUsage]
            gate_out,
        )
        self._mixture_logits = torch.nn.Parameter(beta0)  # type: ignore[reportPrivateImportUsage]

        self._module = torch.nn.ModuleList(  # type: ignore[reportUnknownMemberType]
            [
                self._trunk,
                self._scalar_head,
                self._factor_head,
                self._gate_head,
                torch.nn.ParameterList([self._mixture_logits]),  # type: ignore[reportPrivateImportUsage]
            ]
        )
        self._optimizer = torch.optim.Adam(  # type: ignore[reportPrivateImportUsage]
            self._module.parameters(), lr=cfg.learning_rate
        )

    # ------------------------------------------------------------------ properties

    @property
    def factor_nodes(self) -> list[str]:
        """The ordered SCM parent names, one per value head."""
        return list(self._factor_nodes)

    @property
    def n_factors(self) -> int:
        """``K`` — the number of causal factors, hence of value heads."""
        return len(self._factor_nodes)

    @property
    def config(self) -> CGFACriticConfig:
        """The architecture / loss-coefficient bundle this critic was built with."""
        return self._config

    @property
    def module(self) -> torch.nn.Module:
        """The underlying ``torch.nn.Module`` — for ``state_dict``, ``.to(device)``, saving."""
        return self._module

    def parameters(self) -> Iterator[Tensor]:
        """Every learnable tensor: trunk, three heads, and the mixture logits ``beta``.

        Hand these to a joint optimiser alongside the actor's parameters when following
        Algorithm 1 line 22 literally (one Adam step on the whole of Eq. 13).
        """
        return self._module.parameters()  # type: ignore[reportUnknownMemberType,no-any-return]

    # ------------------------------------------------------------------ inference

    def _as_input(self, observations: NDArray[np.float64]) -> Tensor:
        import torch

        obs = np.asarray(observations, dtype=np.float32)
        if obs.ndim != 2:
            raise ValueError(f"observations must be 2-D (T, obs_dim); got shape {obs.shape}")
        return torch.as_tensor(obs)  # type: ignore[reportPrivateImportUsage]

    def _heads(self, obs: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """``(V(s), V_k(s), g(s))`` from the shared trunk (§E.1)."""
        import torch

        features = self._trunk(obs)
        scalar = self._scalar_head(features).squeeze(-1)
        factors = self._factor_head(features)
        gate = torch.sigmoid(self._gate_head(features).squeeze(-1))  # type: ignore[reportPrivateImportUsage]
        return scalar, factors, gate

    def values(
        self, observations: NDArray[np.float64]
    ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
        """``(V(s_t), V_k(s_t))`` for a batch of observations, shapes ``(T,)`` and ``(T, K)``."""
        import torch

        with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
            scalar, factors, _ = self._heads(self._as_input(observations))
        return (
            scalar.numpy().astype(np.float64),
            factors.numpy().astype(np.float64),
        )

    def gate(self, observations: NDArray[np.float64]) -> NDArray[np.float64]:
        """The state-conditional residual gate ``g(s_t) in (0, 1)``, shape ``(T,)`` (§E.3)."""
        import torch

        with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
            _, _, gate = self._heads(self._as_input(observations))
        return gate.numpy().astype(np.float64)

    def mixture_weights(self) -> NDArray[np.float64]:
        """``w = softmax(beta)``, shape ``(K,)`` — the learned per-factor mixture (§E.1)."""
        import torch

        with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
            w = torch.softmax(self._mixture_logits, dim=0)  # type: ignore[reportPrivateImportUsage]
        return w.numpy().astype(np.float64)

    # ------------------------------------------------------------------ rollout arithmetic

    def advantages(
        self,
        observations: NDArray[np.float64],
        factor_rewards: NDArray[np.float64],
        scalar_advantages: NDArray[np.float64],
        *,
        gamma: float = 0.995,
        lam: float = 1.0,
        bootstrap_values: NDArray[np.float64] | None = None,
        dones: NDArray[np.bool_] | None = None,
        normalize: bool = True,
    ) -> CGFAAdvantages:
        """Algorithm 1 lines 11-13: per-factor GAE, then the gated residual blend.

        Runs the critic over the rollout, computes ``A_{k,t}`` and ``G_{k,t}`` against its own
        factor heads (Eq. 8, Eq. 10), blends them into the caller's scalar advantage with the
        learned gate and mixture weights (Eq. 11), and normalises the result.

        Parameters
        ----------
        observations:
            Rollout observations ``s_0 … s_{T-1}``, shape ``(T, obs_dim)``.
        factor_rewards:
            ``r^factor_{k,t}``, shape ``(T, K)`` —
            :func:`~causalrl.agents.factored_advantage.factor_rewards` of the factor trace.
        scalar_advantages:
            The RL framework's own advantage ``A^scalar_t``, shape ``(T,)``.  Passing the
            framework's GAE (rather than this class's scalar head) is the intended path: the
            paper keeps the host PPO's critic and *adds* to it.
        gamma:
            Discount, shared with the scalar return (Table 6: 0.995).
        lam:
            Per-factor GAE ``lambda``; see
            :func:`~causalrl.agents.factored_advantage.factor_gae` for the paper ambiguity.
        bootstrap_values:
            ``V_k(s_T)`` for a truncated rollout, shape ``(K,)``; ``None`` means the episode
            ended.
        dones:
            Episode-termination flags, shape ``(T,)``.
        normalize:
            Standardise ``A_used`` to zero mean / unit variance across the rollout, as
            Algorithm 1 line 13 does.  Set ``False`` to inspect the raw blend.

        Returns
        -------
        CGFAAdvantages
            ``used`` feeds the surrogate; ``returns`` feeds :meth:`update`.
        """
        v_factor = self.values(observations)[1]
        adv_factor, ret_factor = factor_gae(
            factor_rewards,
            v_factor,
            gamma=gamma,
            lam=lam,
            bootstrap_values=bootstrap_values,
            dones=dones,
        )
        gate = self.gate(observations)
        weights = self.mixture_weights()
        used = blend_advantages(scalar_advantages, adv_factor, gate=gate, weights=weights)
        if normalize:
            std = float(used.std())
            used = (used - float(used.mean())) / (std + 1e-8)
        return CGFAAdvantages(
            used=used,
            scalar=np.asarray(scalar_advantages, dtype=np.float64),
            factor=adv_factor,
            returns=ret_factor,
            gate=gate,
            weights=weights,
        )

    # ------------------------------------------------------------------ losses

    def losses(
        self,
        observations: NDArray[np.float64],
        factor_returns: NDArray[np.float64],
        *,
        scalar_returns: NDArray[np.float64] | None = None,
        scm_effects: NDArray[np.float64] | None = None,
    ) -> CGFALosses:
        """The differentiable Eq. 13 components this class owns, as torch scalars.

        ``L_ppo`` and ``L_ent`` are the RL framework's; pass its surrogate to
        :meth:`objective` to assemble the whole of Eq. 13.

        Parameters
        ----------
        observations:
            Shape ``(T, obs_dim)``.
        factor_returns:
            ``G_{k,t}``, shape ``(T, K)`` — :attr:`CGFAAdvantages.returns`.
        scalar_returns:
            Bootstrapped scalar returns ``R_t``, shape ``(T,)``.  ``None`` (default) zeroes
            ``L_value``, which is correct when the RL framework owns the scalar critic.
        scm_effects:
            ``eps_{k,t}``, the SCM-predicted per-factor change, shape ``(T, K)``.  ``None``
            zeroes ``L_cal`` — the paper's "CGFA without calibration" ablation.

        Raises
        ------
        ValueError
            On any shape mismatch, or if ``factor_returns`` has a column count other than
            :attr:`n_factors`.
        """
        import torch

        obs = self._as_input(observations)
        t_steps = int(obs.shape[0])
        g_target = self._as_matrix(factor_returns, t_steps, "factor_returns")

        scalar, factors, gate = self._heads(obs)

        # Eq. 9 — per-factor MSE, averaged over the K heads.
        factor_loss = ((g_target - factors) ** 2).mean(dim=0).mean()

        # Alg. 1 line 17 — the scalar critic, only when the caller wants us to own it.
        if scalar_returns is None:
            value_loss = torch.zeros((), dtype=scalar.dtype)  # type: ignore[reportPrivateImportUsage]
        else:
            r_target = np.asarray(scalar_returns, dtype=np.float32)
            if r_target.shape != (t_steps,):
                raise ValueError(
                    f"scalar_returns must have shape ({t_steps},); got {r_target.shape}"
                )
            value_loss = ((torch.as_tensor(r_target) - scalar) ** 2).mean()  # type: ignore[reportPrivateImportUsage]

        # Eq. 12 — negative mean Pearson correlation between A_k and the SCM's effect.
        if scm_effects is None:
            calibration = torch.zeros((), dtype=scalar.dtype)  # type: ignore[reportPrivateImportUsage]
        else:
            eps = self._as_matrix(scm_effects, t_steps, "scm_effects")
            # A_k recomputed under the current parameters so Eq. 12 has a gradient path.
            adv = g_target - factors
            calibration = -self._pearson(adv, eps).mean()

        # Alg. 1 line 20 — Bernoulli entropy of the gate, in nats.
        eps_g = 1e-8
        gate_entropy = -(
            gate * torch.log(gate + eps_g)  # type: ignore[reportPrivateImportUsage]
            + (1.0 - gate) * torch.log(1.0 - gate + eps_g)  # type: ignore[reportPrivateImportUsage]
        ).mean()

        return CGFALosses(
            value=value_loss,
            factor=factor_loss,
            calibration=calibration,
            gate_entropy=gate_entropy,
        )

    def blend(
        self,
        observations: NDArray[np.float64],
        scalar_advantages: NDArray[np.float64],
        factor_advantages: NDArray[np.float64],
    ) -> Tensor:
        """Eq. 11 as a **differentiable** torch tensor, shape ``(T,)``.

        ``A_used = (1 - g(s_t)) A^scalar_t + g(s_t) sum_k w_k A_{k,t}``, evaluated under the
        current gate and mixture parameters so gradients reach ``g`` and ``beta``.

        .. note:: **Paper ambiguity, and why this method exists.** Algorithm 1 computes
           ``A_used`` once at line 12 — outside the epoch loop — and stores it; a surrogate
           built on that stored value gives ``g(s)`` and ``beta`` **no gradient at all**, so
           with the Table 6 default ``c_e = 0`` the only two parameters the paper calls
           "learnable" would never move.  Yet §E.3 speaks of the gate collapsing and logs its
           evolution, and Algorithm 1 line 22 updates ``beta``.  The consistent reading is
           that ``A_used`` is recomputed inside the minibatch loop; use this method to build
           the surrogate, and both parameters train.  Feeding
           :attr:`CGFAAdvantages.used` (a detached NumPy array) to the surrogate instead
           reproduces the literal, frozen-gate reading.

        Parameters
        ----------
        observations:
            Shape ``(T, obs_dim)``.
        scalar_advantages:
            ``A^scalar_t``, shape ``(T,)`` — treated as a constant, as PPO always does.
        factor_advantages:
            ``A_{k,t}``, shape ``(T, K)`` — likewise constant; the gradient path here is
            through ``g`` and ``w``, not through the factor critic.
        """
        import torch

        obs = self._as_input(observations)
        t_steps = int(obs.shape[0])
        a_scalar = np.asarray(scalar_advantages, dtype=np.float32)
        if a_scalar.shape != (t_steps,):
            raise ValueError(
                f"scalar_advantages must have shape ({t_steps},); got {a_scalar.shape}"
            )
        a_factor = self._as_matrix(factor_advantages, t_steps, "factor_advantages")
        _, _, gate = self._heads(obs)
        weights = torch.softmax(self._mixture_logits, dim=0)  # type: ignore[reportPrivateImportUsage]
        mixed = a_factor @ weights
        return (1.0 - gate) * torch.as_tensor(a_scalar) + gate * mixed  # type: ignore[reportPrivateImportUsage]

    def objective(
        self,
        losses: CGFALosses,
        *,
        policy_loss: Tensor | None = None,
    ) -> Tensor:
        """Assemble Eq. 13 from :meth:`losses` output and the framework's surrogate.

        ``L_CGFA = policy_loss + c_v L_value + c_f L_factor + c_c L_cal - c_e L_gate``.

        ``policy_loss`` is the caller's ``L_ppo - c_H L_ent`` (the clipped surrogate net of
        its entropy bonus), ideally built on :meth:`blend` so the gate and mixture logits
        receive gradient.  Omit it to optimise only the critic-side terms — a deviation from
        Algorithm 1 line 22, which steps actor and critic together; see the module docstring.
        """
        cfg = self._config
        total = (
            cfg.value_coef * losses.value
            + cfg.factor_coef * losses.factor
            + cfg.calibration_coef * losses.calibration
            - cfg.gate_entropy_coef * losses.gate_entropy
        )
        if policy_loss is not None:
            total = total + policy_loss
        return total

    # ------------------------------------------------------------------ training

    def update(
        self,
        observations: NDArray[np.float64],
        factor_returns: NDArray[np.float64],
        *,
        scalar_returns: NDArray[np.float64] | None = None,
        scm_effects: NDArray[np.float64] | None = None,
        epochs: int = 1,
        batch_size: int | None = None,
        seed: int | None = None,
    ) -> CGFAUpdateStats:
        """Algorithm 1 lines 14-22, restricted to the terms this class owns.

        Each minibatch step evaluates :meth:`losses`, assembles :meth:`objective`, clips the
        global gradient norm to :attr:`CGFACriticConfig.max_grad_norm`, and takes one Adam
        step over the trunk, the three heads, and ``beta``.

        Parameters
        ----------
        observations, factor_returns, scalar_returns, scm_effects:
            As in :meth:`losses`.
        epochs:
            Passes over the rollout (Table 6: 10).
        batch_size:
            Minibatch size (Table 6: 256).  ``None`` uses the whole rollout.
        seed:
            Seed for the minibatch shuffle.

        Returns
        -------
        CGFAUpdateStats
            Mean losses over the steps taken, plus the Algorithm 1 line 25 diagnostics
            (per-factor calibration correlation, credit share, gate distribution) recomputed
            on the full batch after the update.

        Raises
        ------
        ValueError
            If ``epochs`` is not positive, ``batch_size`` is not positive, or any array shape
            disagrees.
        """
        import torch

        if epochs <= 0:
            raise ValueError(f"epochs must be positive; got {epochs}")
        if batch_size is not None and batch_size <= 0:
            raise ValueError(f"batch_size must be positive; got {batch_size}")

        obs = np.asarray(observations, dtype=np.float64)
        if obs.ndim != 2:
            raise ValueError(f"observations must be 2-D (T, obs_dim); got shape {obs.shape}")
        t_steps = obs.shape[0]
        if t_steps == 0:
            raise ValueError("observations must contain at least one rollout step")
        size = t_steps if batch_size is None else min(batch_size, t_steps)
        rng = np.random.default_rng(seed)
        targets = np.asarray(factor_returns, dtype=np.float64)

        totals = {"total": 0.0, "value": 0.0, "factor": 0.0, "calibration": 0.0, "gate": 0.0}
        n_steps = 0
        for _ in range(epochs):
            order = rng.permutation(t_steps)
            for start in range(0, t_steps, size):
                idx = order[start : start + size]
                parts = self.losses(
                    obs[idx],
                    targets[idx],
                    scalar_returns=(None if scalar_returns is None else scalar_returns[idx]),
                    scm_effects=(None if scm_effects is None else scm_effects[idx]),
                )
                loss = self.objective(parts)
                self._optimizer.zero_grad()
                loss.backward()  # type: ignore[reportUnknownMemberType]
                torch.nn.utils.clip_grad_norm_(  # type: ignore[reportPrivateImportUsage]
                    self._module.parameters(), self._config.max_grad_norm
                )
                self._optimizer.step()  # type: ignore[reportUnknownMemberType]
                totals["total"] += float(loss.item())
                totals["value"] += float(parts.value.item())
                totals["factor"] += float(parts.factor.item())
                totals["calibration"] += float(parts.calibration.item())
                totals["gate"] += float(parts.gate_entropy.item())
                n_steps += 1

        return self._diagnostics(obs, factor_returns, scm_effects, totals, n_steps)

    # ------------------------------------------------------------------ internals

    def _as_matrix(self, array: NDArray[np.float64], t_steps: int, name: str) -> Tensor:
        import torch

        arr = np.asarray(array, dtype=np.float32)
        if arr.shape != (t_steps, self.n_factors):
            raise ValueError(
                f"{name} must have shape ({t_steps}, {self.n_factors}); got {arr.shape}"
            )
        return torch.as_tensor(arr)  # type: ignore[reportPrivateImportUsage]

    def _pearson(self, left: Tensor, right: Tensor) -> Tensor:
        """Column-wise Pearson correlation with the §E.4 standard-deviation clamp."""
        cfg = self._config
        lc = left - left.mean(dim=0, keepdim=True)
        rc = right - right.mean(dim=0, keepdim=True)
        cov = (lc * rc).mean(dim=0)
        std_l = lc.pow(2).mean(dim=0).sqrt().clamp(min=cfg.min_std)
        std_r = rc.pow(2).mean(dim=0).sqrt().clamp(min=cfg.min_std)
        return cov / (std_l * std_r + cfg.calibration_eps)

    def _diagnostics(
        self,
        observations: NDArray[np.float64],
        factor_returns: NDArray[np.float64],
        scm_effects: NDArray[np.float64] | None,
        totals: dict[str, float],
        n_steps: int,
    ) -> CGFAUpdateStats:
        """Algorithm 1 line 25 — calibration correlation, credit share, gate distribution."""
        weights = self.mixture_weights()
        gate = self.gate(observations)
        v_factor = self.values(observations)[1]
        adv = np.asarray(factor_returns, dtype=np.float64) - v_factor

        if scm_effects is None:
            corr = np.full(self.n_factors, np.nan)
        else:
            eps = np.asarray(scm_effects, dtype=np.float64)
            ac = adv - adv.mean(axis=0, keepdims=True)
            ec = eps - eps.mean(axis=0, keepdims=True)
            denom = np.sqrt((ac**2).mean(axis=0)) * np.sqrt((ec**2).mean(axis=0))
            with np.errstate(invalid="ignore", divide="ignore"):
                corr = np.where(denom > 0.0, (ac * ec).mean(axis=0) / denom, np.nan)

        contrib = np.abs(weights[np.newaxis, :] * adv)
        row_total = contrib.sum(axis=1)
        live = row_total > 0.0
        if bool(live.any()):
            credit_share = (contrib[live] / row_total[live, np.newaxis]).mean(axis=0)
        else:
            credit_share = np.full(self.n_factors, np.nan)

        scale = float(max(n_steps, 1))
        return CGFAUpdateStats(
            total=totals["total"] / scale,
            value=totals["value"] / scale,
            factor=totals["factor"] / scale,
            calibration=totals["calibration"] / scale,
            gate_entropy=totals["gate"] / scale,
            mixture_weights=weights,
            factor_correlation=np.asarray(corr, dtype=np.float64),
            credit_share=np.asarray(credit_share, dtype=np.float64),
            gate_mean=float(gate.mean()),
        )

factor_nodes property

The ordered SCM parent names, one per value head.

n_factors property

K — the number of causal factors, hence of value heads.

config property

The architecture / loss-coefficient bundle this critic was built with.

module property

The underlying torch.nn.Module — for state_dict, .to(device), saving.

parameters()

Every learnable tensor: trunk, three heads, and the mixture logits beta.

Hand these to a joint optimiser alongside the actor's parameters when following Algorithm 1 line 22 literally (one Adam step on the whole of Eq. 13).

Source code in src/causalrl/agents/cgfa_critic.py

def parameters(self) -> Iterator[Tensor]:
    """Every learnable tensor: trunk, three heads, and the mixture logits ``beta``.

    Hand these to a joint optimiser alongside the actor's parameters when following
    Algorithm 1 line 22 literally (one Adam step on the whole of Eq. 13).
    """
    return self._module.parameters()  # type: ignore[reportUnknownMemberType,no-any-return]

values(observations)

(V(s_t), V_k(s_t)) for a batch of observations, shapes (T,) and (T, K).

Source code in src/causalrl/agents/cgfa_critic.py

def values(
    self, observations: NDArray[np.float64]
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """``(V(s_t), V_k(s_t))`` for a batch of observations, shapes ``(T,)`` and ``(T, K)``."""
    import torch

    with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
        scalar, factors, _ = self._heads(self._as_input(observations))
    return (
        scalar.numpy().astype(np.float64),
        factors.numpy().astype(np.float64),
    )

gate(observations)

The state-conditional residual gate g(s_t) in (0, 1), shape (T,) (§E.3).

Source code in src/causalrl/agents/cgfa_critic.py

def gate(self, observations: NDArray[np.float64]) -> NDArray[np.float64]:
    """The state-conditional residual gate ``g(s_t) in (0, 1)``, shape ``(T,)`` (§E.3)."""
    import torch

    with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
        _, _, gate = self._heads(self._as_input(observations))
    return gate.numpy().astype(np.float64)

mixture_weights()

w = softmax(beta), shape (K,) — the learned per-factor mixture (§E.1).

Source code in src/causalrl/agents/cgfa_critic.py

def mixture_weights(self) -> NDArray[np.float64]:
    """``w = softmax(beta)``, shape ``(K,)`` — the learned per-factor mixture (§E.1)."""
    import torch

    with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
        w = torch.softmax(self._mixture_logits, dim=0)  # type: ignore[reportPrivateImportUsage]
    return w.numpy().astype(np.float64)

advantages(observations, factor_rewards, scalar_advantages, *, gamma=0.995, lam=1.0, bootstrap_values=None, dones=None, normalize=True)

Algorithm 1 lines 11-13: per-factor GAE, then the gated residual blend.

Runs the critic over the rollout, computes A_{k,t} and G_{k,t} against its own factor heads (Eq. 8, Eq. 10), blends them into the caller's scalar advantage with the learned gate and mixture weights (Eq. 11), and normalises the result.

Parameters

observations: Rollout observations s_0 … s_{T-1}, shape (T, obs_dim). factor_rewards: r^factor_{k,t}, shape (T, K) — :func:~causalrl.agents.factored_advantage.factor_rewards of the factor trace. scalar_advantages: The RL framework's own advantage A^scalar_t, shape (T,). Passing the framework's GAE (rather than this class's scalar head) is the intended path: the paper keeps the host PPO's critic and adds to it. gamma: Discount, shared with the scalar return (Table 6: 0.995). lam: Per-factor GAE lambda; see :func:~causalrl.agents.factored_advantage.factor_gae for the paper ambiguity. bootstrap_values: V_k(s_T) for a truncated rollout, shape (K,); None means the episode ended. dones: Episode-termination flags, shape (T,). normalize: Standardise A_used to zero mean / unit variance across the rollout, as Algorithm 1 line 13 does. Set False to inspect the raw blend.

Returns

CGFAAdvantages used feeds the surrogate; returns feeds :meth:update.

Source code in src/causalrl/agents/cgfa_critic.py

def advantages(
    self,
    observations: NDArray[np.float64],
    factor_rewards: NDArray[np.float64],
    scalar_advantages: NDArray[np.float64],
    *,
    gamma: float = 0.995,
    lam: float = 1.0,
    bootstrap_values: NDArray[np.float64] | None = None,
    dones: NDArray[np.bool_] | None = None,
    normalize: bool = True,
) -> CGFAAdvantages:
    """Algorithm 1 lines 11-13: per-factor GAE, then the gated residual blend.

    Runs the critic over the rollout, computes ``A_{k,t}`` and ``G_{k,t}`` against its own
    factor heads (Eq. 8, Eq. 10), blends them into the caller's scalar advantage with the
    learned gate and mixture weights (Eq. 11), and normalises the result.

    Parameters
    ----------
    observations:
        Rollout observations ``s_0 … s_{T-1}``, shape ``(T, obs_dim)``.
    factor_rewards:
        ``r^factor_{k,t}``, shape ``(T, K)`` —
        :func:`~causalrl.agents.factored_advantage.factor_rewards` of the factor trace.
    scalar_advantages:
        The RL framework's own advantage ``A^scalar_t``, shape ``(T,)``.  Passing the
        framework's GAE (rather than this class's scalar head) is the intended path: the
        paper keeps the host PPO's critic and *adds* to it.
    gamma:
        Discount, shared with the scalar return (Table 6: 0.995).
    lam:
        Per-factor GAE ``lambda``; see
        :func:`~causalrl.agents.factored_advantage.factor_gae` for the paper ambiguity.
    bootstrap_values:
        ``V_k(s_T)`` for a truncated rollout, shape ``(K,)``; ``None`` means the episode
        ended.
    dones:
        Episode-termination flags, shape ``(T,)``.
    normalize:
        Standardise ``A_used`` to zero mean / unit variance across the rollout, as
        Algorithm 1 line 13 does.  Set ``False`` to inspect the raw blend.

    Returns
    -------
    CGFAAdvantages
        ``used`` feeds the surrogate; ``returns`` feeds :meth:`update`.
    """
    v_factor = self.values(observations)[1]
    adv_factor, ret_factor = factor_gae(
        factor_rewards,
        v_factor,
        gamma=gamma,
        lam=lam,
        bootstrap_values=bootstrap_values,
        dones=dones,
    )
    gate = self.gate(observations)
    weights = self.mixture_weights()
    used = blend_advantages(scalar_advantages, adv_factor, gate=gate, weights=weights)
    if normalize:
        std = float(used.std())
        used = (used - float(used.mean())) / (std + 1e-8)
    return CGFAAdvantages(
        used=used,
        scalar=np.asarray(scalar_advantages, dtype=np.float64),
        factor=adv_factor,
        returns=ret_factor,
        gate=gate,
        weights=weights,
    )

losses(observations, factor_returns, *, scalar_returns=None, scm_effects=None)

The differentiable Eq. 13 components this class owns, as torch scalars.

L_ppo and L_ent are the RL framework's; pass its surrogate to :meth:objective to assemble the whole of Eq. 13.

Parameters

observations: Shape (T, obs_dim). factor_returns: G_{k,t}, shape (T, K) — :attr:CGFAAdvantages.returns. scalar_returns: Bootstrapped scalar returns R_t, shape (T,). None (default) zeroes L_value, which is correct when the RL framework owns the scalar critic. scm_effects: eps_{k,t}, the SCM-predicted per-factor change, shape (T, K). None zeroes L_cal — the paper's "CGFA without calibration" ablation.

Raises

ValueError On any shape mismatch, or if factor_returns has a column count other than :attr:n_factors.

Source code in src/causalrl/agents/cgfa_critic.py

def losses(
    self,
    observations: NDArray[np.float64],
    factor_returns: NDArray[np.float64],
    *,
    scalar_returns: NDArray[np.float64] | None = None,
    scm_effects: NDArray[np.float64] | None = None,
) -> CGFALosses:
    """The differentiable Eq. 13 components this class owns, as torch scalars.

    ``L_ppo`` and ``L_ent`` are the RL framework's; pass its surrogate to
    :meth:`objective` to assemble the whole of Eq. 13.

    Parameters
    ----------
    observations:
        Shape ``(T, obs_dim)``.
    factor_returns:
        ``G_{k,t}``, shape ``(T, K)`` — :attr:`CGFAAdvantages.returns`.
    scalar_returns:
        Bootstrapped scalar returns ``R_t``, shape ``(T,)``.  ``None`` (default) zeroes
        ``L_value``, which is correct when the RL framework owns the scalar critic.
    scm_effects:
        ``eps_{k,t}``, the SCM-predicted per-factor change, shape ``(T, K)``.  ``None``
        zeroes ``L_cal`` — the paper's "CGFA without calibration" ablation.

    Raises
    ------
    ValueError
        On any shape mismatch, or if ``factor_returns`` has a column count other than
        :attr:`n_factors`.
    """
    import torch

    obs = self._as_input(observations)
    t_steps = int(obs.shape[0])
    g_target = self._as_matrix(factor_returns, t_steps, "factor_returns")

    scalar, factors, gate = self._heads(obs)

    # Eq. 9 — per-factor MSE, averaged over the K heads.
    factor_loss = ((g_target - factors) ** 2).mean(dim=0).mean()

    # Alg. 1 line 17 — the scalar critic, only when the caller wants us to own it.
    if scalar_returns is None:
        value_loss = torch.zeros((), dtype=scalar.dtype)  # type: ignore[reportPrivateImportUsage]
    else:
        r_target = np.asarray(scalar_returns, dtype=np.float32)
        if r_target.shape != (t_steps,):
            raise ValueError(
                f"scalar_returns must have shape ({t_steps},); got {r_target.shape}"
            )
        value_loss = ((torch.as_tensor(r_target) - scalar) ** 2).mean()  # type: ignore[reportPrivateImportUsage]

    # Eq. 12 — negative mean Pearson correlation between A_k and the SCM's effect.
    if scm_effects is None:
        calibration = torch.zeros((), dtype=scalar.dtype)  # type: ignore[reportPrivateImportUsage]
    else:
        eps = self._as_matrix(scm_effects, t_steps, "scm_effects")
        # A_k recomputed under the current parameters so Eq. 12 has a gradient path.
        adv = g_target - factors
        calibration = -self._pearson(adv, eps).mean()

    # Alg. 1 line 20 — Bernoulli entropy of the gate, in nats.
    eps_g = 1e-8
    gate_entropy = -(
        gate * torch.log(gate + eps_g)  # type: ignore[reportPrivateImportUsage]
        + (1.0 - gate) * torch.log(1.0 - gate + eps_g)  # type: ignore[reportPrivateImportUsage]
    ).mean()

    return CGFALosses(
        value=value_loss,
        factor=factor_loss,
        calibration=calibration,
        gate_entropy=gate_entropy,
    )

blend(observations, scalar_advantages, factor_advantages)

Eq. 11 as a differentiable torch tensor, shape (T,).

A_used = (1 - g(s_t)) A^scalar_t + g(s_t) sum_k w_k A_{k,t}, evaluated under the current gate and mixture parameters so gradients reach g and beta.

.. note:: Paper ambiguity, and why this method exists. Algorithm 1 computes A_used once at line 12 — outside the epoch loop — and stores it; a surrogate built on that stored value gives g(s) and beta no gradient at all, so with the Table 6 default c_e = 0 the only two parameters the paper calls "learnable" would never move. Yet §E.3 speaks of the gate collapsing and logs its evolution, and Algorithm 1 line 22 updates beta. The consistent reading is that A_used is recomputed inside the minibatch loop; use this method to build the surrogate, and both parameters train. Feeding :attr:CGFAAdvantages.used (a detached NumPy array) to the surrogate instead reproduces the literal, frozen-gate reading.

Parameters

observations: Shape (T, obs_dim). scalar_advantages: A^scalar_t, shape (T,) — treated as a constant, as PPO always does. factor_advantages: A_{k,t}, shape (T, K) — likewise constant; the gradient path here is through g and w, not through the factor critic.

Source code in src/causalrl/agents/cgfa_critic.py

def blend(
    self,
    observations: NDArray[np.float64],
    scalar_advantages: NDArray[np.float64],
    factor_advantages: NDArray[np.float64],
) -> Tensor:
    """Eq. 11 as a **differentiable** torch tensor, shape ``(T,)``.

    ``A_used = (1 - g(s_t)) A^scalar_t + g(s_t) sum_k w_k A_{k,t}``, evaluated under the
    current gate and mixture parameters so gradients reach ``g`` and ``beta``.

    .. note:: **Paper ambiguity, and why this method exists.** Algorithm 1 computes
       ``A_used`` once at line 12 — outside the epoch loop — and stores it; a surrogate
       built on that stored value gives ``g(s)`` and ``beta`` **no gradient at all**, so
       with the Table 6 default ``c_e = 0`` the only two parameters the paper calls
       "learnable" would never move.  Yet §E.3 speaks of the gate collapsing and logs its
       evolution, and Algorithm 1 line 22 updates ``beta``.  The consistent reading is
       that ``A_used`` is recomputed inside the minibatch loop; use this method to build
       the surrogate, and both parameters train.  Feeding
       :attr:`CGFAAdvantages.used` (a detached NumPy array) to the surrogate instead
       reproduces the literal, frozen-gate reading.

    Parameters
    ----------
    observations:
        Shape ``(T, obs_dim)``.
    scalar_advantages:
        ``A^scalar_t``, shape ``(T,)`` — treated as a constant, as PPO always does.
    factor_advantages:
        ``A_{k,t}``, shape ``(T, K)`` — likewise constant; the gradient path here is
        through ``g`` and ``w``, not through the factor critic.
    """
    import torch

    obs = self._as_input(observations)
    t_steps = int(obs.shape[0])
    a_scalar = np.asarray(scalar_advantages, dtype=np.float32)
    if a_scalar.shape != (t_steps,):
        raise ValueError(
            f"scalar_advantages must have shape ({t_steps},); got {a_scalar.shape}"
        )
    a_factor = self._as_matrix(factor_advantages, t_steps, "factor_advantages")
    _, _, gate = self._heads(obs)
    weights = torch.softmax(self._mixture_logits, dim=0)  # type: ignore[reportPrivateImportUsage]
    mixed = a_factor @ weights
    return (1.0 - gate) * torch.as_tensor(a_scalar) + gate * mixed  # type: ignore[reportPrivateImportUsage]

objective(losses, *, policy_loss=None)

Assemble Eq. 13 from :meth:losses output and the framework's surrogate.

L_CGFA = policy_loss + c_v L_value + c_f L_factor + c_c L_cal - c_e L_gate.

policy_loss is the caller's L_ppo - c_H L_ent (the clipped surrogate net of its entropy bonus), ideally built on :meth:blend so the gate and mixture logits receive gradient. Omit it to optimise only the critic-side terms — a deviation from Algorithm 1 line 22, which steps actor and critic together; see the module docstring.

Source code in src/causalrl/agents/cgfa_critic.py

def objective(
    self,
    losses: CGFALosses,
    *,
    policy_loss: Tensor | None = None,
) -> Tensor:
    """Assemble Eq. 13 from :meth:`losses` output and the framework's surrogate.

    ``L_CGFA = policy_loss + c_v L_value + c_f L_factor + c_c L_cal - c_e L_gate``.

    ``policy_loss`` is the caller's ``L_ppo - c_H L_ent`` (the clipped surrogate net of
    its entropy bonus), ideally built on :meth:`blend` so the gate and mixture logits
    receive gradient.  Omit it to optimise only the critic-side terms — a deviation from
    Algorithm 1 line 22, which steps actor and critic together; see the module docstring.
    """
    cfg = self._config
    total = (
        cfg.value_coef * losses.value
        + cfg.factor_coef * losses.factor
        + cfg.calibration_coef * losses.calibration
        - cfg.gate_entropy_coef * losses.gate_entropy
    )
    if policy_loss is not None:
        total = total + policy_loss
    return total

update(observations, factor_returns, *, scalar_returns=None, scm_effects=None, epochs=1, batch_size=None, seed=None)

Algorithm 1 lines 14-22, restricted to the terms this class owns.

Each minibatch step evaluates :meth:losses, assembles :meth:objective, clips the global gradient norm to :attr:CGFACriticConfig.max_grad_norm, and takes one Adam step over the trunk, the three heads, and beta.

Parameters

observations, factor_returns, scalar_returns, scm_effects: As in :meth:losses. epochs: Passes over the rollout (Table 6: 10). batch_size: Minibatch size (Table 6: 256). None uses the whole rollout. seed: Seed for the minibatch shuffle.

Returns

CGFAUpdateStats Mean losses over the steps taken, plus the Algorithm 1 line 25 diagnostics (per-factor calibration correlation, credit share, gate distribution) recomputed on the full batch after the update.

Raises

ValueError If epochs is not positive, batch_size is not positive, or any array shape disagrees.

Source code in src/causalrl/agents/cgfa_critic.py

def update(
    self,
    observations: NDArray[np.float64],
    factor_returns: NDArray[np.float64],
    *,
    scalar_returns: NDArray[np.float64] | None = None,
    scm_effects: NDArray[np.float64] | None = None,
    epochs: int = 1,
    batch_size: int | None = None,
    seed: int | None = None,
) -> CGFAUpdateStats:
    """Algorithm 1 lines 14-22, restricted to the terms this class owns.

    Each minibatch step evaluates :meth:`losses`, assembles :meth:`objective`, clips the
    global gradient norm to :attr:`CGFACriticConfig.max_grad_norm`, and takes one Adam
    step over the trunk, the three heads, and ``beta``.

    Parameters
    ----------
    observations, factor_returns, scalar_returns, scm_effects:
        As in :meth:`losses`.
    epochs:
        Passes over the rollout (Table 6: 10).
    batch_size:
        Minibatch size (Table 6: 256).  ``None`` uses the whole rollout.
    seed:
        Seed for the minibatch shuffle.

    Returns
    -------
    CGFAUpdateStats
        Mean losses over the steps taken, plus the Algorithm 1 line 25 diagnostics
        (per-factor calibration correlation, credit share, gate distribution) recomputed
        on the full batch after the update.

    Raises
    ------
    ValueError
        If ``epochs`` is not positive, ``batch_size`` is not positive, or any array shape
        disagrees.
    """
    import torch

    if epochs <= 0:
        raise ValueError(f"epochs must be positive; got {epochs}")
    if batch_size is not None and batch_size <= 0:
        raise ValueError(f"batch_size must be positive; got {batch_size}")

    obs = np.asarray(observations, dtype=np.float64)
    if obs.ndim != 2:
        raise ValueError(f"observations must be 2-D (T, obs_dim); got shape {obs.shape}")
    t_steps = obs.shape[0]
    if t_steps == 0:
        raise ValueError("observations must contain at least one rollout step")
    size = t_steps if batch_size is None else min(batch_size, t_steps)
    rng = np.random.default_rng(seed)
    targets = np.asarray(factor_returns, dtype=np.float64)

    totals = {"total": 0.0, "value": 0.0, "factor": 0.0, "calibration": 0.0, "gate": 0.0}
    n_steps = 0
    for _ in range(epochs):
        order = rng.permutation(t_steps)
        for start in range(0, t_steps, size):
            idx = order[start : start + size]
            parts = self.losses(
                obs[idx],
                targets[idx],
                scalar_returns=(None if scalar_returns is None else scalar_returns[idx]),
                scm_effects=(None if scm_effects is None else scm_effects[idx]),
            )
            loss = self.objective(parts)
            self._optimizer.zero_grad()
            loss.backward()  # type: ignore[reportUnknownMemberType]
            torch.nn.utils.clip_grad_norm_(  # type: ignore[reportPrivateImportUsage]
                self._module.parameters(), self._config.max_grad_norm
            )
            self._optimizer.step()  # type: ignore[reportUnknownMemberType]
            totals["total"] += float(loss.item())
            totals["value"] += float(parts.value.item())
            totals["factor"] += float(parts.factor.item())
            totals["calibration"] += float(parts.calibration.item())
            totals["gate"] += float(parts.gate_entropy.item())
            n_steps += 1

    return self._diagnostics(obs, factor_returns, scm_effects, totals, n_steps)

causalrl.agents.cgfa_critic.CGFACriticConfig dataclass

Architecture and loss coefficients for :class:FactoredCritic (arXiv:2605.06066).

Defaults follow the paper's Table 6 where it states a value, and are flagged in the module docstring where it does not.

Parameters

hidden: Widths of the shared critic trunk. Table 6 uses (512, 256) for a 3,077-dimensional observation; the default here is small enough for toy problems. gate_hidden: Hidden width of the state-conditional gate MLP (Table 6: 32). gate_init: Value of g(s) at initialisation, uniform over states (Table 6: 0.5). The gate head's output layer is zero-initialised with the matching bias, so the gate starts state-independent and differentiates as it trains. learning_rate: Adam learning rate for the critic parameters (Table 6 backbone: 3e-4). value_coef: c_v — weight on the scalar value loss. Only active when scalar_returns are supplied to :meth:FactoredCritic.update; a caller whose RL framework owns the scalar critic leaves them out and this term is exactly zero. factor_coef: c_f — weight on the per-factor MSE, Eq. 9 (Table 6: 0.5). calibration_coef: c_c — weight on the intervention-calibration loss, Eq. 12 (Table 6: 0.1). gate_entropy_coef: c_e — weight on the Bernoulli gate-entropy bonus (Table 6 default: 0.0, swept in the ablation). Eq. 13 subtracts it, so a positive value rewards an undecided gate. calibration_eps: delta — the stability constant in Eq. 12's denominator. min_std: Lower clamp on the per-factor standard deviations in Eq. 12 (§E.4 requires a clamp but gives no value). max_grad_norm: Global gradient-norm clip, Algorithm 1 line 22 (Table 6: 0.5).

Source code in src/causalrl/agents/cgfa_critic.py

@dataclass(frozen=True)
class CGFACriticConfig:
    """Architecture and loss coefficients for :class:`FactoredCritic` (arXiv:2605.06066).

    Defaults follow the paper's Table 6 where it states a value, and are flagged in the
    module docstring where it does not.

    Parameters
    ----------
    hidden:
        Widths of the shared critic trunk.  Table 6 uses ``(512, 256)`` for a
        3,077-dimensional observation; the default here is small enough for toy problems.
    gate_hidden:
        Hidden width of the state-conditional gate MLP (Table 6: 32).
    gate_init:
        Value of ``g(s)`` at initialisation, uniform over states (Table 6: 0.5).  The gate
        head's output layer is zero-initialised with the matching bias, so the gate starts
        state-independent and differentiates as it trains.
    learning_rate:
        Adam learning rate for the critic parameters (Table 6 backbone: ``3e-4``).
    value_coef:
        ``c_v`` — weight on the scalar value loss.  Only active when ``scalar_returns`` are
        supplied to :meth:`FactoredCritic.update`; a caller whose RL framework owns the
        scalar critic leaves them out and this term is exactly zero.
    factor_coef:
        ``c_f`` — weight on the per-factor MSE, Eq. 9 (Table 6: 0.5).
    calibration_coef:
        ``c_c`` — weight on the intervention-calibration loss, Eq. 12 (Table 6: 0.1).
    gate_entropy_coef:
        ``c_e`` — weight on the Bernoulli gate-entropy bonus (Table 6 default: 0.0, swept in
        the ablation).  Eq. 13 *subtracts* it, so a positive value rewards an undecided gate.
    calibration_eps:
        ``delta`` — the stability constant in Eq. 12's denominator.
    min_std:
        Lower clamp on the per-factor standard deviations in Eq. 12 (§E.4 requires a clamp
        but gives no value).
    max_grad_norm:
        Global gradient-norm clip, Algorithm 1 line 22 (Table 6: 0.5).
    """

    hidden: tuple[int, ...] = (64, 64)
    gate_hidden: int = 32
    gate_init: float = 0.5
    learning_rate: float = 3e-4
    value_coef: float = 0.5
    factor_coef: float = 0.5
    calibration_coef: float = 0.1
    gate_entropy_coef: float = 0.0
    calibration_eps: float = 1e-8
    min_std: float = 1e-3
    max_grad_norm: float = 0.5

    def __post_init__(self) -> None:
        if not self.hidden:
            raise ValueError("hidden must contain at least one layer width")
        if not 0.0 < self.gate_init < 1.0:
            raise ValueError(f"gate_init must lie strictly in (0, 1); got {self.gate_init}")

causalrl.agents.cgfa_critic.CGFAAdvantages dataclass

One rollout's worth of CGFA advantage arithmetic (Algorithm 1 lines 11-13).

Attributes

used: A_used, shape (T,) — the advantage to hand to the policy-gradient surrogate. scalar: The scalar-critic advantage that was blended in, shape (T,). factor: Per-factor advantages A_{k,t}, shape (T, K) (Eq. 10). returns: Per-factor returns G_{k,t}, shape (T, K) (Eq. 8) — the regression targets :meth:FactoredCritic.update expects. gate: g(s_t), shape (T,). weights: w = softmax(beta), shape (K,).

Source code in src/causalrl/agents/cgfa_critic.py

@dataclass(frozen=True)
class CGFAAdvantages:
    """One rollout's worth of CGFA advantage arithmetic (Algorithm 1 lines 11-13).

    Attributes
    ----------
    used:
        ``A_used``, shape ``(T,)`` — the advantage to hand to the policy-gradient surrogate.
    scalar:
        The scalar-critic advantage that was blended in, shape ``(T,)``.
    factor:
        Per-factor advantages ``A_{k,t}``, shape ``(T, K)`` (Eq. 10).
    returns:
        Per-factor returns ``G_{k,t}``, shape ``(T, K)`` (Eq. 8) — the regression targets
        :meth:`FactoredCritic.update` expects.
    gate:
        ``g(s_t)``, shape ``(T,)``.
    weights:
        ``w = softmax(beta)``, shape ``(K,)``.
    """

    used: NDArray[np.float64]
    scalar: NDArray[np.float64]
    factor: NDArray[np.float64]
    returns: NDArray[np.float64]
    gate: NDArray[np.float64]
    weights: NDArray[np.float64]

causalrl.agents.cgfa_critic.CGFALosses dataclass

The differentiable Eq. 13 terms :class:FactoredCritic owns, as torch scalars.

Attributes

value: L_value — the scalar critic's MSE (Algorithm 1 line 17). Exactly zero when the caller's RL framework owns the scalar critic and supplies no scalar_returns. factor: L_factor — the per-factor MSE of Eq. 9, averaged over the K heads. calibration: L_cal — the negative mean Pearson correlation of Eq. 12. Exactly zero when no SCM-predicted effects are supplied (the paper's "without calibration" ablation). gate_entropy: L_gate — the mean Bernoulli entropy of g(s_t) (Algorithm 1 line 20).

Source code in src/causalrl/agents/cgfa_critic.py

@dataclass(frozen=True)
class CGFALosses:
    """The differentiable Eq. 13 terms :class:`FactoredCritic` owns, as torch scalars.

    Attributes
    ----------
    value:
        ``L_value`` — the scalar critic's MSE (Algorithm 1 line 17).  Exactly zero when the
        caller's RL framework owns the scalar critic and supplies no ``scalar_returns``.
    factor:
        ``L_factor`` — the per-factor MSE of Eq. 9, averaged over the ``K`` heads.
    calibration:
        ``L_cal`` — the negative mean Pearson correlation of Eq. 12.  Exactly zero when no
        SCM-predicted effects are supplied (the paper's "without calibration" ablation).
    gate_entropy:
        ``L_gate`` — the mean Bernoulli entropy of ``g(s_t)`` (Algorithm 1 line 20).
    """

    value: Tensor
    factor: Tensor
    calibration: Tensor
    gate_entropy: Tensor

causalrl.agents.cgfa_critic.CGFAUpdateStats dataclass

Loss values and the per-factor diagnostics of Algorithm 1 line 25.

Attributes

total: Mean of the optimised objective over the update's minibatch steps. value, factor, calibration, gate_entropy: Mean of each Eq. 13 component, before its coefficient. mixture_weights: softmax(beta) after the update, shape (K,). factor_correlation: Per-factor Pearson correlation corr(A_{k,t}, eps_{k,t}) after the update, shape (K,); all-nan when no SCM effects were supplied. This is the paper's headline calibration diagnostic: it says whether each learned factor advantage moves in the same direction as the SCM's predicted intervention effect. credit_share: mean_t |w_k A_{k,t}| / sum_j |w_j A_{j,t}|, shape (K,) — which causal factors dominate the policy update. gate_mean: Mean g(s_t) after the update. Near 0 means the gate has collapsed to vanilla PPO; near 1 means the update rides entirely on the factor-aligned advantage.

Source code in src/causalrl/agents/cgfa_critic.py

@dataclass(frozen=True)
class CGFAUpdateStats:
    """Loss values and the per-factor diagnostics of Algorithm 1 line 25.

    Attributes
    ----------
    total:
        Mean of the optimised objective over the update's minibatch steps.
    value, factor, calibration, gate_entropy:
        Mean of each Eq. 13 component, *before* its coefficient.
    mixture_weights:
        ``softmax(beta)`` after the update, shape ``(K,)``.
    factor_correlation:
        Per-factor Pearson correlation ``corr(A_{k,t}, eps_{k,t})`` after the update, shape
        ``(K,)``; all-``nan`` when no SCM effects were supplied.  This is the paper's
        headline calibration diagnostic: it says whether each learned factor advantage moves
        in the same direction as the SCM's predicted intervention effect.
    credit_share:
        ``mean_t |w_k A_{k,t}| / sum_j |w_j A_{j,t}|``, shape ``(K,)`` — which causal factors
        dominate the policy update.
    gate_mean:
        Mean ``g(s_t)`` after the update.  Near 0 means the gate has collapsed to vanilla
        PPO; near 1 means the update rides entirely on the factor-aligned advantage.
    """

    total: float
    value: float
    factor: float
    calibration: float
    gate_entropy: float
    mixture_weights: NDArray[np.float64]
    factor_correlation: NDArray[np.float64]
    credit_share: NDArray[np.float64]
    gate_mean: float

Intervention Sets

causalrl.identification.intervention_sets.pomis(graph, reward, manipulable=None)

All POMISs for reward: a deduplicated, canonically sorted list of frozensets.

frozenset() (the observational regime) appears when it is possibly optimal. When manipulable is given, only those variables may be intervened on: by r40's Theorem 4 this is the unconstrained POMIS of the latent projection onto manipulable | {reward}.

Source code in src/causalrl/identification/intervention_sets.py

def pomis(
    graph: CausalGraph, reward: str, manipulable: Iterable[str] | None = None
) -> list[frozenset[str]]:
    """All POMISs for `reward`: a deduplicated, canonically sorted list of frozensets.

    ``frozenset()`` (the observational regime) appears when it is possibly optimal. When
    ``manipulable`` is given, only those variables may be intervened on: by r40's Theorem 4
    this is the unconstrained POMIS of the latent projection onto ``manipulable | {reward}``.
    """
    if manipulable is not None:
        graph = graph.latent_projection(set(manipulable) | {reward})
    graph = graph.induced_subgraph(graph.ancestors(reward))
    territory, border = _muct_ib(graph, reward)
    sub = graph.do_mutilate(border).induced_subgraph(territory | border)
    ws = [w for w in _backward_order(sub) if w in (territory - {reward})]
    result = _sub_pomis(sub, reward, ws, frozenset()) | {border}
    return _canonical(result)

causalrl.identification.intervention_sets.minimal_intervention_sets(graph, reward, manipulable=None)

All MISs for reward: a deduplicated, canonically sorted list of frozensets.

When manipulable is given, the result is filtered to sets that avoid the non-manipulable variables (r40: a constrained MIS is just the filtered unconstrained MIS).

Source code in src/causalrl/identification/intervention_sets.py

def minimal_intervention_sets(
    graph: CausalGraph, reward: str, manipulable: Iterable[str] | None = None
) -> list[frozenset[str]]:
    """All MISs for `reward`: a deduplicated, canonically sorted list of frozensets.

    When ``manipulable`` is given, the result is filtered to sets that avoid the
    non-manipulable variables (r40: a constrained MIS is just the filtered unconstrained MIS).
    """
    graph = graph.induced_subgraph(graph.ancestors(reward))
    ws = [w for w in _backward_order(graph) if w != reward]
    result = _canonical(_sub_miss(graph, reward, frozenset(), ws))
    if manipulable is None:
        return result
    allowed = set(manipulable)
    return [s for s in result if s <= allowed]

causalrl.identification.intervention_sets.requires_experiment(graph, treatment, outcome)

Whether learning P(outcome | do(treatment)) requires experimentation (Task 2, "when").

Returns True exactly when the effect is not identifiable from observational (L1) data, so an online experiment (an L2 intervention) is necessary, and False when offline data already suffices. This is the "when to intervene" companion to POMIS's "where": :func:pomis narrows which intervention sets could be optimal, while this decides whether you must intervene at all. Delegates to the complete ID algorithm (:func:causalrl.identification.id_algorithm.is_identifiable_effect).

Source code in src/causalrl/identification/intervention_sets.py

def requires_experiment(graph: CausalGraph, treatment: Iterable[str], outcome: str) -> bool:
    """Whether learning ``P(outcome | do(treatment))`` *requires* experimentation (Task 2, "when").

    Returns ``True`` exactly when the effect is **not** identifiable from observational (L1) data,
    so an online experiment (an L2 intervention) is necessary, and ``False`` when offline data
    already suffices. This is the "when to intervene" companion to POMIS's "where": :func:`pomis`
    narrows *which* intervention sets could be optimal, while this decides *whether* you must
    intervene at all. Delegates to the complete ID algorithm
    (:func:`causalrl.identification.id_algorithm.is_identifiable_effect`).
    """
    return not is_identifiable_effect(graph, treatment, {outcome})

causalrl.identification.intervention_sets.AdmissibleInterventions

POMIS for a graph, recomputed for the variables each context actually permits.

:func:pomis answers "where could it be optimal to intervene" once, at design time. A live agent needs the same answer per decision, because feasibility moves with the state: a lever can be unavailable this step and available the next. This holds the graph fixed and takes the manipulable set as the varying input.

Restricting the manipulable set is not a filter of the unconstrained POMIS. It is the unconstrained POMIS of the latent projection onto the manipulable variables (Lee & Bareinboim, AAAI 2019, Theorem 4 — the r40 result :func:pomis already implements). Projecting away a variable turns paths through it into bidirected edges, which can change the territory and border and so introduce possibly-optimal sets that the unconstrained enumeration never listed. Dropping the sets that mention an infeasible variable would therefore lose optimality, not merely prune; that is why :meth:sets re-enters :func:pomis rather than filtering a cached list.

Because re-entering is not free and the manipulable set usually changes rarely, results are memoised on it. cache_size bounds the memo; the oldest entry is evicted when it is exceeded, which suits the common pattern of a feasible set that oscillates among a handful of configurations.

Source code in src/causalrl/identification/intervention_sets.py

class AdmissibleInterventions:
    """POMIS for a graph, recomputed for the variables each context actually permits.

    :func:`pomis` answers "where could it be optimal to intervene" once, at design time. A live
    agent needs the same answer per decision, because feasibility moves with the state: a lever
    can be unavailable this step and available the next. This holds the graph fixed and takes the
    manipulable set as the varying input.

    **Restricting the manipulable set is not a filter of the unconstrained POMIS.** It is the
    unconstrained POMIS of the *latent projection* onto the manipulable variables (Lee &
    Bareinboim, AAAI 2019, Theorem 4 — the r40 result :func:`pomis` already implements). Projecting
    away a variable turns paths through it into bidirected edges, which can change the territory
    and border and so *introduce* possibly-optimal sets that the unconstrained enumeration never
    listed. Dropping the sets that mention an infeasible variable would therefore lose optimality,
    not merely prune; that is why :meth:`sets` re-enters :func:`pomis` rather than filtering a
    cached list.

    Because re-entering is not free and the manipulable set usually changes rarely, results are
    memoised on it. ``cache_size`` bounds the memo; the oldest entry is evicted when it is
    exceeded, which suits the common pattern of a feasible set that oscillates among a handful of
    configurations.
    """

    def __init__(self, graph: CausalGraph, reward: str, *, cache_size: int = 32) -> None:
        if reward not in graph.nodes:
            raise ValueError(f"reward {reward!r} is not a node of the graph")
        if cache_size < 1:
            raise ValueError(
                f"cache_size={cache_size} must be at least 1: the memo holds the result for the "
                "manipulable set most recently asked about, and a zero-size cache would recompute "
                "POMIS on every call while still paying for the bookkeeping."
            )
        self._graph = graph
        self._reward = reward
        self._cache_size = cache_size
        self._memo: dict[frozenset[str], list[frozenset[str]]] = {}

    @property
    def reward(self) -> str:
        """The reward variable these intervention sets target."""
        return self._reward

    def sets(self, manipulable: Iterable[str]) -> list[frozenset[str]]:
        """The POMISs available when exactly ``manipulable`` may be intervened on.

        The reward itself is never manipulable and is ignored if present, matching :func:`pomis`.
        """
        key = frozenset(manipulable) - {self._reward}
        cached = self._memo.get(key)
        if cached is not None:
            return [frozenset(s) for s in cached]
        result = pomis(self._graph, self._reward, manipulable=key)
        if len(self._memo) >= self._cache_size:
            del self._memo[next(iter(self._memo))]  # dicts preserve insertion order: oldest first
        self._memo[key] = result
        return [frozenset(s) for s in result]

    def arms(self, space: InterventionSpace) -> list[Intervention]:
        """Every admissible intervention worth considering in ``space``, deduplicated.

        The two halves of the decision, composed: :meth:`sets` says which variables could be worth
        setting given what ``space`` permits, and
        :meth:`~causalrl.intervention.InterventionSpace.assignments` turns each of those sets into
        the concrete assignments an agent chooses between. The observational regime appears as the
        empty intervention whenever the empty set is possibly optimal.

        Order is deterministic — sets in :func:`pomis`'s canonical order (by size, then name), and
        assignments in the order the space enumerates them — so a tie broken by position is stable
        across runs rather than dependent on set iteration order.
        """
        seen: set[tuple[tuple[str, object], ...]] = set()
        out: list[Intervention] = []
        for variables in self.sets(space.variables):
            for assignment in space.assignments(variables):
                key = canonical(assignment)
                if key not in seen:
                    seen.add(key)
                    out.append(assignment)
        return out

reward property

The reward variable these intervention sets target.

sets(manipulable)

The POMISs available when exactly manipulable may be intervened on.

The reward itself is never manipulable and is ignored if present, matching :func:pomis.

Source code in src/causalrl/identification/intervention_sets.py

def sets(self, manipulable: Iterable[str]) -> list[frozenset[str]]:
    """The POMISs available when exactly ``manipulable`` may be intervened on.

    The reward itself is never manipulable and is ignored if present, matching :func:`pomis`.
    """
    key = frozenset(manipulable) - {self._reward}
    cached = self._memo.get(key)
    if cached is not None:
        return [frozenset(s) for s in cached]
    result = pomis(self._graph, self._reward, manipulable=key)
    if len(self._memo) >= self._cache_size:
        del self._memo[next(iter(self._memo))]  # dicts preserve insertion order: oldest first
    self._memo[key] = result
    return [frozenset(s) for s in result]

arms(space)

Every admissible intervention worth considering in space, deduplicated.

The two halves of the decision, composed: :meth:sets says which variables could be worth setting given what space permits, and :meth:~causalrl.intervention.InterventionSpace.assignments turns each of those sets into the concrete assignments an agent chooses between. The observational regime appears as the empty intervention whenever the empty set is possibly optimal.

Order is deterministic — sets in :func:pomis's canonical order (by size, then name), and assignments in the order the space enumerates them — so a tie broken by position is stable across runs rather than dependent on set iteration order.

Source code in src/causalrl/identification/intervention_sets.py

def arms(self, space: InterventionSpace) -> list[Intervention]:
    """Every admissible intervention worth considering in ``space``, deduplicated.

    The two halves of the decision, composed: :meth:`sets` says which variables could be worth
    setting given what ``space`` permits, and
    :meth:`~causalrl.intervention.InterventionSpace.assignments` turns each of those sets into
    the concrete assignments an agent chooses between. The observational regime appears as the
    empty intervention whenever the empty set is possibly optimal.

    Order is deterministic — sets in :func:`pomis`'s canonical order (by size, then name), and
    assignments in the order the space enumerates them — so a tie broken by position is stable
    across runs rather than dependent on set iteration order.
    """
    seen: set[tuple[tuple[str, object], ...]] = set()
    out: list[Intervention] = []
    for variables in self.sets(space.variables):
        for assignment in space.assignments(variables):
            key = canonical(assignment)
            if key not in seen:
                seen.add(key)
                out.append(assignment)
    return out

Benchmark Reports

causalrl.eval.benchmark.BenchmarkEstimate dataclass

A per-seed benchmark measurement with simple descriptive uncertainty.

Source code in src/causalrl/eval/benchmark.py

@dataclass(frozen=True)
class BenchmarkEstimate:
    """A per-seed benchmark measurement with simple descriptive uncertainty."""

    name: str
    seeds: tuple[int, ...]
    values: tuple[float, ...]
    mean: float
    std: float
    ci95_low: float
    ci95_high: float

    @classmethod
    def from_values(
        cls, name: str, *, seeds: Sequence[int], values: Sequence[float]
    ) -> BenchmarkEstimate:
        seed_tuple = tuple(seeds)
        value_tuple = tuple(float(value) for value in values)
        if not seed_tuple or len(seed_tuple) != len(value_tuple):
            raise ValueError("seeds and values must be non-empty sequences of equal length")
        mean = statistics.fmean(value_tuple)
        std = statistics.stdev(value_tuple) if len(value_tuple) > 1 else 0.0
        margin = 1.96 * std / math.sqrt(len(value_tuple))
        return cls(name, seed_tuple, value_tuple, mean, std, mean - margin, mean + margin)

    def to_dict(self) -> dict[str, object]:
        return {
            "name": self.name,
            "seeds": list(self.seeds),
            "values": list(self.values),
            "mean": self.mean,
            "std": self.std,
            "ci95_low": self.ci95_low,
            "ci95_high": self.ci95_high,
        }

causalrl.eval.benchmark.run_confounded_chain_benchmark(*, seeds=(0, 1, 2, 3, 4), n_steps=8000, tail_window=2000, n_mc=2000)

Report POMIS, brute-force, and fixed-set behavior on the confounded-chain demo.

Source code in src/causalrl/eval/benchmark.py

def run_confounded_chain_benchmark(
    *,
    seeds: Sequence[int] = (0, 1, 2, 3, 4),
    n_steps: int = 8000,
    tail_window: int = 2000,
    n_mc: int = 2000,
) -> dict[str, BenchmarkEstimate]:
    """Report POMIS, brute-force, and fixed-set behavior on the confounded-chain demo."""

    def env_factory(seed: int) -> StructuralCausalBanditEnv:
        return make_confounded_chain_env(seed=seed, n_mc=n_mc)

    return {
        "pomis": _tail_mean_rewards(
            "pomis",
            seeds=seeds,
            n_steps=n_steps,
            tail_window=tail_window,
            env_factory=env_factory,
            agent_factory=lambda env, seed: POMISThompsonSampling(
                env.graph, env.reward, env.arms, seed=seed, manipulable=env.manipulable
            ),
        ),
        "brute_force": _tail_mean_rewards(
            "brute_force",
            seeds=seeds,
            n_steps=n_steps,
            tail_window=tail_window,
            env_factory=env_factory,
            agent_factory=lambda env, seed: BruteForceInterventionTS(env.arms, seed=seed),
        ),
        "fixed_set": _tail_mean_rewards(
            "fixed_set",
            seeds=seeds,
            n_steps=n_steps,
            tail_window=tail_window,
            env_factory=env_factory,
            agent_factory=lambda env, seed: FixedSetThompsonSampling(env.arms, {"X3"}, seed=seed),
        ),
    }

causalrl.eval.benchmark.run_frontdoor_benchmark(*, seeds=(0, 1, 2, 3, 4), n_steps=30000, tail_window=10000, n_mc=20000)

Report manipulability-aware and naive-filter behavior on the front-door demo.

Source code in src/causalrl/eval/benchmark.py

def run_frontdoor_benchmark(
    *,
    seeds: Sequence[int] = (0, 1, 2, 3, 4),
    n_steps: int = 30000,
    tail_window: int = 10000,
    n_mc: int = 20000,
) -> dict[str, BenchmarkEstimate]:
    """Report manipulability-aware and naive-filter behavior on the front-door demo."""

    def env_factory(seed: int) -> StructuralCausalBanditEnv:
        return make_frontdoor_env(seed=seed, n_mc=n_mc)

    return {
        "manipulability_aware": _tail_mean_rewards(
            "manipulability_aware",
            seeds=seeds,
            n_steps=n_steps,
            tail_window=tail_window,
            env_factory=env_factory,
            agent_factory=lambda env, seed: POMISThompsonSampling(
                env.graph, env.reward, env.arms, seed=seed, manipulable=env.manipulable
            ),
        ),
        "naive_filter": _tail_mean_rewards(
            "naive_filter",
            seeds=seeds,
            n_steps=n_steps,
            tail_window=tail_window,
            env_factory=env_factory,
            agent_factory=lambda env, seed: NaivePOMISThompsonSampling(
                env.graph, env.reward, env.arms, seed=seed
            ),
        ),
    }

Counterfactual Decision-Making (L3 / ETT)

causalrl.identification.counterfactual.counterfactual_expectation(scm, *, outcome, intervention, evidence, n=20000, seed=None)

Return E[ outcome_{do(intervention)} | evidence ] (a Layer-3 counterfactual mean).

Wraps :meth:StructuralCausalModel.counterfactual (abduction-action-prediction) and averages the outcome over the retained, evidence-consistent units. With empty evidence this reduces to the interventional mean E[outcome_{do(intervention)}].

Source code in src/causalrl/identification/counterfactual.py

def counterfactual_expectation(
    scm: StructuralCausalModel,
    *,
    outcome: str,
    intervention: dict[str, float],
    evidence: dict[str, float],
    n: int = 20_000,
    seed: int | None = None,
) -> float:
    """Return ``E[ outcome_{do(intervention)} | evidence ]`` (a Layer-3 counterfactual mean).

    Wraps :meth:`StructuralCausalModel.counterfactual` (abduction-action-prediction) and averages
    the outcome over the retained, evidence-consistent units. With empty ``evidence`` this reduces
    to the interventional mean ``E[outcome_{do(intervention)}]``.
    """
    _validate_nodes(scm, {outcome} | set(intervention) | set(evidence))
    result = scm.counterfactual(evidence, intervention, n, seed=seed)
    return float(result[outcome].float().mean().item())

causalrl.identification.counterfactual.effect_of_treatment_on_treated(scm, *, treatment, outcome, treated, control, n=20000, seed=None)

Effect of Treatment on the Treated: E[Y_{treated} - Y_{control} | treatment = treated].

The treatment effect among the subpopulation that actually received treated (Pearl, Causality §8.2.1). Under confounding this differs from the average treatment effect. Both potential outcomes use the same seed (common random numbers), so they are evaluated on the same abducted units and the difference is matched and low-variance.

Source code in src/causalrl/identification/counterfactual.py

def effect_of_treatment_on_treated(
    scm: StructuralCausalModel,
    *,
    treatment: str,
    outcome: str,
    treated: float,
    control: float,
    n: int = 20_000,
    seed: int | None = None,
) -> float:
    """Effect of Treatment on the Treated: ``E[Y_{treated} - Y_{control} | treatment = treated]``.

    The treatment effect among the subpopulation that actually received ``treated`` (Pearl,
    *Causality* §8.2.1). Under confounding this differs from the average treatment effect. Both
    potential outcomes use the same ``seed`` (common random numbers), so they are evaluated on the
    same abducted units and the difference is matched and low-variance.
    """
    factual = counterfactual_expectation(
        scm,
        outcome=outcome,
        intervention={treatment: treated},
        evidence={treatment: treated},
        n=n,
        seed=seed,
    )
    counterfactual = counterfactual_expectation(
        scm,
        outcome=outcome,
        intervention={treatment: control},
        evidence={treatment: treated},
        n=n,
        seed=seed,
    )
    return factual - counterfactual

causalrl.agents.counterfactual.CounterfactualOptimalPolicy

Bases: BatchAgent

Plays argmax_a E[Y_{do(action_node=a)} | intent] from a known SCM.

The Layer-3 oracle: it precomputes the Regret Decision Criterion table once at construction and then acts greedily on the observed intent. update is a no-op — the model is known, so there is nothing to learn online. The computed table is exposed as decision_table for inspection.

Source code in src/causalrl/agents/counterfactual.py

class CounterfactualOptimalPolicy(BatchAgent):
    """Plays ``argmax_a E[Y_{do(action_node=a)} | intent]`` from a known SCM.

    The Layer-3 oracle: it precomputes the Regret Decision Criterion table once at construction and
    then acts greedily on the observed intent. ``update`` is a no-op — the model is known, so there
    is nothing to learn online. The computed table is exposed as ``decision_table`` for inspection.
    """

    def __init__(
        self,
        scm: StructuralCausalModel,
        *,
        outcome: str,
        action_node: str,
        intent_node: str,
        arms: Sequence[int],
        intents: Sequence[int],
        intent_key: str = "intuition",
        n: int = 20_000,
        seed: int | None = None,
    ) -> None:
        self._intent_key = intent_key
        self.decision_table = regret_decision_table(
            scm,
            outcome=outcome,
            action_node=action_node,
            intent_node=intent_node,
            arms=arms,
            intents=intents,
            n=n,
            seed=seed,
        )
        self._best_arm: dict[int, int] = {
            intent: max(row, key=lambda arm: row[arm])
            for intent, row in self.decision_table.items()
        }

    def act(self, observation: dict[str, Any]) -> int:
        return self._best_arm[int(observation[self._intent_key])]

Transportability

causalrl.identification.transport.SelectionDiagram dataclass

A causal graph plus the variables whose mechanism differs between source and target.

Each selection variable carries an implicit selection node S -> variable (Bareinboim & Pearl). selection_variables must be a subset of graph.nodes.

Source code in src/causalrl/identification/transport.py

@dataclass(frozen=True)
class SelectionDiagram:
    """A causal graph plus the variables whose mechanism differs between source and target.

    Each selection variable carries an implicit selection node ``S -> variable`` (Bareinboim &
    Pearl). ``selection_variables`` must be a subset of ``graph.nodes``.
    """

    graph: CausalGraph
    selection_variables: frozenset[str]

    def __post_init__(self) -> None:
        unknown = set(self.selection_variables) - set(self.graph.nodes)
        if unknown:
            raise CausalGraphError(f"selection variables not in graph: {sorted(unknown)}")
        if any(n.startswith(SEL) or n.startswith(LAT) for n in self.graph.nodes):
            raise CausalGraphError(f"node names must not start with {SEL!r} or {LAT!r}")

causalrl.identification.transport.transport_formula(diagram, *, treatment, outcome, max_adjustment_size=3)

Return the transport formula for P*(outcome | do(treatment)), or None if it is not provably transportable within the supported class (direct / S-admissible adjustment).

Source code in src/causalrl/identification/transport.py

def transport_formula(
    diagram: SelectionDiagram,
    *,
    treatment: str,
    outcome: str,
    max_adjustment_size: int = 3,
) -> TransportFormula | None:
    """Return the transport formula for ``P*(outcome | do(treatment))``, or ``None`` if it is not
    provably transportable within the supported class (direct / S-admissible adjustment)."""
    graph = diagram.graph
    for name in (treatment, outcome):
        if name not in graph.nodes:
            raise CausalGraphError(f"unknown node: {name!r}")
    selection = diagram.selection_variables
    g_bar_x = graph.do_mutilate(treatment)  # the interventional graph G_{\bar X}
    s_nodes = selection_nodes(selection)

    # Case 1: direct transportability — the interventional law is invariant (S ⊥ Y | X).
    if not selection or d_separated(g_bar_x, s_nodes, {outcome}, {treatment}, selection):
        return TransportFormula(
            "direct",
            frozenset(),
            f"P*({outcome}|do({treatment})) = P({outcome}|do({treatment}))",
        )

    # Case 2: an S-admissible adjustment set Z (back-door admissible and S ⊥ Y | Z, X).
    descendants = graph.descendants(treatment)
    candidates = sorted(
        v for v in graph.nodes if v not in descendants and v not in {treatment, outcome}
    )
    for size in range(min(max_adjustment_size, len(candidates)) + 1):
        for combo in combinations(candidates, size):
            z = set(combo)
            if not is_backdoor_admissible(graph, treatment, outcome, z):
                continue
            if not d_separated(g_bar_x, s_nodes, {outcome}, z | {treatment}, selection):
                continue
            expr = f"P*({outcome}|do {treatment}) = sum_z P({outcome}|{treatment},z) P*(z)"
            return TransportFormula("adjustment", frozenset(z), expr)
    return None

causalrl.identification.transport.is_transportable(diagram, *, treatment, outcome)

Whether the target effect is provably transportable (see :func:transport_formula).

Source code in src/causalrl/identification/transport.py

def is_transportable(diagram: SelectionDiagram, *, treatment: str, outcome: str) -> bool:
    """Whether the target effect is provably transportable (see :func:`transport_formula`)."""
    return transport_formula(diagram, treatment=treatment, outcome=outcome) is not None

causalrl.identification.transport.transported_effect(formula, *, treatment, treated_value, outcome, source, target, n=40000, seed=None)

Compute E*[outcome | do(treatment=treated_value)] via the transport formula.

direct: the source interventional mean transfers unchanged. adjustment: reweight the source conditionals E[outcome | treatment, z] by the target covariate marginal P*(z) (discrete z assumed; the demo uses binary covariates). Strata absent from the source sample are skipped.

Source code in src/causalrl/identification/transport.py

def transported_effect(
    formula: TransportFormula,
    *,
    treatment: str,
    treated_value: float,
    outcome: str,
    source: StructuralCausalModel,
    target: StructuralCausalModel,
    n: int = 40_000,
    seed: int | None = None,
) -> float:
    """Compute ``E*[outcome | do(treatment=treated_value)]`` via the transport `formula`.

    ``direct``: the source interventional mean transfers unchanged. ``adjustment``: reweight the
    source conditionals ``E[outcome | treatment, z]`` by the *target* covariate marginal ``P*(z)``
    (discrete `z` assumed; the demo uses binary covariates). Strata absent from the source sample
    are skipped.
    """
    if formula.kind == "direct":
        return counterfactual_expectation(
            source,
            outcome=outcome,
            intervention={treatment: treated_value},
            evidence={},
            n=n,
            seed=seed,
        )

    z_vars = sorted(formula.adjustment_set)
    src = source.see(n, seed=seed)
    tgt = target.see(n, seed=None if seed is None else seed + 1)
    x_mask = src[treatment] == float(treated_value)
    if not z_vars:
        kept = int(x_mask.sum().item())
        return float(src[outcome][x_mask].float().mean().item()) if kept else 0.0

    value_lists: list[list[float]] = []
    for z in z_vars:
        column: list[float] = src[z].tolist()  # type: ignore[reportUnknownMemberType]
        value_lists.append(sorted(set(column)))
    total = 0.0
    for combo in product(*value_lists):
        src_sel = x_mask.clone()
        tgt_sel: torch.Tensor | None = None
        for z, value in zip(z_vars, combo, strict=True):
            src_sel = src_sel & (src[z] == value)
            cond = tgt[z] == value
            tgt_sel = cond if tgt_sel is None else tgt_sel & cond
        assert tgt_sel is not None
        if int(src_sel.sum().item()) == 0:
            continue
        e_y = float(src[outcome][src_sel].float().mean().item())
        p_z = float(tgt_sel.float().mean().item())
        total += e_y * p_z
    return total

causalrl.identification.transport.is_backdoor_admissible(graph, treatment, outcome, z)

Back-door criterion: z has no descendant of treatment and blocks every back-door path (treatment ⊥ outcome | z in the graph with treatment's outgoing edges removed).

Source code in src/causalrl/identification/transport.py

def is_backdoor_admissible(graph: CausalGraph, treatment: str, outcome: str, z: set[str]) -> bool:
    """Back-door criterion: `z` has no descendant of `treatment` and blocks every back-door path
    (``treatment ⊥ outcome | z`` in the graph with `treatment`'s outgoing edges removed)."""
    if z & graph.descendants(treatment):
        return False
    underline = CausalGraph(
        [(u, v) for u, v in graph.directed_edges if u != treatment],
        graph.bidirected_edges,
        nodes=graph.nodes,
    )
    return d_separated(underline, {treatment}, {outcome}, z)

General Transportability (sID)

causalrl.identification.transport.transport_estimand(diagram, *, treatment, outcome)

The general (sID) transport estimand for P*(outcome | do(treatment)) over diagram.

A :class:SelectionDiagram adapter over :func:causalrl.identification.id_algorithm.identify_transport: each target c-factor is taken from the source if its mechanism is invariant, else identified from the target. Raises :class:~causalrl.exceptions.NotIdentifiableError when not transportable. Generalizes the direct / S-admissible-adjustment :func:transport_formula (which returns a readable closed form for those two cases).

Source code in src/causalrl/identification/transport.py

def transport_estimand(diagram: SelectionDiagram, *, treatment: str, outcome: str) -> Estimand:
    """The general (sID) transport estimand for ``P*(outcome | do(treatment))`` over ``diagram``.

    A :class:`SelectionDiagram` adapter over
    :func:`causalrl.identification.id_algorithm.identify_transport`: each target c-factor is taken
    from the source if its mechanism is invariant, else identified from the target. Raises
    :class:`~causalrl.exceptions.NotIdentifiableError` when not transportable. Generalizes the
    direct / S-admissible-adjustment :func:`transport_formula` (which returns a readable closed form
    for those two cases).
    """
    return identify_transport(diagram.graph, [treatment], [outcome], diagram.selection_variables)

causalrl.identification.id_algorithm.identify_transport(graph, treatment, outcome, selection)

Return an :class:Estimand for the target effect P*(outcome | do(treatment)) across a selection diagram, or raise :class:NotIdentifiableError.

selection lists the variables whose mechanism differs between source and target (each carries an implicit selection node). The target effect is decomposed into c-factors; a c-factor with no selection-marked variable is invariant and taken from the source distribution, while one that contains a selection-marked variable is identified from the target's observational distribution. With an empty selection this reduces to source identification (the ID algorithm).

This is the sound c-factor-routing transportability algorithm; it subsumes direct and S-admissible-adjustment transportability. It is not the complete sID (a c-factor identifiable only by combining source and target is reported as non-transportable rather than guessed).

Faithful to J. Pearl, E. Bareinboim, Transportability of Causal and Statistical Relations, AAAI 2011 / External Validity, Statistical Science 2014, via Tian's c-factor decomposition.

Source code in src/causalrl/identification/id_algorithm.py

def identify_transport(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    selection: Iterable[str],
) -> Estimand:
    """Return an :class:`Estimand` for the *target* effect ``P*(outcome | do(treatment))`` across a
    selection diagram, or raise :class:`NotIdentifiableError`.

    ``selection`` lists the variables whose mechanism differs between source and target (each
    carries an implicit selection node). The target effect is decomposed into c-factors; a c-factor
    with no selection-marked variable is invariant and taken from the **source** distribution, while
    one that contains a selection-marked variable is identified from the **target**'s observational
    distribution. With an empty ``selection`` this reduces to source identification (the ID
    algorithm).

    This is the sound c-factor-routing transportability algorithm; it subsumes direct and
    S-admissible-adjustment transportability. It is not the complete sID (a c-factor identifiable
    only by combining source and target is reported as non-transportable rather than guessed).

    Faithful to J. Pearl, E. Bareinboim, *Transportability of Causal and Statistical Relations*,
    AAAI 2011 / *External Validity*, Statistical Science 2014, via Tian's c-factor decomposition.
    """
    x, y, s = frozenset(treatment), frozenset(outcome), frozenset(selection)
    unknown = (x | y | s) - set(graph.nodes)
    if unknown:
        raise CausalGraphError(f"unknown nodes: {sorted(unknown)}")
    if not y:
        raise CausalGraphError("outcome must be non-empty")
    if x & y:
        raise CausalGraphError(f"treatment and outcome overlap: {sorted(x & y)}")
    return _transport(y, x, graph, s)

causalrl.identification.id_algorithm.is_transportable_effect(graph, treatment, outcome, selection)

Whether the target effect P*(outcome | do(treatment)) is transportable (see :func:identify_transport).

Source code in src/causalrl/identification/id_algorithm.py

def is_transportable_effect(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    selection: Iterable[str],
) -> bool:
    """Whether the target effect ``P*(outcome | do(treatment))`` is transportable (see
    :func:`identify_transport`)."""
    try:
        identify_transport(graph, treatment, outcome, selection)
    except NotIdentifiableError:
        return False
    return True

causalrl.identification.id_algorithm.estimate_transported_effect(graph, treatment, outcome, selection, source_data, target_data, *, do)

Estimate the target effect P*(outcome | do(treatment = do)) from source + target data.

Identifies the transport estimand (:func:identify_transport), then evaluates it with source c-factors read from source_data and target c-factors from target_data (both discrete integer columns over graph.nodes). Returns the outcome distribution keyed in sorted(outcome) order.

Source code in src/causalrl/identification/id_algorithm.py

def estimate_transported_effect(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    selection: Iterable[str],
    source_data: Mapping[str, np.ndarray],
    target_data: Mapping[str, np.ndarray],
    *,
    do: Mapping[str, int],
) -> dict[tuple[int, ...], float]:
    """Estimate the target effect ``P*(outcome | do(treatment = do))`` from source + target data.

    Identifies the transport estimand (:func:`identify_transport`), then evaluates it with source
    c-factors read from ``source_data`` and target c-factors from ``target_data`` (both discrete
    integer columns over ``graph.nodes``). Returns the outcome distribution keyed in
    ``sorted(outcome)`` order.
    """
    estimand = identify_transport(graph, treatment, outcome, selection)
    domains = {
        "source": _empirical_joint(source_data, graph.nodes),
        "target": _empirical_joint(target_data, graph.nodes),
    }
    context = _EvalContext(domains["target"], {}, domains, {})
    factor = _evaluate(estimand, context).condition(do)
    targets = sorted(outcome)
    factor = factor.marginalize(set(factor.variables) - set(targets)).reorder(targets)
    return dict(factor.table)

Multiple Domains And Experiments (mz / meta)

causalrl.identification.id_algorithm.Domain dataclass

A source domain relative to the target: which mechanisms differ, and what data it offers.

selection lists the variables whose mechanism differs from the target (each carries an implicit selection node S -> v); a c-factor transfers from this domain only if it touches no selection-marked variable. experiments lists intervention targets available here (each a set of variables that can be randomized). The target domain is always implicitly available observationally as a fallback, so it need not be listed.

Source code in src/causalrl/identification/id_algorithm.py

@dataclass(frozen=True)
class Domain:
    """A source domain relative to the target: which mechanisms differ, and what data it offers.

    ``selection`` lists the variables whose mechanism differs from the target (each carries an
    implicit selection node ``S -> v``); a c-factor transfers from this domain only if it touches no
    selection-marked variable. ``experiments`` lists intervention targets available here (each a set
    of variables that can be randomized). The target domain is always implicitly available
    observationally as a fallback, so it need not be listed.
    """

    name: str
    selection: frozenset[str] = frozenset()
    experiments: frozenset[frozenset[str]] = frozenset()

causalrl.identification.id_algorithm.identify_transport_general(target_graph, treatment, outcome, domains)

Return an :class:Estimand for the target effect P*(outcome | do(treatment)) combining one or more source domains with the target's own data, or raise :class:NotIdentifiableError.

Generalizes :func:identify_transport to multiple domains (meta-transportability) and to surrogate experiments (mz-transportability): each target c-factor is taken from any domain whose mechanism for it is invariant and that can identify it from its observational or experimental data, with the target as the fallback. With a single observational source it coincides with :func:identify_transport; with no selection and no experiments it reduces to the ID algorithm.

Faithful to E. Bareinboim & J. Pearl, A General Algorithm for Deciding Transportability of Experimental Results (Journal of Causal Inference 2013) and Meta-Transportability of Causal Effects: A Formal Approach (AISTATS 2013), unified via the surrogate-experiment view of S. Lee, J. Correa & E. Bareinboim (UAI 2019). No external code is ported.

Source code in src/causalrl/identification/id_algorithm.py

def identify_transport_general(
    target_graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    domains: Iterable[Domain],
) -> Estimand:
    """Return an :class:`Estimand` for the target effect ``P*(outcome | do(treatment))`` combining
    one or more source ``domains`` with the target's own data, or raise
    :class:`NotIdentifiableError`.

    Generalizes :func:`identify_transport` to multiple domains (meta-transportability) and to
    surrogate experiments (mz-transportability): each target c-factor is taken from any domain whose
    mechanism for it is invariant and that can identify it from its observational or experimental
    data, with the target as the fallback. With a single observational source it coincides with
    :func:`identify_transport`; with no selection and no experiments it reduces to the ID algorithm.

    Faithful to E. Bareinboim & J. Pearl, *A General Algorithm for Deciding Transportability of
    Experimental Results* (Journal of Causal Inference 2013) and *Meta-Transportability of Causal
    Effects: A Formal Approach* (AISTATS 2013), unified via the surrogate-experiment view of S. Lee,
    J. Correa & E. Bareinboim (UAI 2019). No external code is ported.
    """
    x, y = frozenset(treatment), frozenset(outcome)
    doms = list(domains)
    referenced = set(x) | set(y)
    for d in doms:
        referenced |= d.selection
        for exp in d.experiments:
            referenced |= exp
    unknown = referenced - set(target_graph.nodes)
    if unknown:
        raise CausalGraphError(f"unknown nodes: {sorted(unknown)}")
    if not y:
        raise CausalGraphError("outcome must be non-empty")
    if x & y:
        raise CausalGraphError(f"treatment and outcome overlap: {sorted(x & y)}")
    return _general_transport(y, x, target_graph, doms)

causalrl.identification.id_algorithm.is_transportable_general(target_graph, treatment, outcome, domains)

Whether P*(outcome | do(treatment)) is transportable from domains (see :func:identify_transport_general).

Source code in src/causalrl/identification/id_algorithm.py

def is_transportable_general(
    target_graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    domains: Iterable[Domain],
) -> bool:
    """Whether ``P*(outcome | do(treatment))`` is transportable from ``domains`` (see
    :func:`identify_transport_general`)."""
    try:
        identify_transport_general(target_graph, treatment, outcome, domains)
    except NotIdentifiableError:
        return False
    return True

causalrl.identification.id_algorithm.estimate_transport_general(target_graph, treatment, outcome, domains, *, domain_data, experiment_data=None, do)

Estimate the target effect P*(outcome | do(treatment = do)) across domains.

domain_data maps each domain name (including "target") to an observational dataset over target_graph.nodes; experiment_data maps (domain, frozenset(target)) to a randomized experimental dataset. Identifies the estimand via :func:identify_transport_general, evaluates it against the supplied data, and returns the outcome distribution keyed in sorted(outcome) order.

Source code in src/causalrl/identification/id_algorithm.py

def estimate_transport_general(
    target_graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    domains: Iterable[Domain],
    *,
    domain_data: Mapping[str, Mapping[str, np.ndarray]],
    experiment_data: Mapping[tuple[str, frozenset[str]], Mapping[str, np.ndarray]] | None = None,
    do: Mapping[str, int],
) -> dict[tuple[int, ...], float]:
    """Estimate the target effect ``P*(outcome | do(treatment = do))`` across ``domains``.

    ``domain_data`` maps each domain name (including ``"target"``) to an observational dataset over
    ``target_graph.nodes``; ``experiment_data`` maps ``(domain, frozenset(target))`` to a randomized
    experimental dataset. Identifies the estimand via :func:`identify_transport_general`, evaluates
    it against the supplied data, and returns the outcome distribution keyed in ``sorted(outcome)``
    order.
    """
    doms = list(domains)
    estimand = identify_transport_general(target_graph, treatment, outcome, doms)
    if "target" not in domain_data:
        raise CausalGraphError("domain_data must include 'target'")
    nodes = target_graph.nodes
    domains_f = {name: _empirical_joint(rows, nodes) for name, rows in domain_data.items()}
    exps_f = {key: _empirical_joint(rows, nodes) for key, rows in (experiment_data or {}).items()}
    ctx = _EvalContext(domains_f["target"], {}, domains_f, exps_f)
    factor = _evaluate(estimand, ctx).condition(do)
    targets = sorted(outcome)
    factor = factor.marginalize(set(factor.variables) - set(targets)).reorder(targets)
    return dict(factor.table)

General Identification (ID Algorithm)

causalrl.identification.id_algorithm.identify_effect(graph, treatment, outcome, *, return_certificate=None)

identify_effect(graph: CausalGraph, treatment: Iterable[str], outcome: Iterable[str], *, return_certificate: Literal[True] | None = ...) -> Certificate

identify_effect(graph: CausalGraph, treatment: Iterable[str], outcome: Iterable[str], *, return_certificate: Literal[False]) -> Estimand

Return a do-free :class:Estimand for P(outcome | do(treatment)), or raise.

Runs the ID algorithm. Raises :class:NotIdentifiableError (with the witnessing hedge attached as .witness) when the effect is not non-parametrically identifiable, and :class:CausalGraphError for malformed inputs (unknown nodes, empty outcome, or a treatment and outcome that overlap).

return_certificate (2.0 default): returns a :class:Certificate; pass False for the legacy :class:Estimand, or True (equivalently, identify_effect_certified) to be explicit.

Source code in src/causalrl/identification/id_algorithm.py

def identify_effect(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    *,
    return_certificate: bool | None = None,
) -> Estimand | Certificate:
    """Return a do-free :class:`Estimand` for ``P(outcome | do(treatment))``, or raise.

    Runs the ID algorithm. Raises :class:`NotIdentifiableError` (with the witnessing hedge attached
    as ``.witness``) when the effect is not non-parametrically identifiable, and
    :class:`CausalGraphError` for malformed inputs (unknown nodes, empty outcome, or a treatment and
    outcome that overlap).

    ``return_certificate`` (2.0 default): returns a :class:`Certificate`; pass ``False`` for the
    legacy :class:`Estimand`, or ``True`` (equivalently, ``identify_effect_certified``) to be
    explicit.
    """
    if return_certificate is None:
        return_certificate = True  # causalrl 2.0: return a Certificate by default
    if return_certificate:
        from causalrl.certify.routines import identify_effect_certified

        return identify_effect_certified(graph, treatment, outcome)
    x, y = frozenset(treatment), frozenset(outcome)
    unknown = (x | y) - set(graph.nodes)
    if unknown:
        raise CausalGraphError(f"unknown nodes: {sorted(unknown)}")
    if not y:
        raise CausalGraphError("outcome must be non-empty")
    if x & y:
        raise CausalGraphError(f"treatment and outcome overlap: {sorted(x & y)}")
    return _id(y, x, graph, _Joint(frozenset(graph.nodes)))

causalrl.identification.id_algorithm.is_identifiable_effect(graph, treatment, outcome)

Whether P(outcome | do(treatment)) is identifiable from observational data.

Source code in src/causalrl/identification/id_algorithm.py

def is_identifiable_effect(
    graph: CausalGraph, treatment: Iterable[str], outcome: Iterable[str]
) -> bool:
    """Whether ``P(outcome | do(treatment))`` is identifiable from observational data."""
    try:
        identify_effect(graph, treatment, outcome, return_certificate=False)
    except NotIdentifiableError:
        return False
    return True

causalrl.identification.id_algorithm.estimate_effect(graph, treatment, outcome, data, *, do)

Estimate P(outcome | do(treatment = do)) from observational data via the ID estimand.

Identifies the effect (raising if it cannot), then evaluates the resulting estimand on the empirical joint of data (discrete integer columns over graph.nodes) at the intervention do. Returns the outcome distribution as {assignment: probability} with assignments keyed in sorted(outcome) order.

Source code in src/causalrl/identification/id_algorithm.py

def estimate_effect(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    data: Mapping[str, np.ndarray],
    *,
    do: Mapping[str, int],
) -> dict[tuple[int, ...], float]:
    """Estimate ``P(outcome | do(treatment = do))`` from observational ``data`` via the ID estimand.

    Identifies the effect (raising if it cannot), then evaluates the resulting estimand on the
    empirical joint of ``data`` (discrete integer columns over ``graph.nodes``) at the intervention
    ``do``. Returns the outcome distribution as ``{assignment: probability}`` with assignments keyed
    in ``sorted(outcome)`` order.
    """
    estimand = identify_effect(graph, treatment, outcome, return_certificate=False)
    factor = _eval_observational(estimand, _empirical_joint(data, graph.nodes))
    factor = factor.condition(do)
    targets = sorted(outcome)
    factor = factor.marginalize(set(factor.variables) - set(targets)).reorder(targets)
    return dict(factor.table)

causalrl.identification.id_algorithm.Estimand

A do-free expression for an interventional distribution over the observational law P(V).

Render it with :meth:render; evaluate it on data with :func:estimate_effect.

Source code in src/causalrl/identification/id_algorithm.py

class Estimand:
    """A do-free expression for an interventional distribution over the observational law ``P(V)``.

    Render it with :meth:`render`; evaluate it on data with :func:`estimate_effect`.
    """

    def render(self) -> str:
        raise NotImplementedError

From Surrogate Experiments (gID)

causalrl.identification.id_algorithm.identify_effect_with_experiments(graph, treatment, outcome, experiments)

Return an :class:Estimand for P(outcome | do(treatment)) using surrogate experiments.

This is general identification (gID): it runs the ID recursion but, where observational data hits a hedge, it tries to obtain the needed c-factor from one of the available experiments (each a set of variables you can intervene on; observational data is always available too). Raises :class:NotIdentifiableError if no combination of observational data and experiments identifies the effect.

Faithful to S. Lee, J. Correa, E. Bareinboim, General Identifiability with Arbitrary Surrogate Experiments, UAI 2019, building on Tian & Pearl's c-factor identification. No code is ported.

Source code in src/causalrl/identification/id_algorithm.py

def identify_effect_with_experiments(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    experiments: Iterable[Iterable[str]],
) -> Estimand:
    """Return an :class:`Estimand` for ``P(outcome | do(treatment))`` using surrogate experiments.

    This is general identification (gID): it runs the ID recursion but, where observational data
    hits a hedge, it tries to obtain the needed c-factor from one of the available ``experiments``
    (each a set of variables you can intervene on; observational data is always available too).
    Raises :class:`NotIdentifiableError` if no combination of observational data and experiments
    identifies the effect.

    Faithful to S. Lee, J. Correa, E. Bareinboim, *General Identifiability with Arbitrary Surrogate
    Experiments*, UAI 2019, building on Tian & Pearl's c-factor identification. No code is ported.
    """
    x, y = frozenset(treatment), frozenset(outcome)
    nodes = set(graph.nodes)
    exps = [frozenset(z) for z in experiments]
    referenced = set(x) | set(y)
    for z in exps:
        referenced |= z
    unknown = referenced - nodes
    if unknown:
        raise CausalGraphError(f"unknown nodes: {sorted(unknown)}")
    if not y:
        raise CausalGraphError("outcome must be non-empty")
    if x & y:
        raise CausalGraphError(f"treatment and outcome overlap: {sorted(x & y)}")
    return _id(y, x, graph, _Joint(frozenset(graph.nodes)), experiments=exps, original=graph)

causalrl.identification.id_algorithm.is_gid_identifiable(graph, treatment, outcome, experiments)

Whether P(outcome | do(treatment)) is identifiable from data plus those experiments.

Source code in src/causalrl/identification/id_algorithm.py

def is_gid_identifiable(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    experiments: Iterable[Iterable[str]],
) -> bool:
    """Whether ``P(outcome | do(treatment))`` is identifiable from data plus those experiments."""
    try:
        identify_effect_with_experiments(graph, treatment, outcome, experiments)
    except NotIdentifiableError:
        return False
    return True

causalrl.identification.id_algorithm.estimate_effect_with_experiments(graph, treatment, outcome, data, experiments_data, *, do)

Estimate P(outcome | do(treatment = do)) from observational data plus experiments.

experiments_data maps each available intervention target (a frozenset of variables) to a dataset drawn from do(target) over graph.nodes. Identifies the effect via gID (:func:identify_effect_with_experiments), then evaluates the estimand against the empirical observational joint and each experiment's empirical c-factor. Returns the outcome distribution keyed in sorted(outcome) order.

Source code in src/causalrl/identification/id_algorithm.py

def estimate_effect_with_experiments(
    graph: CausalGraph,
    treatment: Iterable[str],
    outcome: Iterable[str],
    data: Mapping[str, np.ndarray],
    experiments_data: Mapping[frozenset[str], Mapping[str, np.ndarray]],
    *,
    do: Mapping[str, int],
) -> dict[tuple[int, ...], float]:
    """Estimate ``P(outcome | do(treatment = do))`` from observational ``data`` plus experiments.

    ``experiments_data`` maps each available intervention target (a ``frozenset`` of variables) to a
    dataset drawn from ``do(target)`` over ``graph.nodes``. Identifies the effect via gID
    (:func:`identify_effect_with_experiments`), then evaluates the estimand against the empirical
    observational joint and each experiment's empirical c-factor. Returns the outcome distribution
    keyed in ``sorted(outcome)`` order.
    """
    estimand = identify_effect_with_experiments(graph, treatment, outcome, experiments_data.keys())
    obs = _empirical_joint(data, graph.nodes)
    experiments = {
        target: _empirical_joint(rows, graph.nodes) for target, rows in experiments_data.items()
    }
    factor = _evaluate(estimand, _EvalContext(obs, experiments, {}, {})).condition(do)
    targets = sorted(outcome)
    factor = factor.marginalize(set(factor.variables) - set(targets)).reorder(targets)
    return dict(factor.table)

Causal Discovery

causalrl.discovery.discover(data, variables, *, threshold=0.01, max_conditioning_size=3)

Discover the CPDAG over variables from data via the PC algorithm.

threshold is the conditional-mutual-information cutoff below which two variables are judged independent; max_conditioning_size caps the separating-set search.

Source code in src/causalrl/discovery.py

def discover(
    data: Mapping[str, np.ndarray],
    variables: Sequence[str],
    *,
    threshold: float = 0.01,
    max_conditioning_size: int = 3,
) -> CPDAG:
    """Discover the CPDAG over `variables` from `data` via the PC algorithm.

    ``threshold`` is the conditional-mutual-information cutoff below which two variables are judged
    independent; ``max_conditioning_size`` caps the separating-set search.
    """
    nodes = list(variables)
    adj, sepset = _pc_skeleton(
        data, nodes, threshold=threshold, max_conditioning_size=max_conditioning_size
    )

    # Orient unshielded colliders a -> c <- b (a, b non-adjacent and c not in their separating set).
    directed: set[tuple[str, str]] = set()
    undirected: set[frozenset[str]] = {frozenset((a, b)) for a in nodes for b in adj[a] if a < b}
    for c in nodes:
        for a, b in combinations(sorted(adj[c]), 2):
            if b in adj[a]:
                continue  # shielded triple
            if c not in sepset.get(frozenset((a, b)), ()):
                _orient(a, c, directed, undirected)
                _orient(b, c, directed, undirected)

    _apply_meek_rules(nodes, directed, undirected)
    return CPDAG(tuple(nodes), frozenset(directed), frozenset(undirected))

causalrl.discovery.discover_interventional(observational, interventions, variables, *, threshold=0.01, shift_threshold=0.05, max_conditioning_size=3, min_interventional_samples=20)

Discover the interventional essential graph from observational and experimental data.

Runs the observational PC algorithm (:func:discover), then orients the edges incident to each intervention target by the invariance principle: under a perfect intervention do(T) a child of T shifts its marginal, while a parent (a non-descendant) stays invariant. Each incident edge T - B is oriented T -> B if B shifts and B -> T if not; Meek's rules R1-R3 then propagate the new orientations, so the result refines the observational CPDAG toward the true DAG as more targets are experimented on.

interventions maps each intervened target T to a dataset drawn from do(T) (a perfect intervention covering every variable in variables). shift_threshold is the total-variation cutoff above which an endpoint's marginal is judged to have shifted.

min_interventional_samples is the smallest per-target sample this will orient from; below it the call raises rather than guessing. The shift test is an empirical total variation, so a tiny sample clears any threshold by chance and an empty one clears it unconditionally — orienting every incident edge from no evidence at all. Lower it only if you accept orientations backed by that little data.

Source code in src/causalrl/discovery.py

def discover_interventional(
    observational: Mapping[str, np.ndarray],
    interventions: Mapping[str, Mapping[str, np.ndarray]],
    variables: Sequence[str],
    *,
    threshold: float = 0.01,
    shift_threshold: float = 0.05,
    max_conditioning_size: int = 3,
    min_interventional_samples: int = 20,
) -> CPDAG:
    """Discover the interventional essential graph from observational and experimental data.

    Runs the observational PC algorithm (:func:`discover`), then orients the edges incident to each
    intervention target by the invariance principle: under a perfect intervention ``do(T)`` a
    *child* of ``T`` shifts its marginal, while a *parent* (a non-descendant) stays invariant. Each
    incident edge ``T - B`` is oriented ``T -> B`` if ``B`` shifts and ``B -> T`` if not; Meek's
    rules R1-R3 then propagate the new orientations, so the result refines the observational CPDAG
    toward the true DAG as more targets are experimented on.

    ``interventions`` maps each intervened target ``T`` to a dataset drawn from ``do(T)`` (a perfect
    intervention covering every variable in ``variables``). ``shift_threshold`` is the
    total-variation cutoff above which an endpoint's marginal is judged to have shifted.

    ``min_interventional_samples`` is the smallest per-target sample this will orient from; below it
    the call raises rather than guessing. The shift test is an empirical total variation, so a tiny
    sample clears any threshold by chance and an empty one clears it *unconditionally* — orienting
    every incident edge from no evidence at all. Lower it only if you accept orientations backed by
    that little data.
    """
    nodes = list(variables)
    cpdag = discover(
        observational, nodes, threshold=threshold, max_conditioning_size=max_conditioning_size
    )
    directed = set(cpdag.directed_edges)
    undirected = set(cpdag.undirected_edges)

    for target, idata in interventions.items():
        if target not in nodes:
            raise CausalGraphError(f"intervention target not in variables: {target!r}")
        # Refuse to orient from a sample too small to carry the invariance signal. The marginal
        # shift test is an empirical total variation, and on an empty column the pmf is {} so the
        # distance to ANY observational marginal is 0.5 -- ten times the default threshold. An
        # empty do-sample would therefore orient every edge incident to `target` from zero data,
        # and Meek's rules would propagate outward from those fabricated orientations. A single
        # sample is barely better: its pmf is a point mass, so the distance is 1 - p_obs(v), which
        # clears the threshold for almost any value. Orienting an edge is a causal claim; it needs
        # evidence, and silently manufacturing one from no data is the failure this library exists
        # to prevent.
        observed = min((len(column) for column in idata.values()), default=0)
        if observed < min_interventional_samples:
            raise CausalGraphError(
                f"intervention target {target!r} has only {observed} sample(s), below "
                f"min_interventional_samples={min_interventional_samples}. The marginal-shift "
                f"test cannot distinguish a real interventional shift from sampling noise at "
                f"this size, and on an empty sample it reports a maximal shift regardless of "
                f"the data. Collect more samples under do({target}), drop the target, or lower "
                f"min_interventional_samples if you accept orientations backed by this little "
                f"evidence."
            )
        for edge in list(undirected):
            if target not in edge:
                continue
            other = next(iter(edge - {target}))
            if other not in observational or other not in idata:
                raise CausalGraphError(f"variable not in data: {other!r}")
            shifted = _total_variation(observational[other], idata[other]) >= shift_threshold
            undirected.discard(edge)
            directed.add((target, other) if shifted else (other, target))

    _apply_meek_rules(nodes, directed, undirected)
    return CPDAG(tuple(nodes), frozenset(directed), frozenset(undirected))

causalrl.discovery.discover_latent(data, variables, *, threshold=0.01, max_conditioning_size=3)

Discover a PAG over variables from data via the FCI algorithm (allows latents).

Unlike :func:discover, FCI does not assume causal sufficiency: it learns the PC skeleton, then refines it with the Possible-D-SEP step (sound under latent confounders), re-orients unshielded colliders, and applies the complete orientation rules R1-R10 (Zhang 2008 — sound and complete for latent confounders and selection bias). The result is a :class:PAG: a <-> b witnesses a latent confounder; a circle endpoint is undetermined by the equivalence class.

threshold and max_conditioning_size mirror :func:discover.

Source code in src/causalrl/discovery.py

def discover_latent(
    data: Mapping[str, np.ndarray],
    variables: Sequence[str],
    *,
    threshold: float = 0.01,
    max_conditioning_size: int = 3,
) -> PAG:
    """Discover a PAG over ``variables`` from ``data`` via the FCI algorithm (allows latents).

    Unlike :func:`discover`, FCI does not assume causal sufficiency: it learns the PC skeleton, then
    refines it with the Possible-D-SEP step (sound under latent confounders), re-orients unshielded
    colliders, and applies the complete orientation rules R1-R10 (Zhang 2008 — sound and complete
    for latent confounders and selection bias). The result is a :class:`PAG`: ``a <-> b``
    witnesses a latent confounder; a circle endpoint is undetermined by the equivalence class.

    ``threshold`` and ``max_conditioning_size`` mirror :func:`discover`.
    """
    nodes = list(variables)
    adj, sepset = _pc_skeleton(
        data, nodes, threshold=threshold, max_conditioning_size=max_conditioning_size
    )
    marks: dict[tuple[str, str], str] = {}
    for a in nodes:
        for b in adj[a]:
            marks[(a, b)] = _CIRCLE
    _orient_colliders(marks, nodes, sepset)
    _refine_skeleton_pds(
        data, marks, sepset, nodes, threshold=threshold, max_conditioning_size=max_conditioning_size
    )
    for key in marks:
        marks[key] = _CIRCLE
    _orient_colliders(marks, nodes, sepset)
    _apply_fci_rules(marks, sepset)
    return PAG(tuple(nodes), marks)

causalrl.discovery.CPDAG dataclass

A completed partially directed acyclic graph (a Markov equivalence class).

Source code in src/causalrl/discovery.py

@dataclass(frozen=True)
class CPDAG:
    """A completed partially directed acyclic graph (a Markov equivalence class)."""

    variables: tuple[str, ...]
    directed_edges: frozenset[tuple[str, str]]
    undirected_edges: frozenset[frozenset[str]]

    def to_causal_graph(self) -> CausalGraph:
        """Convert to a :class:`CausalGraph`; raises if any edge is still unoriented."""
        if self.undirected_edges:
            remaining = sorted(tuple(sorted(e)) for e in self.undirected_edges)
            raise CausalGraphError(
                f"CPDAG is not fully oriented; undirected edges remain: {remaining}"
            )
        return CausalGraph(directed_edges=sorted(self.directed_edges), nodes=list(self.variables))

to_causal_graph()

Convert to a :class:CausalGraph; raises if any edge is still unoriented.

Source code in src/causalrl/discovery.py

def to_causal_graph(self) -> CausalGraph:
    """Convert to a :class:`CausalGraph`; raises if any edge is still unoriented."""
    if self.undirected_edges:
        remaining = sorted(tuple(sorted(e)) for e in self.undirected_edges)
        raise CausalGraphError(
            f"CPDAG is not fully oriented; undirected edges remain: {remaining}"
        )
    return CausalGraph(directed_edges=sorted(self.directed_edges), nodes=list(self.variables))

causalrl.discovery.PAG dataclass

A partial ancestral graph: the complete Markov-equivalence class of MAGs (the FCI output).

marks[(a, b)] is the mark on the b end of edge a—b — a circle o (undetermined by the equivalence class), an arrowhead >, or a tail -. An edge exists iff both (a, b) and (b, a) are present. a -> b is tail-at-a / arrow-at-b; a <-> b (arrowheads at both ends) witnesses a latent confounder; a o-o b is fully unoriented.

Source code in src/causalrl/discovery.py

@dataclass(frozen=True)
class PAG:
    """A partial ancestral graph: the complete Markov-equivalence class of MAGs (the FCI output).

    ``marks[(a, b)]`` is the mark on the ``b`` end of edge ``a—b`` — a circle ``o`` (undetermined by
    the equivalence class), an arrowhead ``>``, or a tail ``-``. An edge exists iff both ``(a, b)``
    and ``(b, a)`` are present. ``a -> b`` is tail-at-``a`` / arrow-at-``b``; ``a <-> b``
    (arrowheads at both ends) witnesses a latent confounder; ``a o-o b`` is fully unoriented.
    """

    variables: tuple[str, ...]
    marks: Mapping[tuple[str, str], str]

    def __post_init__(self) -> None:
        for (a, b), mark in self.marks.items():
            if mark not in (_CIRCLE, _ARROW, _TAIL):
                raise CausalGraphError(f"invalid PAG mark {mark!r} on edge {a}-{b}")
            if (b, a) not in self.marks:
                raise CausalGraphError(f"PAG edge {a}-{b} is missing endpoint ({b!r}, {a!r})")

    def adjacent(self, a: str, b: str) -> bool:
        return (a, b) in self.marks

    def is_directed(self, a: str, b: str) -> bool:
        """Whether ``a -> b`` (tail at ``a``, arrowhead at ``b``)."""
        return self.marks.get((a, b)) == _ARROW and self.marks.get((b, a)) == _TAIL

    def is_bidirected(self, a: str, b: str) -> bool:
        """Whether ``a <-> b`` (arrowheads at both ends — a latent confounder)."""
        return self.marks.get((a, b)) == _ARROW and self.marks.get((b, a)) == _ARROW

    def edges(self) -> list[tuple[str, str, str, str]]:
        """``(a, b, mark_at_a, mark_at_b)`` for each edge, with ``a < b``."""
        out: list[tuple[str, str, str, str]] = []
        seen: set[frozenset[str]] = set()
        for a, b in self.marks:
            key = frozenset((a, b))
            if key in seen:
                continue
            seen.add(key)
            x, y = sorted((a, b))
            out.append((x, y, self.marks[(y, x)], self.marks[(x, y)]))
        return sorted(out)

    def render(self) -> str:
        left = {_CIRCLE: "o", _ARROW: "<", _TAIL: "-"}
        right = {_CIRCLE: "o", _ARROW: ">", _TAIL: "-"}
        return ", ".join(f"{x} {left[ma]}-{right[mb]} {y}" for x, y, ma, mb in self.edges())

is_directed(a, b)

Whether a -> b (tail at a, arrowhead at b).

Source code in src/causalrl/discovery.py

def is_directed(self, a: str, b: str) -> bool:
    """Whether ``a -> b`` (tail at ``a``, arrowhead at ``b``)."""
    return self.marks.get((a, b)) == _ARROW and self.marks.get((b, a)) == _TAIL

is_bidirected(a, b)

Whether a <-> b (arrowheads at both ends — a latent confounder).

Source code in src/causalrl/discovery.py

def is_bidirected(self, a: str, b: str) -> bool:
    """Whether ``a <-> b`` (arrowheads at both ends — a latent confounder)."""
    return self.marks.get((a, b)) == _ARROW and self.marks.get((b, a)) == _ARROW

edges()

(a, b, mark_at_a, mark_at_b) for each edge, with a < b.

Source code in src/causalrl/discovery.py

def edges(self) -> list[tuple[str, str, str, str]]:
    """``(a, b, mark_at_a, mark_at_b)`` for each edge, with ``a < b``."""
    out: list[tuple[str, str, str, str]] = []
    seen: set[frozenset[str]] = set()
    for a, b in self.marks:
        key = frozenset((a, b))
        if key in seen:
            continue
        seen.add(key)
        x, y = sorted((a, b))
        out.append((x, y, self.marks[(y, x)], self.marks[(x, y)]))
    return sorted(out)

causalrl.discovery.conditional_mutual_information(data, x, y, z)

Empirical I(X; Y | Z) in nats (discrete columns; 0 iff X ⊥ Y | Z).

Counted with :func:numpy.unique over integer-encoded cells rather than a Python loop over rows. This is the inner loop of every discovery routine here -- PC, FCI and the interventional and invariance variants all bottom out in it, once per (pair, conditioning set) triple -- so its constant factor sets how large a problem discover can take at all. The counts are sparse (only cells that occur are materialised), so the memory cost follows the data rather than the product of the variables' level counts.

Source code in src/causalrl/discovery.py

def conditional_mutual_information(
    data: Mapping[str, np.ndarray], x: str, y: str, z: Sequence[str]
) -> float:
    """Empirical ``I(X; Y | Z)`` in nats (discrete columns; ``0`` iff ``X ⊥ Y | Z``).

    Counted with :func:`numpy.unique` over integer-encoded cells rather than a Python loop over
    rows. This is the inner loop of every discovery routine here -- PC, FCI and the interventional
    and invariance variants all bottom out in it, once per (pair, conditioning set) triple -- so its
    constant factor sets how large a problem ``discover`` can take at all. The counts are sparse
    (only cells that occur are materialised), so the memory cost follows the data rather than the
    product of the variables' level counts.
    """
    xs = np.asarray(data[x]).ravel()
    ys = np.asarray(data[y]).ravel()
    n = int(xs.size)
    if n == 0:
        return 0.0
    _, xcode = np.unique(xs.astype(np.int64), return_inverse=True)
    _, ycode = np.unique(ys.astype(np.int64), return_inverse=True)
    # Mix the conditioning columns into one integer key, then re-factorise it. A row-wise
    # np.unique(axis=0) would do the same job but sorts a structured view of the whole block,
    # which costs several times more than sorting the int64 keys it produces.
    zcode = np.zeros(n, dtype=np.int64)
    for name in z:
        _, column = np.unique(np.asarray(data[name]).ravel().astype(np.int64), return_inverse=True)
        column = column.astype(np.int64).ravel()
        zcode = zcode * (int(column.max()) + 1) + column
    if z:
        _, zcode = np.unique(zcode, return_inverse=True)
        zcode = zcode.astype(np.int64).ravel()
    n_x = int(xcode.max()) + 1
    n_y = int(ycode.max()) + 1

    # One integer key per cell, so every count is a single sorted pass.
    xz_key = zcode * n_x + xcode
    yz_key = zcode * n_y + ycode
    joint_key = xz_key * n_y + ycode

    n_z = int(zcode.max()) + 1
    dense = n_z * n_x * n_y
    if dense <= max(4 * n, 1024):
        # Dense counting is a linear pass instead of a sort, and at this size the table is smaller
        # than the data. Above it, the sparse path keeps memory proportional to occupied cells.
        table = np.bincount(joint_key, minlength=dense)
        keys = np.flatnonzero(table)
        joint = table[keys]
        count_xz_dense = np.bincount(xz_key, minlength=n_z * n_x)
        count_yz_dense = np.bincount(yz_key, minlength=n_z * n_y)
        count_z_dense = np.bincount(zcode, minlength=n_z)
        cell_y = keys % n_y
        cell_rest = keys // n_y
        cell_x = cell_rest % n_x
        cell_z = cell_rest // n_x
        c_xz = count_xz_dense[cell_z * n_x + cell_x]
        c_yz = count_yz_dense[cell_z * n_y + cell_y]
        c_z = count_z_dense[cell_z]
    else:
        keys, joint = np.unique(joint_key, return_counts=True)
        uniq_xz, count_xz = np.unique(xz_key, return_counts=True)
        uniq_yz, count_yz = np.unique(yz_key, return_counts=True)
        uniq_z, count_z = np.unique(zcode, return_counts=True)
        cell_y = keys % n_y
        cell_rest = keys // n_y
        cell_x = cell_rest % n_x
        cell_z = cell_rest // n_x
        c_xz = count_xz[np.searchsorted(uniq_xz, cell_z * n_x + cell_x)]
        c_yz = count_yz[np.searchsorted(uniq_yz, cell_z * n_y + cell_y)]
        c_z = count_z[np.searchsorted(uniq_z, cell_z)]

    ratio = (joint * c_z) / (c_xz * c_yz)
    cmi = float(np.sum((joint / n) * np.log(ratio)))
    return max(cmi, 0.0)

Causal Imitation Learning

causalrl.imitation.is_imitable(graph, *, action, outcome, observable)

Whether the expert is imitable: an observed back-door-admissible set exists.

Source code in src/causalrl/imitation.py

def is_imitable(
    graph: CausalGraph, *, action: str, outcome: str, observable: Iterable[str]
) -> bool:
    """Whether the expert is imitable: an observed back-door-admissible set exists."""
    return (
        imitation_backdoor_set(graph, action=action, outcome=outcome, observable=observable)
        is not None
    )

causalrl.imitation.imitation_backdoor_set(graph, *, action, outcome, observable, max_size=3)

Smallest observed back-door-admissible set for action -> outcome, or None.

A set Z ⊆ observable \ {action, outcome} is admissible when it has no descendant of action and blocks every back-door path (action ⊥ outcome | Z with action's outgoing edges removed). Cloning P(action | Z) then reproduces the expert's outcome distribution.

Source code in src/causalrl/imitation.py

def imitation_backdoor_set(
    graph: CausalGraph,
    *,
    action: str,
    outcome: str,
    observable: Iterable[str],
    max_size: int = 3,
) -> frozenset[str] | None:
    """Smallest observed back-door-admissible set for ``action -> outcome``, or ``None``.

    A set ``Z ⊆ observable \\ {action, outcome}`` is admissible when it has no descendant of
    ``action`` and blocks every back-door path (``action ⊥ outcome | Z`` with ``action``'s outgoing
    edges removed). Cloning ``P(action | Z)`` then reproduces the expert's outcome distribution.
    """
    for name in (action, outcome):
        if name not in graph.nodes:
            raise CausalGraphError(f"unknown node: {name!r}")
    observed = set(observable)
    unknown = observed - set(graph.nodes)
    if unknown:
        raise CausalGraphError(f"unknown observable node(s): {sorted(unknown)}")
    candidates = sorted(observed - {action, outcome})
    for size in range(min(max_size, len(candidates)) + 1):
        for combo in combinations(candidates, size):
            z = set(combo)
            if is_backdoor_admissible(graph, action, outcome, z):
                return frozenset(z)
    return None

causalrl.imitation.CausalImitator

Bases: Agent

Clones P(A | Z) for a back-door-admissible observed set Z (the adjustment set).

Conditioning the cloned policy on Z reproduces the confounding the expert responded to, so deployment matches the expert's reward. Unseen Z values fall back to the marginal P(A).

Source code in src/causalrl/imitation.py

class CausalImitator(Agent):
    """Clones ``P(A | Z)`` for a back-door-admissible observed set ``Z`` (the adjustment set).

    Conditioning the cloned policy on ``Z`` reproduces the confounding the expert responded to, so
    deployment matches the expert's reward. Unseen ``Z`` values fall back to the marginal ``P(A)``.
    """

    def __init__(self, n_actions: int, adjustment: Sequence[str], seed: int | None = None) -> None:
        self.n_actions = n_actions
        self._adjustment = list(adjustment)
        self._rng = np.random.default_rng(seed)
        self._marginal = np.ones(n_actions) / n_actions
        self._conditional: dict[tuple[int, ...], np.ndarray] = {}

    def fit(self, demonstrations: Mapping[str, np.ndarray], *, action: str) -> None:
        actions = np.asarray(demonstrations[action]).astype(int)
        marginal = np.bincount(actions, minlength=self.n_actions).astype(float)
        self._marginal = marginal / marginal.sum()
        z_columns = [np.asarray(demonstrations[z]).astype(int) for z in self._adjustment]
        counts: dict[tuple[int, ...], np.ndarray] = {}
        for i in range(len(actions)):
            key = tuple(int(column[i]) for column in z_columns)
            if key not in counts:
                counts[key] = np.zeros(self.n_actions)
            counts[key][int(actions[i])] += 1.0
        self._conditional = {key: row / row.sum() for key, row in counts.items()}

    def act(self, observation: dict[str, Any]) -> int:
        key = tuple(int(observation[z]) for z in self._adjustment)
        probs = self._conditional.get(key, self._marginal)
        return int(self._rng.choice(self.n_actions, p=probs))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        """No-op: imitation is fit offline from demonstrations."""

update(observation, action, reward)

No-op: imitation is fit offline from demonstrations.

Source code in src/causalrl/imitation.py

def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
    """No-op: imitation is fit offline from demonstrations."""

causalrl.imitation.BehavioralCloning

Bases: Agent

Naive imitator: clones the marginal action distribution P(A), ignoring covariates.

Source code in src/causalrl/imitation.py

class BehavioralCloning(Agent):
    """Naive imitator: clones the marginal action distribution ``P(A)``, ignoring covariates."""

    def __init__(self, n_actions: int, seed: int | None = None) -> None:
        self.n_actions = n_actions
        self._rng = np.random.default_rng(seed)
        self._probs = np.ones(n_actions) / n_actions

    def fit(self, demonstrations: Mapping[str, np.ndarray], *, action: str) -> None:
        actions = np.asarray(demonstrations[action]).astype(int)
        counts = np.bincount(actions, minlength=self.n_actions).astype(float)
        self._probs = counts / counts.sum()

    def act(self, observation: dict[str, Any]) -> int:
        return int(self._rng.choice(self.n_actions, p=self._probs))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        """No-op: imitation is fit offline from demonstrations."""

update(observation, action, reward)

No-op: imitation is fit offline from demonstrations.

Source code in src/causalrl/imitation.py

def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
    """No-op: imitation is fit offline from demonstrations."""

Causal Curriculum Learning

causalrl.curriculum.causal_curriculum(graph, goal=None)

A curriculum (skill order) respecting the causal structure: a topological order in which every parent (prerequisite) precedes its children. With goal, restrict to the goal and its ancestors — the skills the goal depends on — still in topological order.

Source code in src/causalrl/curriculum.py

def causal_curriculum(graph: CausalGraph, goal: str | None = None) -> list[str]:
    """A curriculum (skill order) respecting the causal structure: a topological order in which
    every parent (prerequisite) precedes its children. With ``goal``, restrict to the goal and its
    ancestors — the skills the goal depends on — still in topological order."""
    order = graph.topological_order()
    if goal is None:
        return order
    if goal not in graph.nodes:
        raise CausalGraphError(f"unknown goal skill: {goal!r}")
    relevant = graph.ancestors(goal)  # inclusive of goal
    return [node for node in order if node in relevant]

causalrl.curriculum.curriculum_q_learning(tasks, *, episodes_per_task, alpha=0.5, epsilon=0.1, seed=None)

Learn the target policy by Q-learning through a curriculum of subtasks, easiest first.

tasks is ordered from the simplest subtask to the target (the last element); all tasks share the same state and action spaces. The Q-table is carried forward between stages (warm-start transfer), so value learned on the easy subtasks bootstraps the harder ones. This is the causal curriculum applied to RL: order subgoals by prerequisite structure (see :func:causal_curriculum), then train in that order to reach a target policy that flat learning on the sparse target alone struggles to find. Returns the greedy policy on the target task.

Faithful to Y. Bengio, J. Louradour, R. Collobert, J. Weston, Curriculum Learning, ICML 2009; the Q-learning update matches :func:causalrl.shaping.q_learning. No external code is ported.

Source code in src/causalrl/curriculum.py

def curriculum_q_learning(
    tasks: Sequence[TabularMDP],
    *,
    episodes_per_task: int,
    alpha: float = 0.5,
    epsilon: float = 0.1,
    seed: int | None = None,
) -> dict[int, int]:
    """Learn the target policy by Q-learning through a curriculum of subtasks, easiest first.

    ``tasks`` is ordered from the simplest subtask to the target (the last element); all tasks share
    the same state and action spaces. The Q-table is carried forward between stages (warm-start
    transfer), so value learned on the easy subtasks bootstraps the harder ones. This is the causal
    curriculum applied to RL: order subgoals by prerequisite structure (see
    :func:`causal_curriculum`), then train in that order to reach a target policy that flat learning
    on the sparse target alone struggles to find. Returns the greedy policy on the target task.

    Faithful to Y. Bengio, J. Louradour, R. Collobert, J. Weston, *Curriculum Learning*, ICML 2009;
    the Q-learning update matches :func:`causalrl.shaping.q_learning`. No external code is ported.
    """
    if not tasks:
        raise CausalGraphError("curriculum must contain at least one task")
    target = tasks[-1]
    rng = np.random.default_rng(seed)
    q = np.zeros((target.n_states, target.n_actions))
    for task in tasks:
        for _ in range(episodes_per_task):
            s = 0
            for _ in range(4 * task.n_states):
                if s in task.terminals:
                    break
                if rng.random() < epsilon:
                    a = int(rng.integers(0, task.n_actions))
                else:
                    a = int(np.argmax(q[s]))
                s_next = task.transitions[(s, a)]
                bootstrap = 0.0 if s_next in task.terminals else float(np.max(q[s_next]))
                q[s, a] += alpha * (task.rewards[(s, a)] + task.gamma * bootstrap - q[s, a])
                s = s_next
    return {s: int(np.argmax(q[s])) for s in range(target.n_states) if s not in target.terminals}

causalrl.curriculum.is_valid_curriculum(graph, order)

Whether order respects prerequisites: each skill appears after every parent of it that is also in order.

Source code in src/causalrl/curriculum.py

def is_valid_curriculum(graph: CausalGraph, order: Sequence[str]) -> bool:
    """Whether ``order`` respects prerequisites: each skill appears after every parent of it that is
    also in ``order``."""
    present = set(order)
    seen: set[str] = set()
    for skill in order:
        prerequisites = set(graph.parents(skill)) & present
        if not prerequisites <= seen:
            return False
        seen.add(skill)
    return True

causalrl.curriculum.PrerequisiteLearner

Causally-gated skill acquisition: walking the curriculum left to right, a skill is mastered iff all of its parents (prerequisites) are already mastered. Deterministic and order-faithful, so the effect of the ordering is exactly readable from the mastered set.

Source code in src/causalrl/curriculum.py

class PrerequisiteLearner:
    """Causally-gated skill acquisition: walking the curriculum left to right, a skill is mastered
    iff all of its parents (prerequisites) are already mastered. Deterministic and order-faithful,
    so the effect of the ordering is exactly readable from the mastered set."""

    def __init__(self, graph: CausalGraph) -> None:
        self._graph = graph
        self._mastered: set[str] = set()

    def train(self, curriculum: Sequence[str]) -> frozenset[str]:
        """Process the curriculum once and return the set of mastered skills."""
        self._mastered = set()
        for skill in curriculum:
            if set(self._graph.parents(skill)) <= self._mastered:
                self._mastered.add(skill)
        return frozenset(self._mastered)

    def masters(self, skill: str) -> bool:
        """Whether ``skill`` was mastered by the most recent :meth:`train`."""
        return skill in self._mastered

train(curriculum)

Process the curriculum once and return the set of mastered skills.

Source code in src/causalrl/curriculum.py

def train(self, curriculum: Sequence[str]) -> frozenset[str]:
    """Process the curriculum once and return the set of mastered skills."""
    self._mastered = set()
    for skill in curriculum:
        if set(self._graph.parents(skill)) <= self._mastered:
            self._mastered.add(skill)
    return frozenset(self._mastered)

masters(skill)

Whether skill was mastered by the most recent :meth:train.

Source code in src/causalrl/curriculum.py

def masters(self, skill: str) -> bool:
    """Whether ``skill`` was mastered by the most recent :meth:`train`."""
    return skill in self._mastered

Causal Reward Shaping

causalrl.shaping.apply_potential_shaping(mdp, potential)

Return a new MDP with rewards[s, a] += gamma * Phi(s') - Phi(s) (policy-invariant).

Source code in src/causalrl/shaping.py

def apply_potential_shaping(mdp: TabularMDP, potential: Mapping[int, float]) -> TabularMDP:
    """Return a new MDP with ``rewards[s, a] += gamma * Phi(s') - Phi(s)`` (policy-invariant)."""
    shaped = dict(mdp.rewards)
    for (s, a), s_next in mdp.transitions.items():
        shaped[(s, a)] = mdp.rewards[(s, a)] + mdp.gamma * potential[s_next] - potential[s]
    return replace(mdp, rewards=shaped)

causalrl.shaping.causal_potential(mdp)

The ideal (causal) shaping potential: V* with terminal states pinned to 0 (the condition under which potential shaping is policy-invariant for episodic tasks).

Source code in src/causalrl/shaping.py

def causal_potential(mdp: TabularMDP) -> dict[int, float]:
    """The ideal (causal) shaping potential: ``V*`` with terminal states pinned to ``0`` (the
    condition under which potential shaping is policy-invariant for episodic tasks)."""
    values, _ = value_iteration(mdp)
    return {s: (0.0 if s in mdp.terminals else values[s]) for s in range(mdp.n_states)}

causalrl.shaping.value_iteration(mdp, *, tol=1e-09, max_iter=1000)

Return (V*, greedy optimal policy) for the deterministic MDP.

Source code in src/causalrl/shaping.py

def value_iteration(
    mdp: TabularMDP, *, tol: float = 1e-9, max_iter: int = 1000
) -> tuple[dict[int, float], dict[int, int]]:
    """Return ``(V*, greedy optimal policy)`` for the deterministic MDP."""
    values = dict.fromkeys(range(mdp.n_states), 0.0)
    for _ in range(max_iter):
        delta = 0.0
        for s in range(mdp.n_states):
            if s in mdp.terminals:
                continue
            best = max(
                mdp.rewards[(s, a)] + mdp.gamma * values[mdp.transitions[(s, a)]]
                for a in range(mdp.n_actions)
            )
            delta = max(delta, abs(best - values[s]))
            values[s] = best
        if delta < tol:
            break
    policy: dict[int, int] = {}
    for s in range(mdp.n_states):
        if s in mdp.terminals:
            continue
        q = [
            mdp.rewards[(s, a)] + mdp.gamma * values[mdp.transitions[(s, a)]]
            for a in range(mdp.n_actions)
        ]
        policy[s] = int(np.argmax(q))
    return values, policy

causalrl.shaping.q_learning(mdp, *, episodes, potential=None, alpha=0.5, epsilon=0.1, max_steps=None, seed=None)

Tabular epsilon-greedy Q-learning from state 0. With potential given, the reward is shaped online by gamma * Phi(s') - Phi(s). Returns the greedy policy.

Source code in src/causalrl/shaping.py

def q_learning(
    mdp: TabularMDP,
    *,
    episodes: int,
    potential: Mapping[int, float] | None = None,
    alpha: float = 0.5,
    epsilon: float = 0.1,
    max_steps: int | None = None,
    seed: int | None = None,
) -> dict[int, int]:
    """Tabular epsilon-greedy Q-learning from state ``0``. With ``potential`` given, the reward is
    shaped online by ``gamma * Phi(s') - Phi(s)``. Returns the greedy policy."""
    rng = np.random.default_rng(seed)
    q = np.zeros((mdp.n_states, mdp.n_actions))
    steps_cap = max_steps if max_steps is not None else 4 * mdp.n_states
    for _ in range(episodes):
        s = 0
        for _ in range(steps_cap):
            if s in mdp.terminals:
                break
            if rng.random() < epsilon:
                a = int(rng.integers(0, mdp.n_actions))
            else:
                a = int(np.argmax(q[s]))
            s_next = mdp.transitions[(s, a)]
            reward = mdp.rewards[(s, a)]
            if potential is not None:
                reward = reward + mdp.gamma * potential[s_next] - potential[s]
            bootstrap = 0.0 if s_next in mdp.terminals else float(np.max(q[s_next]))
            q[s, a] += alpha * (reward + mdp.gamma * bootstrap - q[s, a])
            s = s_next
    return {s: int(np.argmax(q[s])) for s in range(mdp.n_states) if s not in mdp.terminals}

causalrl.shaping.TabularMDP dataclass

A finite deterministic MDP: transitions/rewards are keyed by (state, action).

Source code in src/causalrl/shaping.py

@dataclass(frozen=True)
class TabularMDP:
    """A finite deterministic MDP: ``transitions``/``rewards`` are keyed by ``(state, action)``."""

    n_states: int
    n_actions: int
    transitions: dict[tuple[int, int], int]
    rewards: dict[tuple[int, int], float]
    terminals: frozenset[int]
    gamma: float

Causal Game Theory

causalrl.games.CausalGame dataclass

A finite game as a causal influence diagram.

utilities[agent] maps each joint action profile (a tuple in agents order) to agent's payoff; graph is the (M)ACID with a decision and utility node per agent.

Source code in src/causalrl/games.py

@dataclass(frozen=True)
class CausalGame:
    """A finite game as a causal influence diagram.

    ``utilities[agent]`` maps each joint action profile (a tuple in ``agents`` order) to ``agent``'s
    payoff; ``graph`` is the (M)ACID with a decision and utility node per agent.
    """

    agents: tuple[str, ...]
    actions: Mapping[str, tuple[int, ...]]
    utilities: Mapping[str, Mapping[tuple[int, ...], float]]
    graph: CausalGraph

    def __post_init__(self) -> None:
        profiles = set(product(*(self.actions[a] for a in self.agents)))
        for agent in self.agents:
            if agent not in self.utilities:
                raise CausalGraphError(f"missing utilities for agent {agent!r}")
            if set(self.utilities[agent]) != profiles:
                raise CausalGraphError(
                    f"utilities for agent {agent!r} must cover the joint action space exactly"
                )
            for node in (decision_node(agent), utility_node(agent)):
                if node not in self.graph.nodes:
                    raise CausalGraphError(f"influence diagram is missing node {node!r}")

causalrl.games.pure_nash_equilibria(game)

All pure-strategy Nash equilibria, by enumeration over the finite joint action space.

Source code in src/causalrl/games.py

def pure_nash_equilibria(game: CausalGame) -> list[dict[str, int]]:
    """All pure-strategy Nash equilibria, by enumeration over the finite joint action space."""
    equilibria: list[dict[str, int]] = []
    for combo in product(*(game.actions[a] for a in game.agents)):
        profile = dict(zip(game.agents, combo, strict=True))
        if is_nash_equilibrium(game, profile):
            equilibria.append(profile)
    return equilibria

causalrl.games.mixed_nash_equilibria(game)

All mixed-strategy Nash equilibria (pure and properly mixed), by support enumeration.

Two-player games are solved exactly over rational arithmetic (:class:fractions.Fraction), so symmetric games yield exact mixes (matching pennies gives 0.5/0.5). Games with three or more agents are solved by support enumeration with a numerical (Newton) solve of the multilinear indifference system; every returned profile is then verified to be an ε-Nash equilibrium (no agent can gain more than 1e-6 by deviating to a pure action). Each equilibrium maps an agent to {action: probability} with off-support actions at zero.

Assumes a non-degenerate game; degenerate games may admit a continuum of equilibria, of which only support-extreme points are enumerated, and the numerical solver may miss a support whose system is ill-conditioned. Raises :class:CausalGraphError for fewer than two agents (use :func:pure_nash_equilibria for pure equilibria of any game).

Faithful to the support-enumeration method (R. Porter, E. Nudelman, Y. Shoham, Simple Search Methods for Finding a Nash Equilibrium, Games and Economic Behavior 2008; B. von Stengel, Computing Equilibria for Two-Person Games, Handbook of Game Theory 2002). No code is ported.

Source code in src/causalrl/games.py

def mixed_nash_equilibria(game: CausalGame) -> list[dict[str, dict[int, float]]]:
    """All mixed-strategy Nash equilibria (pure and properly mixed), by support enumeration.

    Two-player games are solved **exactly** over rational arithmetic (:class:`fractions.Fraction`),
    so symmetric games yield exact mixes (matching pennies gives ``0.5``/``0.5``). Games with three
    or more agents are solved by support enumeration with a **numerical** (Newton) solve of the
    multilinear indifference system; every returned profile is then verified to be an ε-Nash
    equilibrium (no agent can gain more than ``1e-6`` by deviating to a pure action). Each
    equilibrium maps an agent to ``{action: probability}`` with off-support actions at zero.

    Assumes a non-degenerate game; degenerate games may admit a continuum of equilibria, of which
    only support-extreme points are enumerated, and the numerical solver may miss a support whose
    system is ill-conditioned. Raises :class:`CausalGraphError` for fewer than two agents (use
    :func:`pure_nash_equilibria` for pure equilibria of any game).

    Faithful to the support-enumeration method (R. Porter, E. Nudelman, Y. Shoham, *Simple Search
    Methods for Finding a Nash Equilibrium*, Games and Economic Behavior 2008; B. von Stengel,
    *Computing Equilibria for Two-Person Games*, Handbook of Game Theory 2002). No code is ported.
    """
    if len(game.agents) < 2:
        raise CausalGraphError("mixed_nash_equilibria needs at least two agents")
    if len(game.agents) == 2:
        return _mixed_nash_two_player(game)
    return _mixed_nash_general(game)

causalrl.games.best_response(game, agent, profile)

The actions maximizing agent's payoff given the other agents' actions in profile.

Source code in src/causalrl/games.py

def best_response(game: CausalGame, agent: str, profile: Mapping[str, int]) -> frozenset[int]:
    """The actions maximizing ``agent``'s payoff given the other agents' actions in ``profile``."""
    index = game.agents.index(agent)
    base = [profile[a] for a in game.agents]
    payoffs: dict[int, float] = {}
    for action in game.actions[agent]:
        joint = list(base)
        joint[index] = action
        payoffs[action] = game.utilities[agent][tuple(joint)]
    best = max(payoffs.values())
    return frozenset(action for action, value in payoffs.items() if value == best)

causalrl.games.is_nash_equilibrium(game, profile)

Whether every agent's action in profile is a best response to the others.

Source code in src/causalrl/games.py

def is_nash_equilibrium(game: CausalGame, profile: Mapping[str, int]) -> bool:
    """Whether every agent's action in ``profile`` is a best response to the others."""
    return all(profile[agent] in best_response(game, agent, profile) for agent in game.agents)

Exceptions

causalrl.exceptions.CausalRLError

Bases: Exception

Base class for all causalrl errors.

Source code in src/causalrl/exceptions.py

class CausalRLError(Exception):
    """Base class for all causalrl errors."""

causalrl.exceptions.CausalInterfaceUnavailableError

Bases: CausalRLError

The causal interface is not available on this wrapper.

Raised when a method that requires a live SCM and a named reward node is called on a :class:~causalrl.envs.wrapper.CausalEnvWrapper that was constructed without them (e.g. wrapping a :class:~causalrl.envs.base.ConfoundedMDP that carries scm=None, or without passing a reward_node).

Source code in src/causalrl/exceptions.py

class CausalInterfaceUnavailableError(CausalRLError):
    """The causal interface is not available on this wrapper.

    Raised when a method that requires a live SCM and a named reward node is called
    on a :class:`~causalrl.envs.wrapper.CausalEnvWrapper` that was constructed without
    them (e.g. wrapping a :class:`~causalrl.envs.base.ConfoundedMDP` that carries
    ``scm=None``, or without passing a ``reward_node``).
    """

causalrl.exceptions.NotIdentifiableError

Bases: CausalRLError

A causal query is not identifiable from the available data.

Source code in src/causalrl/exceptions.py

class NotIdentifiableError(CausalRLError):
    """A causal query is not identifiable from the available data."""

    def __init__(self, message: str, witness: object | None = None) -> None:
        super().__init__(message)
        self.witness = witness

causalrl.exceptions.CausalGraphError

Bases: CausalRLError

Invalid graph operation (unknown node, cycle, malformed edge).

Source code in src/causalrl/exceptions.py

class CausalGraphError(CausalRLError):
    """Invalid graph operation (unknown node, cycle, malformed edge)."""

causalrl.exceptions.RealizabilityError

Bases: CausalRLError

A counterfactual query cannot be realized from the given evidence.

Source code in src/causalrl/exceptions.py

class RealizabilityError(CausalRLError):
    """A counterfactual query cannot be realized from the given evidence."""

causalrl.exceptions.UnverifiedAssumptionError

Bases: CausalRLError

A method's claimed guarantee requires an assumption the caller has not declared.

Source code in src/causalrl/exceptions.py

class UnverifiedAssumptionError(CausalRLError):
    """A method's claimed guarantee requires an assumption the caller has not declared."""

Partial-Identification And OPE Bounds

causalrl.identification.bounds.manski_bounds(data, *, treatment, outcome, action, outcome_range=(0.0, 1.0))

Sharp no-assumptions bounds on E[outcome | do(treatment = action)] (Manski 1990).

From observational data (integer treatment column, numeric outcome in outcome_range): the units that took action contribute their observed mean, while the rest are bounded only by the outcome range. With p = P(treatment = action) and m = E[outcome | treatment = action] the bounds are [m*p + y_min*(1-p), m*p + y_max*(1-p)] — sharp, collapsing to a point when every unit took action. The observational counterpart of :func:~causalrl.ope.bounds.causal_q_bounds.

Source code in src/causalrl/identification/bounds.py

def manski_bounds(
    data: Mapping[str, Sequence[float]],
    *,
    treatment: str,
    outcome: str,
    action: int,
    outcome_range: tuple[float, float] = (0.0, 1.0),
) -> Interval:
    """Sharp no-assumptions bounds on ``E[outcome | do(treatment = action)]`` (Manski 1990).

    From observational ``data`` (integer ``treatment`` column, numeric ``outcome`` in
    ``outcome_range``): the units that took ``action`` contribute their observed mean, while
    the rest are bounded only by the outcome range. With ``p = P(treatment = action)`` and
    ``m = E[outcome | treatment = action]`` the bounds are
    ``[m*p + y_min*(1-p), m*p + y_max*(1-p)]`` — sharp, collapsing to a point when every unit took
    ``action``. The observational counterpart of :func:`~causalrl.ope.bounds.causal_q_bounds`.
    """
    x = np.asarray(data[treatment])
    y = np.asarray(data[outcome], dtype=float)
    y_min, y_max = outcome_range
    mask = x == action
    p = float(mask.mean())
    observed = float(y[mask].mean()) if bool(mask.any()) else 0.0
    return Interval(observed * p + y_min * (1.0 - p), observed * p + y_max * (1.0 - p))

causalrl.ope.bounds.ipw_sensitivity_bounds(outcomes, propensities, *, gamma, return_certificate=None)

ipw_sensitivity_bounds(outcomes: Sequence[float], propensities: Sequence[float], *, gamma: float, return_certificate: Literal[True] | None = ...) -> Certificate

ipw_sensitivity_bounds(outcomes: Sequence[float], propensities: Sequence[float], *, gamma: float, return_certificate: Literal[False]) -> Interval

Marginal-sensitivity-model bounds on the treated counterfactual mean E[Y(1)].

outcomes and propensities are the treated units' outcomes Y_i and nominal propensities e(Z_i) = P(treated | Z_i) (what an unconfounded model fits). Under Tan's marginal sensitivity model the true inverse-propensity weight lies within an odds-ratio factor gamma >= 1 of the nominal, giving a_i in [1 + (1/g)(1/e_i - 1), 1 + g(1/e_i - 1)]; the bounds are the extreme stabilized (Hájek) weighted means over that range. At gamma = 1 the interval collapses to the IPW point estimate; it widens monotonically with gamma and contains E[Y(1)] whenever the true confounding odds ratio is at most gamma.

Faithful to Z. Tan, A Distributional Approach for Causal Inference Using Propensity Scores (JASA 2006) and Q. Zhao, D. Small, B. Bhattacharya, Sensitivity Analysis for Inverse Probability Weighting Estimators via the Percentile Bootstrap (JRSS-B 2019). No code is ported.

return_certificate (2.0 default): returns a BOUNDED :class:Certificate; pass False for the legacy :class:Interval, or True (equivalently, the *_certified variant) to be explicit.

Source code in src/causalrl/ope/bounds.py

def ipw_sensitivity_bounds(
    outcomes: Sequence[float],
    propensities: Sequence[float],
    *,
    gamma: float,
    return_certificate: bool | None = None,
) -> Interval | Certificate:
    """Marginal-sensitivity-model bounds on the treated counterfactual mean ``E[Y(1)]``.

    ``outcomes`` and ``propensities`` are the treated units' outcomes ``Y_i`` and *nominal*
    propensities ``e(Z_i) = P(treated | Z_i)`` (what an unconfounded model fits). Under Tan's
    marginal sensitivity model the true inverse-propensity weight lies within an odds-ratio factor
    ``gamma >= 1`` of the nominal, giving ``a_i in [1 + (1/g)(1/e_i - 1), 1 + g(1/e_i - 1)]``; the
    bounds are the extreme stabilized (Hájek) weighted means over that range. At ``gamma = 1`` the
    interval collapses to the IPW point estimate; it widens monotonically with ``gamma`` and
    contains ``E[Y(1)]`` whenever the true confounding odds ratio is at most ``gamma``.

    Faithful to Z. Tan, *A Distributional Approach for Causal Inference Using Propensity Scores*
    (JASA 2006) and Q. Zhao, D. Small, B. Bhattacharya, *Sensitivity Analysis for Inverse
    Probability Weighting Estimators via the Percentile Bootstrap* (JRSS-B 2019). No code is ported.

    ``return_certificate`` (2.0 default): returns a ``BOUNDED`` :class:`Certificate`; pass ``False``
    for the legacy :class:`Interval`, or ``True`` (equivalently, the ``*_certified`` variant) to be
    explicit.
    """
    if return_certificate is None:
        return_certificate = True  # causalrl 2.0: return a Certificate by default
    if return_certificate:
        from causalrl.certify.routines import ipw_sensitivity_bounds_certified

        return ipw_sensitivity_bounds_certified(outcomes, propensities, gamma=gamma)
    if gamma < 1.0:
        raise ValueError("gamma must be >= 1")
    y = np.asarray(outcomes, dtype=float)
    e = np.asarray(propensities, dtype=float)
    odds = (1.0 - e) / e
    lo_w = 1.0 + odds / gamma
    hi_w = 1.0 + odds * gamma
    return Interval(
        _fractional_extreme(y, lo_w, hi_w, maximize=False),
        _fractional_extreme(y, lo_w, hi_w, maximize=True),
    )

causalrl.ope.bounds.causal_q_bounds(dataset, state, action, *, require_identified=False)

Manski natural bounds on E[return | do(action), state] from confounded logs.

For a return in [0, 1] with empirical mean m = E[R|s,a] and propensity p = P(a|s): lower = m * p, upper = m * p + (1 - p). A never-logged action (p = 0) yields the vacuous [0, 1] — not identifiable from the logs alone. With require_identified=True, a vacuous bound raises NotIdentifiableError carrying (state, action) as the witness.

Source code in src/causalrl/ope/bounds.py

def causal_q_bounds(
    dataset: ConfoundedTrajectoryDataset,
    state: int,
    action: int,
    *,
    require_identified: bool = False,
) -> Interval:
    """Manski natural bounds on E[return | do(action), state] from confounded logs.

    For a return in [0, 1] with empirical mean m = E[R|s,a] and propensity p = P(a|s):
        lower = m * p,  upper = m * p + (1 - p).
    A never-logged action (p = 0) yields the vacuous [0, 1] — not identifiable from the
    logs alone. With `require_identified=True`, a vacuous bound raises NotIdentifiableError
    carrying (state, action) as the witness.
    """
    p = dataset.behavior_propensity(state, action)
    m = dataset.mean_reward(state, action)
    lower = m * p
    upper = m * p + (1.0 - p)
    if require_identified and p == 0.0:
        raise NotIdentifiableError(
            f"E[R|do(a={action}), s={state}] is not identifiable: action never logged "
            f"in this state (vacuous bound [0, 1])",
            witness=(state, action),
        )
    return Interval(lower, upper)

causalrl.ope.bounds.msm_policy_value_bounds(outcomes, logging_propensities, target_propensities, *, gamma, return_certificate=None)

msm_policy_value_bounds(outcomes: Sequence[float], logging_propensities: Sequence[float], target_propensities: Sequence[float], *, gamma: float, return_certificate: Literal[True] | None = ...) -> Certificate

msm_policy_value_bounds(outcomes: Sequence[float], logging_propensities: Sequence[float], target_propensities: Sequence[float], *, gamma: float, return_certificate: Literal[False]) -> Interval

Marginal-sensitivity-model bounds on an off-policy value V(pi_t) = E[(pi_t/e0) Y].

Self-normalised (Hájek) off-policy value of a target policy pi_t estimated from logs of a logging policy with nominal propensities e0(a|x) = P(a | x) (a valid probability in (0, 1]). outcomes are the logged rewards Y_i; target_propensities are pi_t(a_i | x_i) at the logged action. Under Tan's marginal sensitivity model the true logging propensity deviates from nominal by an odds-ratio at most gamma >= 1, so the true inverse weight 1/ẽ0 lies in [1 + odds/gamma, 1 + odds*gamma] with odds = (1-e0)/e0; each unit's contribution weight is pi_t(a_i|x_i) * (1/ẽ0). The bounds are the extreme stabilised weighted means of Y over those per-unit weight ranges.

Reduces to :func:ipw_sensitivity_bounds when pi_t is constant across the logged actions (the treated / uniform-target mean — the constant cancels in the self-normalised ratio), and collapses to the self-normalised IPS point at gamma = 1. The off-policy generalisation of Tan's MSM in the spirit of N. Kallus & A. Zhou, Confounding-Robust Policy Evaluation in Infinite-Horizon Reinforcement Learning (NeurIPS 2020). The caller supplies pi_t and the nominal e0; no code is ported.

return_certificate (2.0 default): returns a BOUNDED :class:Certificate; pass False for the legacy :class:Interval, or True (equivalently, the *_certified variant) to be explicit.

Source code in src/causalrl/ope/bounds.py

def msm_policy_value_bounds(
    outcomes: Sequence[float],
    logging_propensities: Sequence[float],
    target_propensities: Sequence[float],
    *,
    gamma: float,
    return_certificate: bool | None = None,
) -> Interval | Certificate:
    """Marginal-sensitivity-model bounds on an off-policy value ``V(pi_t) = E[(pi_t/e0) Y]``.

    Self-normalised (Hájek) off-policy value of a target policy ``pi_t`` estimated from logs of a
    logging policy with *nominal* propensities ``e0(a|x) = P(a | x)`` (a valid probability in
    ``(0, 1]``). ``outcomes`` are the logged rewards ``Y_i``; ``target_propensities`` are
    ``pi_t(a_i | x_i)`` at the logged action. Under Tan's marginal sensitivity model the true
    logging propensity deviates from nominal by an odds-ratio at most ``gamma >= 1``, so the true
    inverse weight ``1/ẽ0`` lies in ``[1 + odds/gamma, 1 + odds*gamma]`` with ``odds = (1-e0)/e0``;
    each unit's contribution weight is ``pi_t(a_i|x_i) * (1/ẽ0)``. The bounds are the extreme
    stabilised weighted means of ``Y`` over those per-unit weight ranges.

    Reduces to :func:`ipw_sensitivity_bounds` when ``pi_t`` is constant across the logged actions
    (the treated / uniform-target mean — the constant cancels in the self-normalised ratio), and
    collapses to the self-normalised IPS point at ``gamma = 1``. The off-policy generalisation of
    Tan's MSM in the spirit of N. Kallus & A. Zhou, *Confounding-Robust Policy Evaluation in
    Infinite-Horizon Reinforcement Learning* (NeurIPS 2020). The caller supplies ``pi_t`` and the
    nominal ``e0``; no code is ported.

    ``return_certificate`` (2.0 default): returns a ``BOUNDED`` :class:`Certificate`; pass ``False``
    for the legacy :class:`Interval`, or ``True`` (equivalently, the ``*_certified`` variant) to be
    explicit.
    """
    if return_certificate is None:
        return_certificate = True  # causalrl 2.0: return a Certificate by default
    if return_certificate:
        from causalrl.certify.routines import msm_policy_value_bounds_certified

        return msm_policy_value_bounds_certified(
            outcomes, logging_propensities, target_propensities, gamma=gamma
        )
    if gamma < 1.0:
        raise ValueError("gamma must be >= 1")
    y = np.asarray(outcomes, dtype=float)
    e0 = np.asarray(logging_propensities, dtype=float)
    pt = np.asarray(target_propensities, dtype=float)
    odds = (1.0 - e0) / e0
    lo_w = pt * (1.0 + odds / gamma)
    hi_w = pt * (1.0 + odds * gamma)
    return Interval(
        _fractional_extreme(y, lo_w, hi_w, maximize=False),
        _fractional_extreme(y, lo_w, hi_w, maximize=True),
    )

causalrl.ope.bounds.msm_contribution_bounds(outcomes, logging_propensities, target_propensities_on, target_propensities_off, *, gamma)

Marginal-sensitivity-model bounds on a contribution V(pi_on) - V(pi_off).

The off-policy value DIFFERENCE between two target rules, estimated from confounded logs under Tan's marginal sensitivity model — e.g. a per-agent credit or per-factor contribution E[Y_{do(F=1)}] - E[Y_{do(F=0)}]. Each arm is bounded by :func:msm_policy_value_bounds (target_propensities_on = pi_on(a_i | x_i) at the logged action, ..._off likewise, shared nominal e0) and the contribution interval is the difference

[ on.lower - off.upper ,  on.upper - off.lower ].

Always valid (it contains the true difference for any targets); sharp when the two target supports are disjoint — e.g. the deterministic one-hot arms 1{F=1} / 1{F=0} that partition the logged units, so the two arms' weight perturbations are independent — and conservative otherwise. Collapses to the difference of the two self-normalised IPS points at gamma = 1 and widens monotonically with gamma.

Source code in src/causalrl/ope/bounds.py

def msm_contribution_bounds(
    outcomes: Sequence[float],
    logging_propensities: Sequence[float],
    target_propensities_on: Sequence[float],
    target_propensities_off: Sequence[float],
    *,
    gamma: float,
) -> Interval:
    """Marginal-sensitivity-model bounds on a *contribution* ``V(pi_on) - V(pi_off)``.

    The off-policy value DIFFERENCE between two target rules, estimated from confounded logs
    under Tan's marginal sensitivity model — e.g. a per-agent credit or per-factor contribution
    ``E[Y_{do(F=1)}] - E[Y_{do(F=0)}]``. Each arm is bounded by :func:`msm_policy_value_bounds`
    (``target_propensities_on`` = ``pi_on(a_i | x_i)`` at the logged action, ``..._off`` likewise,
    shared nominal ``e0``) and the contribution interval is the difference

        [ on.lower - off.upper ,  on.upper - off.lower ].

    Always *valid* (it contains the true difference for any targets); **sharp** when the two
    target supports are disjoint — e.g. the deterministic one-hot arms ``1{F=1}`` / ``1{F=0}``
    that partition the logged units, so the two arms' weight perturbations are independent — and
    *conservative* otherwise. Collapses to the difference of the two self-normalised IPS points at
    ``gamma = 1`` and widens monotonically with ``gamma``.
    """
    if gamma < 1.0:
        raise ValueError("gamma must be >= 1")
    on = msm_policy_value_bounds(
        outcomes,
        logging_propensities,
        target_propensities_on,
        gamma=gamma,
        return_certificate=False,
    )
    off = msm_policy_value_bounds(
        outcomes,
        logging_propensities,
        target_propensities_off,
        gamma=gamma,
        return_certificate=False,
    )
    return Interval(on.lower - off.upper, on.upper - off.lower)

causalrl.ope.bounds.msm_per_step_bounds(rewards_by_step, propensities_by_step, *, gamma)

Per-step marginal-sensitivity-model bounds on a cumulative (summed) reward.

Each element of rewards_by_step / propensities_by_step is one time step's per-unit rewards r_t and nominal propensities e_t; the cumulative-reward MSM bound is the sum over steps of the per-step :func:ipw_sensitivity_bounds. This is the additive (per-step) cumulative-reward MSM: each step is bounded independently under the sensitivity model and the bounds add, which is tight for sparse / per-step rewards.

Reusable kernel of the per-step cumulative-reward MSM used for confounded multi-step OPE (Bennett & Kallus, Efficient and Sharp OPE in Robust MDPs, NeurIPS 2024). The experiment supplies the per-step nominal propensities; no code is ported.

Source code in src/causalrl/ope/bounds.py

def msm_per_step_bounds(
    rewards_by_step: Sequence[Sequence[float]],
    propensities_by_step: Sequence[Sequence[float]],
    *,
    gamma: float,
) -> Interval:
    """Per-step marginal-sensitivity-model bounds on a cumulative (summed) reward.

    Each element of ``rewards_by_step`` / ``propensities_by_step`` is one time step's
    per-unit rewards ``r_t`` and nominal propensities ``e_t``; the cumulative-reward MSM
    bound is the sum over steps of the per-step :func:`ipw_sensitivity_bounds`. This is the
    additive (per-step) cumulative-reward MSM: each step is bounded independently under the
    sensitivity model and the bounds add, which is tight for sparse / per-step rewards.

    Reusable kernel of the per-step cumulative-reward MSM used for confounded multi-step OPE
    (Bennett & Kallus, *Efficient and Sharp OPE in Robust MDPs*, NeurIPS 2024). The
    experiment supplies the per-step nominal propensities; no code is ported.
    """
    if gamma < 1.0:
        raise ValueError("gamma must be >= 1")
    if len(rewards_by_step) != len(propensities_by_step):
        raise ValueError("rewards_by_step and propensities_by_step must have equal length")
    lower = upper = 0.0
    for r_t, e_t in zip(rewards_by_step, propensities_by_step, strict=True):
        iv = ipw_sensitivity_bounds(r_t, e_t, gamma=gamma, return_certificate=False)
        lower += iv.lower
        upper += iv.upper
    return Interval(lower, upper)

causalrl.ope.bounds.msm_stratified_bounds(values, propensities, strata, target_weights, *, gamma)

Stratified marginal-sensitivity-model bounds: Σ_s w_s · MSM_s.

Compute the MSM bound within each stratum (units sharing a strata label) and combine them with target_weights (the target stratum marginal w_s, e.g. a uniform initial state distribution). Strata absent from the data contribute nothing. Because conditioning removes between-stratum weight variation, the stratified bound is never wider than the pooled :func:ipw_sensitivity_bounds and never under-covers (THEORY.md, Prop 1).

The reusable kernel of the stratified cumulative-reward MSM; the experiment supplies the stratum labels (e.g. initial state) and nominal propensities.

Source code in src/causalrl/ope/bounds.py

def msm_stratified_bounds(
    values: Sequence[float],
    propensities: Sequence[float],
    strata: Sequence[int],
    target_weights: Mapping[int, float],
    *,
    gamma: float,
) -> Interval:
    """Stratified marginal-sensitivity-model bounds: ``Σ_s w_s · MSM_s``.

    Compute the MSM bound within each stratum (units sharing a ``strata`` label) and combine
    them with ``target_weights`` (the target stratum marginal ``w_s``, e.g. a uniform initial
    state distribution). Strata absent from the data contribute nothing. Because conditioning
    removes between-stratum weight variation, the stratified bound is never wider than the
    pooled :func:`ipw_sensitivity_bounds` and never under-covers (THEORY.md, Prop 1).

    The reusable kernel of the stratified cumulative-reward MSM; the experiment supplies the
    stratum labels (e.g. initial state) and nominal propensities.
    """
    if gamma < 1.0:
        raise ValueError("gamma must be >= 1")
    v = np.asarray(values, dtype=float)
    e = np.asarray(propensities, dtype=float)
    s = np.asarray(strata)
    lower = upper = 0.0
    for label, w in target_weights.items():
        mask = s == label
        if not bool(mask.any()):
            continue
        iv = ipw_sensitivity_bounds(
            v[mask].tolist(), e[mask].tolist(), gamma=gamma, return_certificate=False
        )
        lower += w * iv.lower
        upper += w * iv.upper
    return Interval(lower, upper)

causalrl.conformal.core.conformal_action_value(dataset, target_actions=None, *, alpha=0.1)

Calibrated band for the return of one decision taken under a target policy.

The off-policy caller of the weighted path. At logged transition i the likelihood ratio dP_target/dP_behavior is pi_target(a_i | s_i) / pi_behavior(a_i | s_i), computed here from the log's :meth:~causalrl.data.logged.LoggedDecisions.logging_propensities and target_actions[i] — the action the policy takes at that transition's state, one per logged transition, exactly the argument :func:causalrl.certify_policy takes. Passing target_actions=None scores the logging policy itself, whose logged returns need no reweighting; that is the reference the lower-bound gate compares against.

Returns a kind=EMPIRICAL certificate whose ci is the band and whose value is None. The claim is coverage of a single fresh return, not a confidence interval for V(pi) = E[return]: each end is a one-sided conformal bound at alpha / 2, so ci.lower is a distribution-free lower confidence bound on the next return under the target policy at level 1 - alpha/2, and the two-sided band covers at 1 - alpha. Too few effectively-weighted samples for the level returns an infinite end rather than a false one.

Valid under (i) weighted exchangeability of the logged returns; (ii) no unmeasured confounding of the logged action, without which the ratio above is not dP_target/dP_behavior; and (iii) positivity — every action the policy reaches must have been played by the logs in that state, else the band is vacuous and carries a positivity :class:~causalrl.certify.Hedge; a log that cannot check this carries the hedge too, rather than passing quietly. The state distribution is taken to be unchanged by the policy: the one-step / terminal-return regime :func:causalrl.certify_policy documents.

Source code in src/causalrl/conformal/core.py

def conformal_action_value(
    dataset: LoggedDecisions[ActionT],
    target_actions: Sequence[ActionT] | None = None,
    *,
    alpha: float = 0.1,
) -> Certificate:
    """Calibrated band for the return of one decision taken under a target policy.

    The off-policy caller of the weighted path. At logged transition ``i`` the likelihood ratio
    ``dP_target/dP_behavior`` is ``pi_target(a_i | s_i) / pi_behavior(a_i | s_i)``, computed here
    from the log's :meth:`~causalrl.data.logged.LoggedDecisions.logging_propensities` and
    ``target_actions[i]`` — the action the policy takes at that transition's state, one per
    logged transition, exactly the argument :func:`causalrl.certify_policy` takes. Passing
    ``target_actions=None`` scores the logging policy itself, whose logged returns need no
    reweighting; that is the reference the lower-bound gate compares against.

    Returns a ``kind=EMPIRICAL`` certificate whose ``ci`` is the band and whose ``value`` is
    ``None``. **The claim is coverage of a single fresh return, not a confidence interval for**
    ``V(pi) = E[return]``: each end is a one-sided conformal bound at ``alpha / 2``, so
    ``ci.lower`` is a distribution-free lower confidence bound on the next return under the target
    policy at level ``1 - alpha/2``, and the two-sided band covers at ``1 - alpha``. Too few
    effectively-weighted samples for the level returns an infinite end rather than a false one.

    Valid under (i) weighted exchangeability of the logged returns; (ii) no unmeasured confounding
    of the logged action, without which the ratio above is not ``dP_target/dP_behavior``; and
    (iii) positivity — every action the policy reaches must have been played by the logs in that
    state, else the band is vacuous and carries a positivity :class:`~causalrl.certify.Hedge`; a
    log that cannot check this carries the hedge too, rather than passing quietly. The
    state distribution is taken to be unchanged by the policy: the one-step / terminal-return
    regime :func:`causalrl.certify_policy` documents.
    """
    if not 0.0 < alpha < 1.0:
        raise ValueError("alpha must be in (0, 1)")
    returns = list(dataset.outcomes())
    weights: list[float] | None = None
    test_weight: float | None = None
    report = PositivityReport(checkable=True)
    if target_actions is not None:
        weights, test_weight, report = _propensity_ratios(dataset, target_actions)
    band = _return_band(returns, alpha, weights=weights, test_weight=test_weight)

    w = np.ones(len(returns)) if weights is None else np.asarray(weights, dtype=np.float64)
    sq = float(np.square(w).sum())
    ess = float(w.sum() ** 2 / sq) if sq > 0.0 else 0.0
    assumptions = [
        Assumption(
            name="weighted-exchangeability" if weights is not None else "exchangeability",
            params={"n_calibration": len(returns), "effective_sample_size": ess},
            checkable=False,
        )
    ]
    hedge: Hedge | None = None
    if target_actions is not None:
        assumptions.append(Assumption(name="no-unmeasured-confounding", params={}, checkable=False))
        assumptions.append(
            Assumption(
                name="positivity",
                params={"max_propensity_ratio": test_weight},
                checkable=report.checkable,
                diagnostic={"unsupported": list(report.gaps)},
            )
        )
        if not report.checkable:
            hedge = Hedge(
                reason=(
                    "positivity could not be checked: this log cannot report the behaviour "
                    "propensity of an action it never recorded, so the test-point ratio is "
                    "bounded by infinity rather than verified and the band is vacuous"
                ),
                detail={},
            )
        elif report.violated:
            hedge = Hedge(
                reason="positivity: the target policy takes actions the logs never played there",
                detail={"unsupported": list(report.gaps)},
            )
    label = "behavior" if target_actions is None else "target"
    return Certificate(
        claim=(
            f"P(return of one decision under the {label} policy in "
            f"[{band.lower:.4g}, {band.upper:.4g}]) >= {1.0 - alpha:.2f} (conformal); "
            "a fresh return, not E[return]"
        ),
        estimand=EstimandSpec(query="policy_value", target="quantile", policy=label),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=alpha,
        assumptions=tuple(assumptions),
        method="weighted-conformal-return-band" if weights is not None else "conformal-return-band",
        witness=None,
        hedge=hedge,
        provenance=Provenance.create(),
        ci=band,
    )

Decision Certificates

The decision stack — certify whether a confounded / off-policy decision ("is action 1 better than action 0?") is robust to hidden confounding, cheapest layer first. certify_decision is the one-call front door over the layers below.

causalrl.identification.decision.certify_decision(rewards=None, actions=None, *, confounder_bins=None, mi_cap=None, propensities=None, gamma_max=10.0, estimate=None)

Certify whether a binary decision from confounded logs is robust to hidden confounding.

rewards are the logged rewards Y_i; actions is the binary arm indicator (1 = the arm under test, 0 = baseline). The naive decision is the sign of the logged contrast E[Y | F=1] - E[Y | F=0]. Supply at least one evidence source — they compose:

• confounder_bins (a measured hidden variable Z) or mi_cap (a structural cap on the information channel MI(I; Z)) runs the sign-robustness layer (:func:pivotality_certificate): certifies iff no hidden confounder consistent with that information can flip the contrast's sign. One-sided — failure to certify is not evidence of a flip.
• propensities (the nominal logging propensities e0(a_i | x_i) at the logged action) runs the MSM tipping layer: the smallest Tan odds-ratio Gamma at which the off-policy (IPS) value contrast band first includes zero (:func:tipping_gamma over the sharp one-hot :func:msm_contribution_bounds). tipping_gamma is None means the decision is robust to confounding at least as strong as gamma_max.

With informative propensities the MSM layer concerns the inverse-propensity-weighted off-policy contrast, which coincides with the raw logged contrast only under uniform logging.

Alternatively pass a pre-built estimate (a :class:PolicyValueContrast, e.g. from an interop adapter) instead of raw rewards / actions; the two forms are mutually exclusive and certify_decision(estimate=c) is exactly certify_estimate(c).

Source code in src/causalrl/identification/decision.py

def certify_decision(
    rewards: Sequence[float] | None = None,
    actions: Sequence[int] | None = None,
    *,
    confounder_bins: Sequence[int] | None = None,
    mi_cap: float | None = None,
    propensities: Sequence[float] | None = None,
    gamma_max: float = 10.0,
    estimate: PolicyValueContrast | None = None,
) -> DecisionCertificate:
    """Certify whether a binary decision from confounded logs is robust to hidden confounding.

    ``rewards`` are the logged rewards ``Y_i``; ``actions`` is the binary arm indicator (1 = the
    arm under test, 0 = baseline). The naive decision is the sign of the logged contrast
    ``E[Y | F=1] - E[Y | F=0]``. Supply at least one evidence source — they compose:

    * ``confounder_bins`` (a measured hidden variable ``Z``) or ``mi_cap`` (a structural cap on
      the information channel ``MI(I; Z)``) runs the **sign-robustness layer**
      (:func:`pivotality_certificate`): certifies iff no hidden confounder consistent with that
      information can flip the contrast's sign. One-sided — failure to certify is not evidence of
      a flip.
    * ``propensities`` (the nominal logging propensities ``e0(a_i | x_i)`` at the logged action)
      runs the **MSM tipping layer**: the smallest Tan odds-ratio ``Gamma`` at which the
      off-policy (IPS) value contrast band first includes zero (:func:`tipping_gamma` over the
      sharp one-hot :func:`msm_contribution_bounds`). ``tipping_gamma is None`` means the decision
      is robust to confounding at least as strong as ``gamma_max``.

    With informative propensities the MSM layer concerns the inverse-propensity-weighted
    off-policy contrast, which coincides with the raw logged contrast only under uniform logging.

    Alternatively pass a pre-built ``estimate`` (a :class:`PolicyValueContrast`, e.g. from an
    interop adapter) instead of raw ``rewards`` / ``actions``; the two forms are mutually
    exclusive and ``certify_decision(estimate=c)`` is exactly ``certify_estimate(c)``.
    """
    if estimate is not None:
        if rewards is not None or actions is not None:
            raise ValueError("pass either raw logs (rewards, actions) or estimate, not both")
        return certify_estimate(estimate, gamma_max=gamma_max)
    if rewards is None or actions is None:
        raise ValueError("pass rewards and actions (or a pre-built estimate)")
    contrast = PolicyValueContrast.from_binary(
        rewards,
        actions,
        propensities=propensities,
        confounder_bins=confounder_bins,
        mi_cap=mi_cap,
    )
    return certify_estimate(contrast, gamma_max=gamma_max, labels=("action 1", "action 0"))

causalrl.identification.decision.DecisionCertificate

Bases: NamedTuple

Result of :func:certify_decision.

certified is the headline: is the naive (logged) decision robust to hidden confounding, by the strongest layer that ran? When the structural/measured layer ran it carries its one-sided guarantee (no confounder consistent with the supplied information can flip the sign); otherwise it reports whether the MSM layer found the decision robust up to gamma_max. When the finite-sample gate also ran (:func:causalrl.certify_policy with alpha), certified requires it too and conformal_lcb reports the bound it gated on. The component fields and summary make the exact guarantee explicit.

Source code in src/causalrl/identification/decision.py

class DecisionCertificate(NamedTuple):
    """Result of :func:`certify_decision`.

    ``certified`` is the headline: is the naive (logged) decision robust to hidden confounding,
    by the strongest layer that ran? When the structural/measured layer ran it carries its
    one-sided guarantee (no confounder consistent with the supplied information can flip the
    sign); otherwise it reports whether the MSM layer found the decision robust up to ``gamma_max``.
    When the finite-sample gate also ran (:func:`causalrl.certify_policy` with ``alpha``),
    ``certified`` requires it too and ``conformal_lcb`` reports the bound it gated on.
    The component fields and ``summary`` make the exact guarantee explicit.
    """

    decision: str  # "prefer action 1" | "prefer action 0" | "indifferent" (sign of naive_contrast)
    naive_contrast: float  # E[Y | F=1] - E[Y | F=0]
    certified: bool
    pivotality: PivotalityCertificate | None  # structural/measured sign-robustness layer, if run
    tipping_gamma: float | None  # MSM odds-ratio at which the decision tips; None if not run/robust
    msm_certified: bool | None  # MSM layer robust to gamma_max? None if the MSM layer did not run
    summary: str
    # Distribution-free lower bound on a fresh return under the policy under test, from the
    # conformal gate; None when that gate did not run. A single return, not E[return].
    conformal_lcb: float | None = None

    def __str__(self) -> str:
        return self.summary

    @property
    def recommendation(self) -> str:
        """Verdict-based action, resolving the naive-vs-adjusted footgun.

        ``"act"`` iff the naive :attr:`decision` is certified robust to hidden confounding by the
        strongest layer that ran; otherwise ``"abstain"``. Read THIS, not :attr:`decision`, to
        decide whether to act — :attr:`decision` is the (possibly confounded) hypothesis under
        test, not the recommendation.
        """
        return "act" if self.certified else "abstain"

recommendation property

Verdict-based action, resolving the naive-vs-adjusted footgun.

"act" iff the naive :attr:decision is certified robust to hidden confounding by the strongest layer that ran; otherwise "abstain". Read THIS, not :attr:decision, to decide whether to act — :attr:decision is the (possibly confounded) hypothesis under test, not the recommendation.

causalrl.identification.bounds.pivotality_certificate(outcomes, treated, confounder_bins=None, *, mi_cap=None)

One-sided sign-robustness certificate for a naive contrast under hidden confounding.

certified=True means: no hidden variable consistent with the supplied information can flip the sign of E[Y|F=1] - E[Y|F=0]. Two modes:

• confounder_bins given (a measured hidden variable, e.g. post-hoc showdown/oracle data): certify iff the TV-form :func:confounding_bias_bound is strictly below |naive|; mi_measured reports the plug-in channel.
• mi_cap given (a structural cap on MI(I;Z) from the environment's information rules — the data-processing route): certify iff mi_cap < mi_flip computed with the outcome-span relaxation (no Z needed anywhere).

The cheapest layer of the decision stack — certificate, then MSM band (:func:~causalrl.ope.bounds.msm_contribution_bounds), then abstention (:func:tipping_gamma). One-sided: failure to certify is NOT evidence of a flip. Verified against measured ground truth in three game-log regimes in experiments/games/theory/verify_pivotality.py.

Source code in src/causalrl/identification/bounds.py

def pivotality_certificate(
    outcomes: Sequence[float],
    treated: Sequence[int],
    confounder_bins: Sequence[int] | None = None,
    *,
    mi_cap: float | None = None,
) -> PivotalityCertificate:
    """One-sided sign-robustness certificate for a naive contrast under hidden confounding.

    ``certified=True`` means: no hidden variable consistent with the supplied information can
    flip the sign of ``E[Y|F=1] - E[Y|F=0]``. Two modes:

    * ``confounder_bins`` given (a *measured* hidden variable, e.g. post-hoc showdown/oracle
      data): certify iff the TV-form :func:`confounding_bias_bound` is strictly below
      ``|naive|``; ``mi_measured`` reports the plug-in channel.
    * ``mi_cap`` given (a *structural* cap on ``MI(I;Z)`` from the environment's information
      rules — the data-processing route): certify iff ``mi_cap < mi_flip`` computed with the
      outcome-span relaxation (no ``Z`` needed anywhere).

    The cheapest layer of the decision stack — certificate, then MSM band
    (:func:`~causalrl.ope.bounds.msm_contribution_bounds`), then abstention
    (:func:`tipping_gamma`). One-sided:
    failure to certify is NOT evidence of a flip. Verified against measured ground truth in
    three game-log regimes in ``experiments/games/theory/verify_pivotality.py``.
    """
    y = np.asarray(outcomes, dtype=float)
    f = np.asarray(treated, dtype=bool)
    if not (f.any() and (~f).any()):
        raise ValueError("both arms must be present")
    p = float(f.mean())
    naive = float(y[f].mean() - y[~f].mean())
    if confounder_bins is not None:
        bound = confounding_bias_bound(outcomes, treated, confounder_bins, form="tv")
        z = np.asarray(confounder_bins)
        bins = np.unique(z)
        pz = np.array([(z == b).mean() for b in bins])
        pz1 = np.array([((z == b) & f).sum() / f.sum() for b in bins])
        pz0 = np.array([((z == b) & ~f).sum() / (~f).sum() for b in bins])
        mi_measured: float | None = _plugin_mi(pz, pz1, pz0, p)
        m1 = [float(y[(z == b) & f].mean()) for b in bins]
        m0 = [float(y[(z == b) & ~f].mean()) for b in bins]
        mi_flip = mi_flip_threshold(naive, max(m1) - min(m1), max(m0) - min(m0), p)
    elif mi_cap is not None:
        if mi_cap < 0.0:
            raise ValueError("mi_cap must be >= 0")
        span = float(y.max() - y.min())
        mi_flip = mi_flip_threshold(naive, span, span, p)
        # sharp MI form with M1 = M0 = span: span * sqrt(C / (2 p (1-p))), trivially capped
        bound = min(span * float(np.sqrt(mi_cap / (2 * p * (1 - p)))), 2 * span)
        mi_measured = None
    else:
        raise ValueError("supply confounder_bins (measured Z) or mi_cap (structural cap)")
    return PivotalityCertificate(
        certified=bool(bound < abs(naive)),
        naive=naive,
        bias_bound=float(bound),
        mi_flip=float(mi_flip),
        mi_measured=mi_measured,
    )

causalrl.identification.bounds.confounding_bias_bound(outcomes, treated, confounder_bins, *, form='tv')

Upper bound on the omitted-variable bias |naive - Z-adjusted| from logged rows.

Omitted-variable bias in total-variation form:

``|bias| <= M1 * TV(P_{Z|F=1}, P_Z) + M0 * TV(P_{Z|F=0}, P_Z)``

with M_f the span over Z-bins of E[Y | F=f, Z] (each signed measure integrates to zero, so the midrange trick gives the span, not the sup; sharp — attained with equality by an explicit two-point family at every parameter value). form="mi" applies Pinsker with the KL budget p*KL1 + (1-p)*KL0 = MI(F;Z) split optimally across arms (sharp small-MI constant):

``|bias| <= sqrt( MI/2 * (M1^2/p + M0^2/(1-p)) )``  (capped at the trivial ``M1+M0``),

the (looser-than-TV) form whose budget the environment's information structure can cap. Every stratum must contain both arms (positivity); restrict to an overlap population first — a logger that conditions hard on Z destroys overlap, which is a finding, not a nuisance.

Source code in src/causalrl/identification/bounds.py

def confounding_bias_bound(
    outcomes: Sequence[float],
    treated: Sequence[int],
    confounder_bins: Sequence[int],
    *,
    form: str = "tv",
) -> float:
    """Upper bound on the omitted-variable bias ``|naive - Z-adjusted|`` from logged rows.

    Omitted-variable bias in total-variation form:

        ``|bias| <= M1 * TV(P_{Z|F=1}, P_Z) + M0 * TV(P_{Z|F=0}, P_Z)``

    with ``M_f`` the span over ``Z``-bins of ``E[Y | F=f, Z]`` (each signed measure integrates
    to zero, so the midrange trick gives the span, not the sup; sharp — attained with equality
    by an explicit two-point family at every parameter value). ``form="mi"`` applies Pinsker
    with the KL budget ``p*KL1 + (1-p)*KL0 = MI(F;Z)`` split optimally across arms (sharp
    small-MI constant):

        ``|bias| <= sqrt( MI/2 * (M1^2/p + M0^2/(1-p)) )``  (capped at the trivial ``M1+M0``),

    the (looser-than-TV) form whose budget the environment's information structure can cap.
    Every stratum must contain both arms (positivity); restrict to an overlap population first —
    a logger that conditions hard on ``Z`` destroys overlap, which is a finding, not a nuisance.
    """
    if form not in ("tv", "mi"):
        raise ValueError("form must be 'tv' or 'mi'")
    y = np.asarray(outcomes, dtype=float)
    f = np.asarray(treated, dtype=bool)
    z = np.asarray(confounder_bins)
    if not (y.shape == f.shape == z.shape):
        raise ValueError("outcomes, treated, confounder_bins must have equal length")
    if not (f.any() and (~f).any()):
        raise ValueError("both arms must be present")
    p = float(f.mean())
    bins = np.unique(z)
    m1: list[float] = []
    m0: list[float] = []
    pz: list[float] = []
    pz1: list[float] = []
    pz0: list[float] = []
    for b in bins:
        in_b = z == b
        if not ((in_b & f).any() and (in_b & ~f).any()):
            raise ValueError(
                f"stratum {b!r} lacks one arm (positivity violated) — restrict to an overlap "
                "population before bounding"
            )
        m1.append(float(y[in_b & f].mean()))
        m0.append(float(y[in_b & ~f].mean()))
        pz.append(float(in_b.mean()))
        pz1.append(float((in_b & f).sum() / f.sum()))
        pz0.append(float((in_b & ~f).sum() / (~f).sum()))
    span1 = max(m1) - min(m1)
    span0 = max(m0) - min(m0)
    pz_a = np.asarray(pz, dtype=float)
    pz1_a = np.asarray(pz1, dtype=float)
    pz0_a = np.asarray(pz0, dtype=float)
    if form == "tv":
        tv1 = 0.5 * float(np.abs(pz1_a - pz_a).sum())
        tv0 = 0.5 * float(np.abs(pz0_a - pz_a).sum())
        return span1 * tv1 + span0 * tv0
    mi = _plugin_mi(pz_a, pz1_a, pz0_a, p)
    sharp = float(np.sqrt(mi / 2.0 * (span1**2 / p + span0**2 / (1.0 - p))))
    return min(sharp, span1 + span0)  # trivial cap: TV_f <= 1

causalrl.identification.bounds.mi_flip_threshold(naive, span_treated, span_control, p_treated)

Channel capacity (nats) below which NO hidden confounder can flip the naive sign.

The decision-pivotality threshold (sharp form)

``MI_flip = 2 naive^2 / (M1^2/p + M0^2/(1-p))``:

if the mutual information between the logged binary action and a hidden variable Z is below this value, the omitted-Z bias is strictly smaller than |naive| (omitted- variable bias in TV form + Pinsker with the KL budget split optimally across arms), so the Z-adjusted contrast has the same sign as the naive one. This constant is SHARP: a binary-channel family attains the corresponding bias bound with ratio -> 1 in the small-MI limit (THEORY_pivotality.md, tightness section); the additive form 2(|naive|/(M1/sqrt(p)+M0/sqrt(1-p)))^2 is its (looser) Cauchy-Schwarz relaxation. Combined with the data-processing inequality MI(F;Z) <= MI(I;Z) — the logging actor's policy reads only its information set I — the information structure of the environment caps the reachable confounding budget.

span_treated / span_control are the spans of E[Y | F=f, Z=z] over z (the regression spans M1 / M0); with Z unmeasured, use the outcome span for both (valid, looser).

Source code in src/causalrl/identification/bounds.py

def mi_flip_threshold(
    naive: float,
    span_treated: float,
    span_control: float,
    p_treated: float,
) -> float:
    """Channel capacity (nats) below which NO hidden confounder can flip the naive sign.

    The decision-pivotality threshold (sharp form)

        ``MI_flip = 2 naive^2 / (M1^2/p + M0^2/(1-p))``:

    if the mutual information between the logged binary action and a hidden variable ``Z`` is
    below this value, the omitted-``Z`` bias is strictly smaller than ``|naive|`` (omitted-
    variable bias in TV form + Pinsker with the KL budget split optimally across arms), so the
    ``Z``-adjusted contrast has the same sign as the naive one. This constant is SHARP: a
    binary-channel family attains the corresponding bias bound with ratio -> 1 in the small-MI
    limit (THEORY_pivotality.md, tightness section); the additive form
    ``2(|naive|/(M1/sqrt(p)+M0/sqrt(1-p)))^2`` is its (looser) Cauchy-Schwarz relaxation.
    Combined with the data-processing inequality ``MI(F;Z) <= MI(I;Z)`` — the logging actor's
    policy reads only its information set ``I`` — the *information structure of the
    environment* caps the reachable confounding budget.

    ``span_treated`` / ``span_control`` are the spans of ``E[Y | F=f, Z=z]`` over ``z`` (the
    regression spans ``M1`` / ``M0``); with ``Z`` unmeasured, use the outcome span for both
    (valid, looser).
    """
    if not 0.0 < p_treated < 1.0:
        raise ValueError("p_treated must be in (0, 1)")
    if span_treated < 0.0 or span_control < 0.0:
        raise ValueError("spans must be non-negative")
    denom2 = span_treated**2 / p_treated + span_control**2 / (1.0 - p_treated)
    if denom2 == 0.0:
        return float("inf")
    return float(2.0 * naive**2 / denom2)

causalrl.identification.bounds.tipping_gamma(bound, *, reference=0.0, gamma_max=10.0, tol=0.001)

Sensitivity tipping point: the smallest gamma >= 1 at which the partial-ID interval bound(gamma) first contains reference.

This is the causal-sensitivity reporting companion to the MSM bound kernels: it answers "how strong would unmeasured confounding have to be (on the MSM/Rosenbaum odds-ratio scale) to overturn the conclusion that the estimand lies strictly on one side of reference?". A larger tipping gamma ⇒ a more robust conclusion — the odds-ratio-scale analog of the E-value (VanderWeele & Ding, Ann. Intern. Med. 2017) and of Rosenbaum's Gamma.

bound maps a sensitivity level gamma to an :class:Interval; it must collapse to a point at gamma = 1 and widen monotonically with gamma (as every MSM kernel here does — e.g. lambda g: msm_contribution_bounds(y, e0, on, off, gamma=g)). Returns 1.0 if the point already sits on reference, and None if the interval never reaches reference by gamma_max (the conclusion is robust to confounding at least that strong).

Source code in src/causalrl/identification/bounds.py

def tipping_gamma(
    bound: Callable[[float], Interval],
    *,
    reference: float = 0.0,
    gamma_max: float = 10.0,
    tol: float = 1e-3,
) -> float | None:
    """Sensitivity tipping point: the smallest ``gamma >= 1`` at which the partial-ID interval
    ``bound(gamma)`` first contains ``reference``.

    This is the causal-sensitivity *reporting* companion to the MSM bound kernels: it answers
    "how strong would unmeasured confounding have to be (on the MSM/Rosenbaum odds-ratio scale)
    to overturn the conclusion that the estimand lies strictly on one side of ``reference``?".
    A larger tipping ``gamma`` ⇒ a more robust conclusion — the odds-ratio-scale analog of the
    E-value (VanderWeele & Ding, *Ann. Intern. Med.* 2017) and of Rosenbaum's ``Gamma``.

    ``bound`` maps a sensitivity level ``gamma`` to an :class:`Interval`; it must collapse to a
    point at ``gamma = 1`` and widen monotonically with ``gamma`` (as every MSM kernel here does —
    e.g. ``lambda g: msm_contribution_bounds(y, e0, on, off, gamma=g)``). Returns ``1.0`` if the
    point already sits on ``reference``, and ``None`` if the interval never reaches ``reference``
    by ``gamma_max`` (the conclusion is robust to confounding at least that strong).
    """
    if gamma_max < 1.0:
        raise ValueError("gamma_max must be >= 1")
    eps = 1e-12
    at1 = bound(1.0)
    if at1.lower - eps <= reference <= at1.upper + eps:
        return 1.0
    if not (bound(gamma_max).lower - eps <= reference <= bound(gamma_max).upper + eps):
        return None  # robust: never reaches `reference` within [1, gamma_max]
    lo, hi = 1.0, gamma_max
    while hi - lo > tol:
        mid = 0.5 * (lo + hi)
        iv = bound(mid)
        if iv.lower - eps <= reference <= iv.upper + eps:
            hi = mid
        else:
            lo = mid
    return float(hi)

Unified Certificates (§5.2)

causalrl.certify.certificate.Certificate dataclass

A serializable certified claim (plan §5.2).

Source code in src/causalrl/certify/certificate.py

@dataclass(frozen=True)
class Certificate:
    """A serializable certified claim (plan §5.2)."""

    claim: str
    estimand: EstimandSpec
    kind: Kind
    value: float | Interval | None
    alpha: float | None
    assumptions: tuple[Assumption, ...]
    method: str
    witness: Witness | None
    hedge: Hedge | None
    provenance: Provenance
    ci: Interval | None = None  # optional confidence interval for `value` at level `alpha`

    def __str__(self) -> str:
        parts = [f"[{self.kind.name}] {self.claim}"]
        v = _format_value(self.value)
        if v is not None:
            parts.append(f"value={v}")
        if self.ci is not None:
            parts.append(f"ci={_format_value(self.ci)}")
        if self.hedge is not None:
            parts.append(f"HEDGE: {self.hedge.reason}")
        return " | ".join(parts)

    def to_dict(self) -> dict[str, Any]:
        return {
            "schema_version": _SCHEMA_VERSION,
            "claim": self.claim,
            "estimand": {
                "query": self.estimand.query,
                "target": self.estimand.target,
                "policy": self.estimand.policy,
                "domains": list(self.estimand.domains),
            },
            "kind": self.kind.value,
            "value": _value_to_json(self.value),
            "ci": _value_to_json(self.ci),
            "alpha": self.alpha,
            "assumptions": [
                {
                    "name": a.name,
                    "params": a.params,
                    "checkable": a.checkable,
                    "diagnostic": a.diagnostic,
                }
                for a in self.assumptions
            ],
            "method": self.method,
            "witness": (
                None
                if self.witness is None
                else {"kind": self.witness.kind, "detail": self.witness.detail}
            ),
            "hedge": (
                None
                if self.hedge is None
                else {
                    "reason": self.hedge.reason,
                    "detail": self.hedge.detail,
                    "downgraded_from": self.hedge.downgraded_from,
                }
            ),
            "provenance": {
                "library_version": self.provenance.library_version,
                "seeds": list(self.provenance.seeds),
                "data_fingerprint": self.provenance.data_fingerprint,
                "graph_hash": self.provenance.graph_hash,
                "timestamp": self.provenance.timestamp,
            },
        }

    def to_json(self) -> str:
        return json.dumps(self.to_dict())

    @classmethod
    def from_dict(cls, d: Mapping[str, Any]) -> Certificate:
        est = d["estimand"]
        prov = d["provenance"]
        wit = d.get("witness")
        hed = d.get("hedge")
        ci_raw = d.get("ci")
        ci: Interval | None = None
        if ci_raw is not None:
            parsed = _value_from_json(ci_raw)
            ci = parsed if isinstance(parsed, Interval) else None
        return cls(
            claim=d["claim"],
            estimand=EstimandSpec(
                query=est["query"],
                target=est.get("target", "mean"),
                policy=est.get("policy"),
                domains=tuple(est.get("domains", ())),
            ),
            kind=Kind(d["kind"]),
            value=_value_from_json(d["value"]),
            alpha=d["alpha"],
            assumptions=tuple(
                Assumption(
                    name=a["name"],
                    params=dict(a.get("params", {})),
                    checkable=a.get("checkable", False),
                    diagnostic=a.get("diagnostic"),
                )
                for a in d["assumptions"]
            ),
            method=d["method"],
            witness=None
            if wit is None
            else Witness(kind=wit["kind"], detail=dict(wit.get("detail", {}))),
            hedge=(
                None
                if hed is None
                else Hedge(
                    reason=hed["reason"],
                    detail=hed.get("detail"),
                    downgraded_from=hed.get("downgraded_from"),
                )
            ),
            provenance=Provenance(
                library_version=prov["library_version"],
                seeds=tuple(prov.get("seeds", ())),
                data_fingerprint=prov.get("data_fingerprint"),
                graph_hash=prov.get("graph_hash"),
                timestamp=prov.get("timestamp", ""),
            ),
            ci=ci,
        )

    @staticmethod
    def from_json(s: str) -> Certificate:
        return Certificate.from_dict(json.loads(s))

causalrl.certify.adapters.as_certificate(obj)

Convert a shipped certificate to a unified :class:Certificate (additive view).

Accepts DecisionCertificate, PivotalityCertificate, TransportRegretCertificate, or a Certificate (returned unchanged). Raises :class:TypeError for anything else.

Source code in src/causalrl/certify/adapters.py

def as_certificate(obj: object) -> Certificate:
    """Convert a shipped certificate to a unified :class:`Certificate` (additive view).

    Accepts ``DecisionCertificate``, ``PivotalityCertificate``, ``TransportRegretCertificate``, or a
    ``Certificate`` (returned unchanged). Raises :class:`TypeError` for anything else.
    """
    if isinstance(obj, Certificate):
        return obj
    if isinstance(obj, DecisionCertificate):
        return _decision_to_certificate(obj)
    if isinstance(obj, PivotalityCertificate):
        return _pivotality_to_certificate(obj)
    if isinstance(obj, TransportRegretCertificate):
        return _transport_regret_to_certificate(obj)
    raise TypeError(f"cannot adapt {type(obj).__name__} to a Certificate")

Identification-Aware Estimation (Phase 1)

causalrl.estimate.compiler.certify_effect(graph, treatment, outcome, data, *, method='dml', alpha=0.05, n_folds=5, seed=0, outcome_model=None, propensity_model=None, overlap_eps=0.01, clip=0.001)

Certify the back-door ATE E[Y|do(X=1)] - E[Y|do(X=0)] of a binary treatment.

Returns a kind=IDENTIFIED :class:Certificate with the point estimate (value), its confidence interval (ci), the adjustment-set witness, and provenance — or a hedged certificate (value=None, hedge set) when the effect is not identifiable, needs an unsupported ID method, or when estimated positivity falls below overlap_eps (I3). data maps variable names to arrays (see :func:~causalrl.estimate.estimators.estimate_ate).

Source code in src/causalrl/estimate/compiler.py

def certify_effect(
    graph: CausalGraph,
    treatment: str,
    outcome: str,
    data: Mapping[str, Any],
    *,
    method: str = "dml",
    alpha: float = 0.05,
    n_folds: int = 5,
    seed: int = 0,
    outcome_model: OutcomeFactory | None = None,
    propensity_model: PropensityFactory | None = None,
    overlap_eps: float = 0.01,
    clip: float = 1e-3,
) -> Certificate:
    """Certify the back-door ATE ``E[Y|do(X=1)] - E[Y|do(X=0)]`` of a binary ``treatment``.

    Returns a ``kind=IDENTIFIED`` :class:`Certificate` with the point estimate (``value``), its
    confidence interval (``ci``), the adjustment-set witness, and provenance — or a hedged
    certificate (``value=None``, ``hedge`` set) when the effect is not identifiable, needs an
    unsupported ID method, or when estimated positivity falls below ``overlap_eps`` (I3). ``data``
    maps variable names to arrays (see :func:`~causalrl.estimate.estimators.estimate_ate`).
    """
    try:
        plan = compile_estimand(graph, treatment, outcome)
    except NotIdentifiableError as exc:
        return _hedged(
            treatment,
            outcome,
            alpha,
            Hedge(reason="not-identifiable", detail={"message": str(exc)}),
            graph,
        )
    except EstimandNotSupportedError as exc:
        return _hedged(
            treatment,
            outcome,
            alpha,
            Hedge(reason="estimand-unsupported", detail={"message": str(exc)}),
            graph,
        )

    est: EffectEstimate = estimate_ate(
        data,
        treatment,
        outcome,
        plan.adjustment_set,
        method=method,
        alpha=alpha,
        n_folds=n_folds,
        seed=seed,
        outcome_model=outcome_model,
        propensity_model=propensity_model,
        clip=clip,
    )

    min_e = est.overlap.get("min_propensity", float("nan"))
    max_e = est.overlap.get("max_propensity", float("nan"))
    if np.isfinite(min_e) and (min_e < overlap_eps or max_e > 1.0 - overlap_eps):
        return _hedged(
            treatment,
            outcome,
            alpha,
            Hedge(
                reason="overlap-violation",
                detail={"overlap_eps": overlap_eps, **est.overlap},
            ),
            graph,
        )

    return Certificate(
        claim=f"E[{outcome}|do({treatment})=1] - E[{outcome}|do({treatment})=0] = {est.value:.4g}",
        estimand=EstimandSpec(query="do", target="mean"),
        kind=Kind.IDENTIFIED,
        value=est.value,
        alpha=alpha,
        assumptions=(
            Assumption(
                name="backdoor",
                params={"adjustment_set": list(plan.adjustment_set)},
                checkable=True,
            ),
            Assumption(
                name="overlap",
                params={"eps": overlap_eps},
                checkable=True,
                diagnostic=est.overlap,
            ),
        ),
        method=f"{est.method}" + (f" (cross-fit K={est.n_folds})" if est.n_folds else ""),
        witness=Witness(
            kind="adjustment",
            detail={"set": list(plan.adjustment_set), "estimand": plan.estimand_render},
        ),
        hedge=None,
        provenance=Provenance.create(
            seeds=(seed,),
            data_fingerprint=_data_fingerprint(data),
            graph_hash=graph_hash(graph),
        ),
        ci=est.ci,
    )

causalrl.bounds.continuous.msm_sensitivity_bounds(outcomes, *, gamma, propensities=None, treatment=None, covariates=None, propensity_model=None)

MSM bounds on E[Y(1)] with estimated logging propensities (Dorn-Guo/Kallus line).

Two modes (note outcomes differs between them):

• propensities= given -> treated units' outcomes + nominal propensities; delegates to ipw_sensitivity_bounds unchanged (the known-propensity case it reduces to exactly).
• treatment= + covariates= given -> outcomes/treatment/covariates are the full per-logged-unit arrays; fits e_hat = P(T=1 | covariates) with propensity_model (default logistic) and applies the same sharp box kernel to the treated subset.

Source code in src/causalrl/bounds/continuous.py

def msm_sensitivity_bounds(
    outcomes: Sequence[float],
    *,
    gamma: float,
    propensities: Sequence[float] | None = None,
    treatment: Sequence[float] | None = None,
    covariates: Any | None = None,
    propensity_model: Classifier | None = None,
) -> Interval:
    """MSM bounds on ``E[Y(1)]`` with *estimated* logging propensities (Dorn-Guo/Kallus line).

    Two modes (note ``outcomes`` differs between them):

    * ``propensities=`` given -> treated units' outcomes + nominal propensities; delegates to
      ``ipw_sensitivity_bounds`` unchanged (the *known*-propensity case it reduces to exactly).
    * ``treatment=`` + ``covariates=`` given -> ``outcomes``/``treatment``/``covariates`` are the
      *full per-logged-unit* arrays; fits ``e_hat = P(T=1 | covariates)`` with ``propensity_model``
      (default logistic) and applies the same sharp box kernel to the treated subset.
    """
    if propensities is not None:
        return ipw_sensitivity_bounds(outcomes, propensities, gamma=gamma, return_certificate=False)
    if treatment is None or covariates is None:
        raise ValueError("provide either propensities=, or both treatment= and covariates=")
    y = np.asarray(outcomes, dtype=np.float64)
    t = np.asarray(treatment, dtype=np.float64)
    z = np.asarray(covariates, dtype=np.float64)
    model: Classifier = propensity_model if propensity_model is not None else LogisticRegressor()
    proba = np.asarray(model.fit(z, t).predict_proba(z), dtype=np.float64)
    e_hat = proba[:, 1] if proba.ndim == 2 else proba
    treated = t == 1
    return ipw_sensitivity_bounds(
        y[treated].tolist(), e_hat[treated].tolist(), gamma=gamma, return_certificate=False
    )

causalrl.transport.estimate.certify_transported_effect(diagram, source_data, target_data, *, treatment, outcome, treated_value=1.0, max_adjustment_size=3)

Certify the transported mean E*[outcome | do(treatment=treated_value)] from data.

Uses the shipped :func:~causalrl.identification.transport.transport_formula to decide transportability, then estimates the closed form from source_data / target_data (numpy, discrete adjustment covariates). Returns a kind=IDENTIFIED certificate with the point value and the transport witness, or a hedged certificate when the effect is not transportable (I3).

Source code in src/causalrl/transport/estimate.py

def certify_transported_effect(
    diagram: SelectionDiagram,
    source_data: Mapping[str, Any],
    target_data: Mapping[str, Any],
    *,
    treatment: str,
    outcome: str,
    treated_value: float = 1.0,
    max_adjustment_size: int = 3,
) -> Certificate:
    """Certify the transported mean ``E*[outcome | do(treatment=treated_value)]`` from data.

    Uses the shipped :func:`~causalrl.identification.transport.transport_formula` to decide
    transportability, then estimates the closed form from ``source_data`` / ``target_data`` (numpy,
    discrete adjustment covariates). Returns a ``kind=IDENTIFIED`` certificate with the point value
    and the transport witness, or a hedged certificate when the effect is not transportable (I3).
    """
    formula = transport_formula(
        diagram, treatment=treatment, outcome=outcome, max_adjustment_size=max_adjustment_size
    )
    selection = sorted(diagram.selection_variables)
    if formula is None:
        return _hedged(treatment, outcome, selection, diagram, "non-transportable")

    if formula.kind == "direct":
        try:
            adjustment = sorted(backdoor_adjustment_set(diagram.graph, treatment, outcome))
        except NotIdentifiableError:
            return _hedged(treatment, outcome, selection, diagram, "source-effect-unsupported")
        marginal: Mapping[str, Any] = source_data  # direct: source interventional law transfers
        witness_kind = "direct-transport"
    else:
        adjustment = sorted(formula.adjustment_set)
        marginal = target_data  # adjustment: reweight by the target covariate marginal
        witness_kind = "s-admissible-adjustment"

    value = transport_gcomp(
        source_data,
        marginal,
        treatment=treatment,
        outcome=outcome,
        adjustment=adjustment,
        treated_value=treated_value,
    )
    return Certificate(
        claim=f"E*[{outcome} | do({treatment}={treated_value:g})] = {value:.4g}",
        estimand=EstimandSpec(query="transport", target="mean", domains=("source", "target")),
        kind=Kind.IDENTIFIED,
        value=value,
        alpha=None,
        assumptions=(
            Assumption(name="selection-diagram", params={"selection": selection}, checkable=True),
        ),
        method=f"transport-{formula.kind}",
        witness=Witness(
            kind=witness_kind,
            detail={
                "adjustment_set": adjustment,
                "selection": selection,
                "formula": formula.expression,
            },
        ),
        hedge=None,
        provenance=Provenance.create(graph_hash=graph_hash(diagram.graph)),
    )

Multi-Agent Core (Phase 2)

causalrl.magames.population.Population dataclass

N named agents instantiated from :class:AgentType templates with explicit sharing.

Source code in src/causalrl/magames/population.py

@dataclass(frozen=True)
class Population:
    """``N`` named agents instantiated from :class:`AgentType` templates with explicit sharing."""

    agents: tuple[str, ...]
    types: Mapping[str, AgentType]
    params: Mapping[str, Mapping[str, float]] = field(default_factory=_empty_params)
    topology: LearnerTopology = LearnerTopology.SINGLE_LEARNER

    def __post_init__(self) -> None:
        missing = [a for a in self.agents if a not in self.types]
        if missing:
            raise ValueError(f"no AgentType for agents {missing}")

    @property
    def max_kind(self) -> Kind:
        """The strongest ``Kind`` a certificate about this population may claim (per topology)."""
        return topology_max_kind(self.topology)

    def actions(self) -> dict[str, tuple[int, ...]]:
        return {a: self.types[a].actions for a in self.agents}

    def to_game(self) -> CausalGame:
        """Materialise the normal-form :class:`~causalrl.games.CausalGame` (small ``N``).

        Per-agent graphs stay factored in the templates; this fully materialises the joint payoff
        tables and a MAID with an edge from every decision to every utility (a general normal-form
        game: each payoff depends on all actions).
        """
        agents = self.agents
        action_map = self.actions()
        utilities: dict[str, dict[tuple[int, ...], float]] = {}
        for a in agents:
            t = self.types[a]
            idx = agents.index(a)
            p: Mapping[str, float] = self.params.get(a, {})
            table: dict[tuple[int, ...], float] = {}
            for combo in product(*(action_map[x] for x in agents)):
                others = tuple(c for j, c in enumerate(combo) if j != idx)
                table[combo] = float(t.payoff(combo[idx], others, p))
            utilities[a] = table
        nodes = [n for a in agents for n in (decision_node(a), utility_node(a))]
        edges = [(decision_node(i), utility_node(j)) for i in agents for j in agents]
        graph = CausalGraph(directed_edges=edges, nodes=nodes)
        return CausalGame(agents=agents, actions=action_map, utilities=utilities, graph=graph)

max_kind property

The strongest Kind a certificate about this population may claim (per topology).

to_game()

Materialise the normal-form :class:~causalrl.games.CausalGame (small N).

Per-agent graphs stay factored in the templates; this fully materialises the joint payoff tables and a MAID with an edge from every decision to every utility (a general normal-form game: each payoff depends on all actions).

Source code in src/causalrl/magames/population.py

def to_game(self) -> CausalGame:
    """Materialise the normal-form :class:`~causalrl.games.CausalGame` (small ``N``).

    Per-agent graphs stay factored in the templates; this fully materialises the joint payoff
    tables and a MAID with an edge from every decision to every utility (a general normal-form
    game: each payoff depends on all actions).
    """
    agents = self.agents
    action_map = self.actions()
    utilities: dict[str, dict[tuple[int, ...], float]] = {}
    for a in agents:
        t = self.types[a]
        idx = agents.index(a)
        p: Mapping[str, float] = self.params.get(a, {})
        table: dict[tuple[int, ...], float] = {}
        for combo in product(*(action_map[x] for x in agents)):
            others = tuple(c for j, c in enumerate(combo) if j != idx)
            table[combo] = float(t.payoff(combo[idx], others, p))
        utilities[a] = table
    nodes = [n for a in agents for n in (decision_node(a), utility_node(a))]
    edges = [(decision_node(i), utility_node(j)) for i in agents for j in agents]
    graph = CausalGraph(directed_edges=edges, nodes=nodes)
    return CausalGame(agents=agents, actions=action_map, utilities=utilities, graph=graph)

causalrl.magames.empirical.EmpiricalGame dataclass

A finite game estimated by replication: per-profile payoff means, standard errors and n.

All three mappings are keyed by a profile of strategy labels in agent order. Build one with :meth:from_samples rather than by hand — it is the constructor that computes the error.

Source code in src/causalrl/magames/empirical.py

@dataclass(frozen=True)
class EmpiricalGame:
    """A finite game estimated by replication: per-profile payoff means, standard errors and ``n``.

    All three mappings are keyed by a profile of *strategy labels* in agent order. Build one with
    :meth:`from_samples` rather than by hand — it is the constructor that computes the error.
    """

    strategies: tuple[str, ...]
    agents: tuple[str, ...]
    means: Mapping[Profile, tuple[float, ...]]
    stderrs: Mapping[Profile, tuple[float, ...]]
    replications: Mapping[Profile, int]

    @classmethod
    def from_samples(
        cls,
        strategies: Sequence[str],
        samples: Samples,
        *,
        agents: Sequence[str] | None = None,
    ) -> EmpiricalGame:
        """Average a complete round-robin into a game, keeping each cell's standard error.

        ``samples`` maps a profile of strategy labels to the replications observed for it, each
        replication being one payoff per agent. Every profile in the product of ``strategies`` must
        be present with at least two replications: a cell with one sample has no error bar, and an
        object whose contract is to carry the error bar must not invent one.
        """
        pool = tuple(strategies)
        duplicates = sorted({s for s in pool if pool.count(s) > 1})
        if duplicates:
            raise ValueError(f"duplicate strategy labels: {duplicates}")
        arity = len(next(iter(samples))) if samples else 0
        names = tuple(agents) if agents is not None else _default_agents(arity)
        if agents is not None and arity and len(names) != arity:
            raise ValueError(f"{len(names)} agents given but profiles have arity {arity}")

        means: dict[Profile, tuple[float, ...]] = {}
        stderrs: dict[Profile, tuple[float, ...]] = {}
        counts: dict[Profile, int] = {}
        for profile in product(pool, repeat=arity):
            if profile not in samples:
                raise ValueError(f"round-robin is missing the profile {profile}")
            replications = samples[profile]
            if len(replications) < 2:
                raise ValueError(
                    f"profile {profile} has {len(replications)} replication(s); at least 2 are "
                    "needed to estimate its error"
                )
            if any(len(r) != arity for r in replications):
                raise ValueError(
                    f"profile {profile} needs one payoff per agent in every replication"
                )
            observed = np.asarray(replications, dtype=np.float64)
            means[profile] = tuple(float(v) for v in observed.mean(axis=0))
            spread = observed.std(axis=0, ddof=1) / np.sqrt(observed.shape[0])
            stderrs[profile] = tuple(float(v) for v in spread)
            counts[profile] = observed.shape[0]
        return cls(pool, names, means, stderrs, counts)

    def payoff_error(self, *, alpha: float = 0.05, functional_terms: int = 1) -> PayoffError:
        """The cellwise error bound this table licenses, at level ``alpha``.

        A normal interval on each cell's mean, Bonferroni-corrected over every (profile, agent)
        cell so the bound holds simultaneously — which is what the certificate needs, since its
        linear programs touch all of them at once. The reported ``utility`` is the widest such
        half-width, because the relaxation it feeds is a single scalar over all constraints.

        ``functional_terms`` is how many estimated payoffs the functional being certified sums: 1
        for one agent's payoff (the usual case), ``len(agents)`` for welfare, 0 for a functional
        that does not read the payoffs at all.
        """
        if functional_terms < 0:
            raise ValueError("functional_terms must be nonnegative")
        cells = sum(len(v) for v in self.stderrs.values())
        widest = max((max(v, default=0.0) for v in self.stderrs.values()), default=0.0)
        z = float(norm_ppf(1.0 - alpha / (2.0 * max(cells, 1))))
        utility = z * widest
        return PayoffError(utility=utility, functional=functional_terms * utility, alpha=alpha)

    def to_game(self) -> CausalGame:
        """Materialise the normal form over integer actions, using the estimated means as payoffs.

        Strategy ``i`` is action ``i``. The returned game says nothing about its own error — pass
        :meth:`payoff_error` alongside it to :func:`~causalrl.magames.cce.certify_cce_do`.
        """
        actions = tuple(range(len(self.strategies)))
        types = {
            agent: AgentType(name=agent, actions=actions, payoff=self._payoff_for(position))
            for position, agent in enumerate(self.agents)
        }
        return Population(agents=self.agents, types=types).to_game()

    def _payoff_for(self, position: int) -> PayoffTemplate:
        labels = self.strategies

        def payoff(own: int, others: tuple[int, ...], params: Mapping[str, float]) -> float:
            actions = list(others)
            actions.insert(position, own)
            return self.means[tuple(labels[a] for a in actions)][position]

        return payoff

from_samples(strategies, samples, *, agents=None) classmethod

Average a complete round-robin into a game, keeping each cell's standard error.

samples maps a profile of strategy labels to the replications observed for it, each replication being one payoff per agent. Every profile in the product of strategies must be present with at least two replications: a cell with one sample has no error bar, and an object whose contract is to carry the error bar must not invent one.

Source code in src/causalrl/magames/empirical.py

@classmethod
def from_samples(
    cls,
    strategies: Sequence[str],
    samples: Samples,
    *,
    agents: Sequence[str] | None = None,
) -> EmpiricalGame:
    """Average a complete round-robin into a game, keeping each cell's standard error.

    ``samples`` maps a profile of strategy labels to the replications observed for it, each
    replication being one payoff per agent. Every profile in the product of ``strategies`` must
    be present with at least two replications: a cell with one sample has no error bar, and an
    object whose contract is to carry the error bar must not invent one.
    """
    pool = tuple(strategies)
    duplicates = sorted({s for s in pool if pool.count(s) > 1})
    if duplicates:
        raise ValueError(f"duplicate strategy labels: {duplicates}")
    arity = len(next(iter(samples))) if samples else 0
    names = tuple(agents) if agents is not None else _default_agents(arity)
    if agents is not None and arity and len(names) != arity:
        raise ValueError(f"{len(names)} agents given but profiles have arity {arity}")

    means: dict[Profile, tuple[float, ...]] = {}
    stderrs: dict[Profile, tuple[float, ...]] = {}
    counts: dict[Profile, int] = {}
    for profile in product(pool, repeat=arity):
        if profile not in samples:
            raise ValueError(f"round-robin is missing the profile {profile}")
        replications = samples[profile]
        if len(replications) < 2:
            raise ValueError(
                f"profile {profile} has {len(replications)} replication(s); at least 2 are "
                "needed to estimate its error"
            )
        if any(len(r) != arity for r in replications):
            raise ValueError(
                f"profile {profile} needs one payoff per agent in every replication"
            )
        observed = np.asarray(replications, dtype=np.float64)
        means[profile] = tuple(float(v) for v in observed.mean(axis=0))
        spread = observed.std(axis=0, ddof=1) / np.sqrt(observed.shape[0])
        stderrs[profile] = tuple(float(v) for v in spread)
        counts[profile] = observed.shape[0]
    return cls(pool, names, means, stderrs, counts)

payoff_error(*, alpha=0.05, functional_terms=1)

The cellwise error bound this table licenses, at level alpha.

A normal interval on each cell's mean, Bonferroni-corrected over every (profile, agent) cell so the bound holds simultaneously — which is what the certificate needs, since its linear programs touch all of them at once. The reported utility is the widest such half-width, because the relaxation it feeds is a single scalar over all constraints.

functional_terms is how many estimated payoffs the functional being certified sums: 1 for one agent's payoff (the usual case), len(agents) for welfare, 0 for a functional that does not read the payoffs at all.

Source code in src/causalrl/magames/empirical.py

def payoff_error(self, *, alpha: float = 0.05, functional_terms: int = 1) -> PayoffError:
    """The cellwise error bound this table licenses, at level ``alpha``.

    A normal interval on each cell's mean, Bonferroni-corrected over every (profile, agent)
    cell so the bound holds simultaneously — which is what the certificate needs, since its
    linear programs touch all of them at once. The reported ``utility`` is the widest such
    half-width, because the relaxation it feeds is a single scalar over all constraints.

    ``functional_terms`` is how many estimated payoffs the functional being certified sums: 1
    for one agent's payoff (the usual case), ``len(agents)`` for welfare, 0 for a functional
    that does not read the payoffs at all.
    """
    if functional_terms < 0:
        raise ValueError("functional_terms must be nonnegative")
    cells = sum(len(v) for v in self.stderrs.values())
    widest = max((max(v, default=0.0) for v in self.stderrs.values()), default=0.0)
    z = float(norm_ppf(1.0 - alpha / (2.0 * max(cells, 1))))
    utility = z * widest
    return PayoffError(utility=utility, functional=functional_terms * utility, alpha=alpha)

to_game()

Materialise the normal form over integer actions, using the estimated means as payoffs.

Strategy i is action i. The returned game says nothing about its own error — pass :meth:payoff_error alongside it to :func:~causalrl.magames.cce.certify_cce_do.

Source code in src/causalrl/magames/empirical.py

def to_game(self) -> CausalGame:
    """Materialise the normal form over integer actions, using the estimated means as payoffs.

    Strategy ``i`` is action ``i``. The returned game says nothing about its own error — pass
    :meth:`payoff_error` alongside it to :func:`~causalrl.magames.cce.certify_cce_do`.
    """
    actions = tuple(range(len(self.strategies)))
    types = {
        agent: AgentType(name=agent, actions=actions, payoff=self._payoff_for(position))
        for position, agent in enumerate(self.agents)
    }
    return Population(agents=self.agents, types=types).to_game()

causalrl.magames.cce.PayoffError dataclass

How far a measured payoff table may be from the game it estimates.

utility bounds |u_hat - u_true| for every agent and profile; functional bounds the same for the values of the functional being certified (0 when the functional does not read the estimated payoffs — an indicator of a profile, say — and k times utility when it sums k of them). alpha records the level at which the two bounds hold when they came from standard errors rather than from a hard envelope; it is reported, not used in the arithmetic.

Passing one to :func:certify_cce_do is what stops an estimated game being certified as if it were exact: see that function for where each of the two numbers enters.

Source code in src/causalrl/magames/cce.py

@dataclass(frozen=True)
class PayoffError:
    """How far a *measured* payoff table may be from the game it estimates.

    ``utility`` bounds ``|u_hat - u_true|`` for every agent and profile; ``functional`` bounds the
    same for the values of the functional being certified (0 when the functional does not read the
    estimated payoffs — an indicator of a profile, say — and ``k`` times ``utility`` when it sums
    ``k`` of them). ``alpha`` records the level at which the two bounds hold when they came from
    standard errors rather than from a hard envelope; it is reported, not used in the arithmetic.

    Passing one to :func:`certify_cce_do` is what stops an estimated game being certified as if it
    were exact: see that function for where each of the two numbers enters.
    """

    utility: float
    functional: float = 0.0
    alpha: float | None = None

    def __post_init__(self) -> None:
        if self.utility < 0.0 or self.functional < 0.0:
            raise ValueError("payoff error bounds must be nonnegative")
        if self.alpha is not None and not 0.0 < self.alpha < 1.0:
            raise ValueError("alpha must lie in (0, 1)")

    @property
    def exact(self) -> bool:
        """Whether this error is no error at all — the table is exact and licenses everything."""
        return self.utility == 0.0 and self.functional == 0.0

exact property

Whether this error is no error at all — the table is exact and licenses everything.

causalrl.magames.learning.run_no_regret(population, rounds, *, do=None, algorithm='regret_matching', learning_rate=None, explore=0.0, seed=0, trace_points=8)

Play rounds rounds of no-regret learning and return the realized empirical joint.

Every agent not pinned by do gets its own learner and, each round, the counterfactual payoff vector of what its actions would have paid against the profile the rest of the population actually played (full-information feedback, available because the game's payoffs are known). Agents pinned by do play their forced action and do not learn — the same restriction :func:~causalrl.magames.cce.cce_polytope applies, so the returned run is feasible for the intervened polytope by construction.

algorithm selects :class:~causalrl.agents.no_regret.RegretMatching (parameter-free) or :class:~causalrl.agents.no_regret.MultiplicativeWeights (Hedge; learning_rate defaults to the theory rate for this horizon and the agent's own payoff range). explore mixes in uniform play, trace_points sets how many log-spaced horizons the regret trace records.

The canonical use is the finite-time certificate, which assumes nothing about the limit::

run = run_no_regret(population, 20_000, do={"A2": 0})
cert = certify_cce_do(game, functional, do={"A2": 0}, no_regret=False, epsilon=run.regret)

Source code in src/causalrl/magames/learning.py

def run_no_regret(
    population: Population | CausalGame,
    rounds: int,
    *,
    do: Mapping[str, int] | None = None,
    algorithm: NoRegretAlgorithm = "regret_matching",
    learning_rate: float | None = None,
    explore: float = 0.0,
    seed: int = 0,
    trace_points: int = 8,
) -> NoRegretRun:
    """Play ``rounds`` rounds of no-regret learning and return the realized empirical joint.

    Every agent not pinned by ``do`` gets its own learner and, each round, the counterfactual payoff
    vector of what its actions would have paid against the profile the rest of the population
    actually played (full-information feedback, available because the game's payoffs are known).
    Agents pinned by ``do`` play their forced action and do not learn — the same restriction
    :func:`~causalrl.magames.cce.cce_polytope` applies, so the returned run is feasible for the
    intervened polytope by construction.

    ``algorithm`` selects :class:`~causalrl.agents.no_regret.RegretMatching` (parameter-free) or
    :class:`~causalrl.agents.no_regret.MultiplicativeWeights` (Hedge; ``learning_rate`` defaults to
    the theory rate for this horizon and the agent's own payoff range). ``explore`` mixes in uniform
    play, ``trace_points`` sets how many log-spaced horizons the regret trace records.

    The canonical use is the finite-time certificate, which assumes nothing about the limit::

        run = run_no_regret(population, 20_000, do={"A2": 0})
        cert = certify_cce_do(game, functional, do={"A2": 0}, no_regret=False, epsilon=run.regret)
    """
    game = population.to_game() if isinstance(population, Population) else population
    if rounds < 1:
        raise ValueError("rounds must be at least 1")
    pinned = dict(do or {})
    polytope = cce_polytope(game, do=pinned)  # validates `do` and fixes the profile order
    agents = game.agents
    action_values = {a: list(game.actions[a]) for a in agents}
    free = [a for a in agents if a not in pinned]
    seeds = np.random.default_rng(seed).integers(0, 2**31 - 1, size=len(free))
    learners = {
        agent: _make_learner(
            game,
            agent,
            algorithm=algorithm,
            rounds=rounds,
            learning_rate=learning_rate,
            explore=explore,
            seed=int(seeds[i]),
        )
        for i, agent in enumerate(free)
    }
    positions = {agent: agents.index(agent) for agent in agents}
    marks = _trace_marks(rounds, trace_points)

    counts: dict[tuple[int, ...], int] = {}
    trace: list[tuple[int, float]] = []
    for played in range(1, rounds + 1):
        profile = tuple(
            pinned[a] if a in pinned else action_values[a][learners[a].act()] for a in agents
        )
        counts[profile] = counts.get(profile, 0) + 1
        for agent in free:
            index = positions[agent]
            table = game.utilities[agent]
            learners[agent].observe(
                np.array(
                    [
                        table[(*profile[:index], action, *profile[index + 1 :])]
                        for action in action_values[agent]
                    ],
                    dtype=np.float64,
                )
            )
        if played in marks:
            running = {p: c / played for p, c in counts.items()}
            trace.append((played, cce_regret(game, running, do=pinned)))

    order = {profile: j for j, profile in enumerate(polytope.profiles)}
    weights = np.zeros(len(polytope.profiles))
    for profile, count in counts.items():
        weights[order[profile]] = count / rounds
    return NoRegretRun(
        agents=agents,
        profiles=polytope.profiles,
        weights=weights,
        rounds=rounds,
        algorithm=algorithm,
        do=pinned,
        regret_trace=tuple(trace),
    )

causalrl.magames.learning.NoRegretRun dataclass

Realized play of a no-regret population: the empirical joint plus its regret trace.

weights is aligned with profiles, which is the profile order of :func:~causalrl.magames.cce.cce_polytope for the same do — so both weights and :attr:empirical_joint drop straight into :func:~causalrl.magames.cce.cce_regret and :func:~causalrl.magames.cce.certify_cce_do. regret_trace pairs a horizon with the measured regret of the time average up to it; it is the evidence that the population is learning.

Source code in src/causalrl/magames/learning.py

@dataclass(frozen=True)
class NoRegretRun:
    """Realized play of a no-regret population: the empirical joint plus its regret trace.

    ``weights`` is aligned with ``profiles``, which is the profile order of
    :func:`~causalrl.magames.cce.cce_polytope` for the same ``do`` — so both ``weights`` and
    :attr:`empirical_joint` drop straight into :func:`~causalrl.magames.cce.cce_regret` and
    :func:`~causalrl.magames.cce.certify_cce_do`. ``regret_trace`` pairs a horizon with the measured
    regret of the time average up to it; it is the evidence that the population is learning.
    """

    agents: tuple[str, ...]
    profiles: tuple[tuple[int, ...], ...]
    weights: FloatArray
    rounds: int
    algorithm: str
    do: Mapping[str, int]
    regret_trace: tuple[tuple[int, float], ...]

    @property
    def empirical_joint(self) -> dict[tuple[int, ...], float]:
        """The realized joint distribution as ``{profile: probability}`` over played profiles."""
        return {
            profile: float(weight)
            for profile, weight in zip(self.profiles, self.weights, strict=True)
            if weight > 0.0
        }

    @property
    def regret(self) -> float:
        """The measured realized regret at the full horizon — pass it as ``epsilon``."""
        return self.regret_trace[-1][1]

    def marginal(self, agent: str) -> dict[int, float]:
        """What one agent played: the realized joint summed over everyone else.

        Keyed by the agent's own actions, including any it never played (weight 0.0), so the values
        always sum to 1. An agent pinned by ``do`` did not learn, so its marginal is the
        intervention itself.
        """
        if agent not in self.agents:
            raise KeyError(f"unknown agent: {agent!r}")
        index = self.agents.index(agent)
        played: dict[int, float] = {}
        for profile, weight in zip(self.profiles, self.weights, strict=True):
            action = profile[index]
            played[action] = played.get(action, 0.0) + float(weight)
        return played

    def boundary_mass(self, agent: str) -> float:
        """Realized mass on the extreme actions of ``agent``'s action set.

        Meaningful only when the actions are an *ordered* grid — a discretisation of a quantity,
        not a set of labelled strategies — in which case this is the truncation diagnostic: mass
        pressed against the smallest or largest action means the play may be capped by where the
        grid was stopped rather than settled where the game puts it, and the honest response is to
        widen the grid and re-run. An agent with a single available action scores 1.0.
        """
        played = self.marginal(agent)
        extremes = {min(played), max(played)}
        return sum(played[action] for action in extremes)

empirical_joint property

The realized joint distribution as {profile: probability} over played profiles.

regret property

The measured realized regret at the full horizon — pass it as epsilon.

marginal(agent)

What one agent played: the realized joint summed over everyone else.

Keyed by the agent's own actions, including any it never played (weight 0.0), so the values always sum to 1. An agent pinned by do did not learn, so its marginal is the intervention itself.

Source code in src/causalrl/magames/learning.py

def marginal(self, agent: str) -> dict[int, float]:
    """What one agent played: the realized joint summed over everyone else.

    Keyed by the agent's own actions, including any it never played (weight 0.0), so the values
    always sum to 1. An agent pinned by ``do`` did not learn, so its marginal is the
    intervention itself.
    """
    if agent not in self.agents:
        raise KeyError(f"unknown agent: {agent!r}")
    index = self.agents.index(agent)
    played: dict[int, float] = {}
    for profile, weight in zip(self.profiles, self.weights, strict=True):
        action = profile[index]
        played[action] = played.get(action, 0.0) + float(weight)
    return played

boundary_mass(agent)

Realized mass on the extreme actions of agent's action set.

Meaningful only when the actions are an ordered grid — a discretisation of a quantity, not a set of labelled strategies — in which case this is the truncation diagnostic: mass pressed against the smallest or largest action means the play may be capped by where the grid was stopped rather than settled where the game puts it, and the honest response is to widen the grid and re-run. An agent with a single available action scores 1.0.

Source code in src/causalrl/magames/learning.py

def boundary_mass(self, agent: str) -> float:
    """Realized mass on the extreme actions of ``agent``'s action set.

    Meaningful only when the actions are an *ordered* grid — a discretisation of a quantity,
    not a set of labelled strategies — in which case this is the truncation diagnostic: mass
    pressed against the smallest or largest action means the play may be capped by where the
    grid was stopped rather than settled where the game puts it, and the honest response is to
    widen the grid and re-run. An agent with a single available action scores 1.0.
    """
    played = self.marginal(agent)
    extremes = {min(played), max(played)}
    return sum(played[action] for action in extremes)

causalrl.agents.no_regret.NoRegretLearner

Bases: Agent

Base class for a learner over n_actions whose external regret grows sublinearly.

explore mixes the algorithm's distribution with the uniform one ((1 - explore) * p + explore / K). Leave it at 0 under full-information feedback (:meth:observe); set it positive under bandit feedback (:meth:update), where it bounds the inverse-propensity weights.

:meth:act samples the current mixed strategy. A repeated finite game has no state to condition on, so observation is ignored here; it is present because it is part of the :class:~causalrl.agents.base.Agent interface, and subclasses that do condition on it should say so.

Source code in src/causalrl/agents/no_regret.py

class NoRegretLearner(Agent):
    """Base class for a learner over ``n_actions`` whose external regret grows sublinearly.

    ``explore`` mixes the algorithm's distribution with the uniform one (``(1 - explore) * p +
    explore / K``). Leave it at ``0`` under full-information feedback (:meth:`observe`); set it
    positive under bandit feedback (:meth:`update`), where it bounds the inverse-propensity weights.

    :meth:`act` samples the current mixed strategy. A repeated finite game has no state to condition
    on, so ``observation`` is ignored here; it is present because it is part of the
    :class:`~causalrl.agents.base.Agent` interface, and subclasses that *do* condition on it should
    say so.
    """

    def __init__(self, n_actions: int, *, explore: float = 0.0, seed: int | None = None) -> None:
        if n_actions < 1:
            raise ValueError("n_actions must be at least 1")
        if not 0.0 <= explore <= 1.0:
            raise ValueError("explore must lie in [0, 1]")
        self.n_actions = int(n_actions)
        self.explore = float(explore)
        self.rounds = 0
        self._rng = np.random.default_rng(seed)

    @abstractmethod
    def _strategy(self) -> FloatArray:
        """The algorithm's own distribution over actions, before the ``explore`` mixing."""

    @abstractmethod
    def _absorb(self, payoffs: FloatArray, strategy: FloatArray) -> None:
        """Fold one counterfactual payoff vector into the algorithm's state."""

    def distribution(self) -> FloatArray:
        """The mixed strategy actually played: the algorithm's, mixed with uniform ``explore``."""
        strategy = self._strategy()
        if self.explore == 0.0:
            return strategy
        return (1.0 - self.explore) * strategy + self.explore / self.n_actions

    def observe(self, payoffs: FloatArray | list[float]) -> None:
        """Learn from the counterfactual payoff vector: ``payoffs[k]`` is what action ``k`` paid."""
        vector = np.asarray(payoffs, dtype=np.float64)
        if vector.shape != (self.n_actions,):
            raise ValueError(f"payoffs must have shape ({self.n_actions},), got {vector.shape}")
        self._absorb(vector, self.distribution())
        self.rounds += 1

    def act(self, observation: dict[str, Any] | None = None) -> int:
        """Sample an action from :meth:`distribution` (``observation`` is not conditioned on)."""
        return int(self._rng.choice(self.n_actions, p=self.distribution()))

    def update(self, observation: dict[str, Any] | None, action: int, reward: float) -> None:
        """Learn from bandit feedback: one realized ``reward`` for the action actually played.

        The single observation is expanded into an unbiased estimate of the full payoff vector by
        inverse-propensity weighting (``reward / p[action]`` on the played coordinate, ``0``
        elsewhere) before the ordinary full-information update. With ``explore > 0`` this is EXP3
        for :class:`MultiplicativeWeights`; ``observation`` is not conditioned on.
        """
        if not 0 <= action < self.n_actions:
            raise ValueError(f"action {action} outside 0..{self.n_actions - 1}")
        played = self.distribution()
        estimate = np.zeros(self.n_actions)
        estimate[action] = float(reward) / max(float(played[action]), 1e-12)
        self._absorb(estimate, played)
        self.rounds += 1

distribution()

The mixed strategy actually played: the algorithm's, mixed with uniform explore.

Source code in src/causalrl/agents/no_regret.py

def distribution(self) -> FloatArray:
    """The mixed strategy actually played: the algorithm's, mixed with uniform ``explore``."""
    strategy = self._strategy()
    if self.explore == 0.0:
        return strategy
    return (1.0 - self.explore) * strategy + self.explore / self.n_actions

observe(payoffs)

Learn from the counterfactual payoff vector: payoffs[k] is what action k paid.

Source code in src/causalrl/agents/no_regret.py

def observe(self, payoffs: FloatArray | list[float]) -> None:
    """Learn from the counterfactual payoff vector: ``payoffs[k]`` is what action ``k`` paid."""
    vector = np.asarray(payoffs, dtype=np.float64)
    if vector.shape != (self.n_actions,):
        raise ValueError(f"payoffs must have shape ({self.n_actions},), got {vector.shape}")
    self._absorb(vector, self.distribution())
    self.rounds += 1

act(observation=None)

Sample an action from :meth:distribution (observation is not conditioned on).

Source code in src/causalrl/agents/no_regret.py

def act(self, observation: dict[str, Any] | None = None) -> int:
    """Sample an action from :meth:`distribution` (``observation`` is not conditioned on)."""
    return int(self._rng.choice(self.n_actions, p=self.distribution()))

update(observation, action, reward)

Learn from bandit feedback: one realized reward for the action actually played.

The single observation is expanded into an unbiased estimate of the full payoff vector by inverse-propensity weighting (reward / p[action] on the played coordinate, 0 elsewhere) before the ordinary full-information update. With explore > 0 this is EXP3 for :class:MultiplicativeWeights; observation is not conditioned on.

Source code in src/causalrl/agents/no_regret.py

def update(self, observation: dict[str, Any] | None, action: int, reward: float) -> None:
    """Learn from bandit feedback: one realized ``reward`` for the action actually played.

    The single observation is expanded into an unbiased estimate of the full payoff vector by
    inverse-propensity weighting (``reward / p[action]`` on the played coordinate, ``0``
    elsewhere) before the ordinary full-information update. With ``explore > 0`` this is EXP3
    for :class:`MultiplicativeWeights`; ``observation`` is not conditioned on.
    """
    if not 0 <= action < self.n_actions:
        raise ValueError(f"action {action} outside 0..{self.n_actions - 1}")
    played = self.distribution()
    estimate = np.zeros(self.n_actions)
    estimate[action] = float(reward) / max(float(played[action]), 1e-12)
    self._absorb(estimate, played)
    self.rounds += 1

causalrl.magames.equilibrium.certify_equilibrium(game, profile, *, do=None, topology=LearnerTopology.SINGLE_LEARNER, tol=1e-09, require_kind=None)

Certify profile is a robust (tol-)Nash equilibrium of game, optionally under do.

do fixes the given agents' actions; the remaining free agents must each best-respond within tol for the profile to be certified. kind is capped by topology (I2). If require_kind is stronger than the topology licenses (e.g. IDENTIFIED under INDEPENDENT_LEARNERS), raise :class:KindNotLicensedError — the requested guarantee is not available, so no certificate is returned (acceptance d).

Source code in src/causalrl/magames/equilibrium.py

def certify_equilibrium(
    game: CausalGame,
    profile: Mapping[str, int],
    *,
    do: Mapping[str, int] | None = None,
    topology: LearnerTopology = LearnerTopology.SINGLE_LEARNER,
    tol: float = 1e-9,
    require_kind: Kind | None = None,
) -> Certificate:
    """Certify ``profile`` is a robust (tol-)Nash equilibrium of ``game``, optionally under ``do``.

    ``do`` fixes the given agents' actions; the remaining *free* agents must each best-respond
    within ``tol`` for the profile to be certified. ``kind`` is capped by ``topology`` (I2). If
    ``require_kind`` is stronger than the topology licenses (e.g. ``IDENTIFIED`` under
    ``INDEPENDENT_LEARNERS``), raise :class:`KindNotLicensedError` — the requested guarantee is not
    available, so no certificate is returned (acceptance d).
    """
    kind = topology_max_kind(topology)
    if require_kind is Kind.IDENTIFIED and kind is not Kind.IDENTIFIED:
        raise KindNotLicensedError(
            f"topology {topology.value!r} licenses at most {kind.name}, not IDENTIFIED; "
            "single-agent identification does not hold for simultaneous/centralised learners"
        )

    do = dict(do or {})
    profile_eff = {a: do.get(a, profile[a]) for a in game.agents}
    free_agents = [a for a in game.agents if a not in do]
    deviator, regret = _max_regret(game, profile_eff, free_agents)
    certified = regret <= tol

    assumptions = (
        Assumption(name="finite-game", params={"agents": list(game.agents)}, checkable=True),
        Assumption(name="tol", params={"tol": tol}, checkable=True),
    )
    provenance = Provenance.create(graph_hash=graph_hash(game.graph))
    intervened = f" | do={dict(do)}" if do else ""
    if not certified:
        return Certificate(
            claim=f"profile {dict(profile)} is NOT a robust equilibrium{intervened}",
            estimand=EstimandSpec(query="equilibrium", target="mean"),
            kind=kind,
            value=None,
            alpha=None,
            assumptions=assumptions,
            method="best-response-exact",
            witness=None,
            hedge=Hedge(
                reason="not-an-equilibrium",
                detail={"deviating_agent": deviator, "regret": regret, "do": dict(do)},
            ),
            provenance=provenance,
        )
    return Certificate(
        claim=f"profile {dict(profile)} is a robust equilibrium{intervened}",
        estimand=EstimandSpec(query="equilibrium", target="mean"),
        kind=kind,
        value=None,
        alpha=None,
        assumptions=assumptions,
        method="best-response-exact",
        witness=Witness(
            kind="nash-equilibrium",
            detail={"profile": dict(profile), "do": dict(do), "max_regret": regret},
        ),
        hedge=None,
        provenance=provenance,
    )

causalrl.magames.views.linear_gaussian_population_env(ego='ego', coplayer='co', **kwargs)

Construct a :class:LinearGaussianPopulationEnv (see it for the DGP coefficients).

Source code in src/causalrl/magames/views.py

def linear_gaussian_population_env(
    ego: str = "ego", coplayer: str = "co", **kwargs: float
) -> LinearGaussianPopulationEnv:
    """Construct a :class:`LinearGaussianPopulationEnv` (see it for the DGP coefficients)."""
    return LinearGaussianPopulationEnv(ego=ego, coplayer=coplayer, **kwargs)

causalrl.magames.views.LinearGaussianPopulationEnv dataclass

The environment one ego faces inside a fixed population: a synthetic DGP, not a learner.

An observed context Z confounds the ego's logging action and drives a co-player's action; the reward is linear in the ego action, the co-player action and Z, with Gaussian noise. The ego's causal action effect E[Y | do(ego=1)] - E[Y | do(ego=0)] = ego_effect (the co-player and context are unaffected by the ego's action) is back-door identified by adjusting for Z. Nothing here updates from experience: the fields below are fixed structural coefficients.

Source code in src/causalrl/magames/views.py

@dataclass(frozen=True)
class LinearGaussianPopulationEnv:
    """The environment one ego faces inside a fixed population: a synthetic DGP, not a learner.

    An observed context ``Z`` confounds the ego's logging action and drives a co-player's action;
    the reward is linear in the ego action, the co-player action and ``Z``, with Gaussian noise. The
    ego's causal action effect ``E[Y | do(ego=1)] - E[Y | do(ego=0)] = ego_effect`` (the co-player
    and context are unaffected by the ego's action) is back-door identified by adjusting for ``Z``.
    Nothing here updates from experience: the fields below are fixed structural coefficients.
    """

    ego: str = "ego"
    coplayer: str = "co"
    ego_effect: float = 1.5
    coplayer_effect: float = 0.8
    context_effect: float = 1.0
    confound: float = 1.0  # Z -> ego logging-action confounding strength
    coplayer_bias: float = 0.7  # Z -> co-player action
    noise: float = 0.5

    def _draw(self, n: int, seed: int | None, ego_action: float | None) -> dict[str, FloatArray]:
        rng = np.random.default_rng(seed)
        z = rng.standard_normal(n)
        if ego_action is None:
            a = rng.binomial(1, _sigmoid(self.confound * z)).astype(np.float64)
        else:
            a = np.full(n, float(ego_action))
        b = rng.binomial(1, _sigmoid(self.coplayer_bias * z)).astype(np.float64)
        y = (
            self.ego_effect * a
            + self.coplayer_effect * b
            + self.context_effect * z
            + self.noise * rng.standard_normal(n)
        )
        return {"Z": z, self.ego: a, self.coplayer: b, "Y": y}

    def _log(self, cols: dict[str, FloatArray], regime_label: str) -> TrajectoryLog:
        kinds = {"Z": "obs", self.ego: "action", self.coplayer: "action", "Y": "reward"}
        rows: list[dict[str, Any]] = []
        for name, vals in cols.items():
            for i, v in enumerate(vals.tolist()):
                rows.append(
                    {
                        "entity_id": i,
                        "episode_id": 0,
                        "t": 0,
                        "kind": kinds[name],
                        "name": name,
                        "value": float(v),
                        "regime": regime_label,
                        "observed": True,
                    }
                )
        return TrajectoryLog.from_rows(rows)

    def sample(
        self, n: int, *, seed: int | None = None, regime: Regime | None = None
    ) -> TrajectoryLog:
        return self._log(self._draw(n, seed, None), "observed")

    def do(
        self,
        interventions: Mapping[str, Any],
        n: int,
        *,
        seed: int | None = None,
        regime: Regime | None = None,
    ) -> TrajectoryLog:
        ego_action = interventions.get(self.ego)
        return self._log(
            self._draw(n, seed, None if ego_action is None else float(ego_action)), "do"
        )

    def noise_ledger(self) -> NoiseLedger | None:
        return None

causalrl.agents.no_regret.RegretMatching

Bases: NoRegretLearner

Hart & Mas-Colell (2000) regret matching on external regrets: parameter-free, no-regret.

The cumulative regret of action k is the total extra payoff it would have earned had it replaced the strategy actually played; the next round is drawn proportionally to the positive part of that vector (uniformly while no regret is positive).

Source code in src/causalrl/agents/no_regret.py

class RegretMatching(NoRegretLearner):
    """Hart & Mas-Colell (2000) regret matching on external regrets: parameter-free, no-regret.

    The cumulative regret of action ``k`` is the total extra payoff it would have earned had it
    replaced the strategy actually played; the next round is drawn proportionally to the positive
    part of that vector (uniformly while no regret is positive).
    """

    def __init__(self, n_actions: int, *, explore: float = 0.0, seed: int | None = None) -> None:
        super().__init__(n_actions, explore=explore, seed=seed)
        self.regret: FloatArray = np.zeros(n_actions)

    def _strategy(self) -> FloatArray:
        positive = np.maximum(self.regret, 0.0)
        total = float(positive.sum())
        if total <= 0.0:
            return np.full(self.n_actions, 1.0 / self.n_actions)
        return positive / total

    def _absorb(self, payoffs: FloatArray, strategy: FloatArray) -> None:
        self.regret += payoffs - float(strategy @ payoffs)

causalrl.agents.no_regret.MultiplicativeWeights

Bases: NoRegretLearner

Hedge / multiplicative weights (Freund & Schapire 1997): p[k] ~ exp(eta * cumulative[k]).

learning_rate is eta. Left at None it is the theory rate sqrt(8 log K / T) / payoff_range: with horizon=T given, the fixed rate tuned to that horizon; without one, the anytime rate that uses the number of rounds seen so far instead (both from N. Cesa-Bianchi, G. Lugosi, Prediction, Learning, and Games, CUP 2006, §2.3). payoff_range is the width of the interval the payoffs live in, which is what the rate must be scaled by; pass the game's own range when it is not 1.

Source code in src/causalrl/agents/no_regret.py

class MultiplicativeWeights(NoRegretLearner):
    """Hedge / multiplicative weights (Freund & Schapire 1997): ``p[k] ~ exp(eta * cumulative[k])``.

    ``learning_rate`` is ``eta``. Left at ``None`` it is the theory rate
    ``sqrt(8 log K / T) / payoff_range``: with ``horizon=T`` given, the fixed rate tuned to that
    horizon; without one, the anytime rate that uses the number of rounds seen so far instead
    (both from N. Cesa-Bianchi, G. Lugosi, *Prediction, Learning, and Games*, CUP 2006, §2.3).
    ``payoff_range`` is the width of the interval the payoffs live in, which is what the rate must
    be scaled by; pass the game's own range when it is not ``1``.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        learning_rate: float | None = None,
        horizon: int | None = None,
        payoff_range: float = 1.0,
        explore: float = 0.0,
        seed: int | None = None,
    ) -> None:
        super().__init__(n_actions, explore=explore, seed=seed)
        if learning_rate is not None and learning_rate <= 0.0:
            raise ValueError("learning_rate must be positive")
        if horizon is not None and horizon < 1:
            raise ValueError("horizon must be at least 1")
        if payoff_range <= 0.0:
            raise ValueError("payoff_range must be positive")
        self.learning_rate = learning_rate
        self.horizon = horizon
        self.payoff_range = float(payoff_range)
        self.cumulative: FloatArray = np.zeros(n_actions)

    def _rate(self) -> float:
        if self.learning_rate is not None:
            return self.learning_rate
        elapsed = self.horizon if self.horizon is not None else max(1, self.rounds)
        log_k = np.log(max(2, self.n_actions))
        return float(np.sqrt(8.0 * log_k / elapsed) / self.payoff_range)

    def _strategy(self) -> FloatArray:
        logits = self._rate() * self.cumulative
        logits -= float(logits.max())
        weights = np.exp(logits)
        return weights / float(weights.sum())

    def _absorb(self, payoffs: FloatArray, strategy: FloatArray) -> None:
        self.cumulative += payoffs

Scale & Streaming (Phase 3)

causalrl.ope.ipw.stream_policy_value(source, *, weight='weight', reward='reward', batch_size=100000, alpha=0.05, min_ess_fraction=0.1, policy=None)

Certify a policy's self-normalised importance-sampling value, streamed over a columnar log.

Each decision contributes an importance weight (the weight value cell, rho = π_target/π_behaviour) and a reward cell; the estimate is the Hájek ratio V = Σ(rho·r) / Σ rho with an influence-function CI. Returns a kind=IDENTIFIED :class:Certificate (value + ci) under logged-propensity positivity, else a hedge when the Kish effective sample size falls below min_ess_fraction of the decision count — the streaming positivity guard (I3). source is a :class:TrajectoryLog or a Parquet path, streamed in row batches of batch_size without materialising it.

Source code in src/causalrl/ope/ipw.py

def stream_policy_value(
    source: LogSource,
    *,
    weight: str = "weight",
    reward: str = "reward",
    batch_size: int = 100_000,
    alpha: float = 0.05,
    min_ess_fraction: float = 0.1,
    policy: str | None = None,
) -> Certificate:
    """Certify a policy's self-normalised importance-sampling value, streamed over a columnar log.

    Each decision contributes an importance weight (the ``weight`` value cell,
    ``rho = π_target/π_behaviour``) and a ``reward`` cell; the estimate is the Hájek ratio
    ``V = Σ(rho·r) / Σ rho`` with an influence-function CI. Returns a ``kind=IDENTIFIED``
    :class:`Certificate` (value + ``ci``) under logged-propensity positivity, else a hedge
    when the Kish effective sample size falls below ``min_ess_fraction`` of the decision count — the
    streaming positivity guard (I3). ``source`` is a :class:`TrajectoryLog` or a Parquet path,
    streamed in row batches of ``batch_size`` without materialising it.
    """
    joiner = KeyJoiner((weight, reward))
    ratio = WeightedStreamingRatio()
    n_batches = 0
    for log in iter_log_batches(source, batch_size):
        n_batches += 1
        cols = joiner.drain(log)
        w = cols[weight]
        if w.shape[0]:
            ratio.update(w, cols[reward])

    n = ratio.count
    fingerprint = _stream_fingerprint(float(n), ratio.sum_weights, ratio.value if n else 0.0)
    if n == 0:
        return _hedged_policy_value(
            alpha, policy, Hedge("no-decisions", {"weight": weight, "reward": reward}), fingerprint
        )
    ess = ratio.effective_sample_size
    ess_fraction = ess / n
    if ess_fraction < min_ess_fraction:
        return _hedged_policy_value(
            alpha,
            policy,
            Hedge(
                "overlap-violation",
                {
                    "ess": ess,
                    "n": n,
                    "ess_fraction": ess_fraction,
                    "min_ess_fraction": min_ess_fraction,
                },
            ),
            fingerprint,
        )

    value = ratio.value
    z = float(norm_ppf(1.0 - alpha / 2.0))
    return Certificate(
        claim=f"IS off-policy value V(π) = {value:.4g}",
        estimand=EstimandSpec(query="policy_value", target="mean", policy=policy),
        kind=Kind.IDENTIFIED,
        value=value,
        alpha=alpha,
        assumptions=(
            Assumption(
                name="overlap",
                params={"min_ess_fraction": min_ess_fraction},
                checkable=True,
                diagnostic={"ess": ess, "ess_fraction": ess_fraction, "n": n},
            ),
            Assumption(name="logged-propensities", params={}, checkable=False),
        ),
        method=f"self-normalised IPW (streaming, {n_batches} batch(es), dropped={joiner.dropped})",
        witness=Witness(kind="importance-weighting", detail={"weight": weight, "reward": reward}),
        hedge=None,
        provenance=Provenance.create(data_fingerprint=fingerprint),
        ci=ratio.ci(z),
    )

causalrl.bounds.streaming.stream_msm_bounds(source, *, gamma, outcome='reward', propensity='propensity', treatment=None, batch_size=100000)

Certify MSM bounds on E[Y(1)] from a streamed log under confounding budget gamma.

Each decision supplies an outcome cell and a nominal propensity cell; when treatment is given, only decisions with that indicator cell > 0.5 (the treated units) enter the bound. Returns a kind=BOUNDED :class:Certificate whose value is the Tan interval: it holds E[Y(1)] whenever the true confounding odds ratio is at most gamma, collapses to the IPW point at gamma = 1, and widens monotonically. source streams in row batches of batch_size without materialising the log (only the two needed columns accumulate).

Source code in src/causalrl/bounds/streaming.py

def stream_msm_bounds(
    source: LogSource,
    *,
    gamma: float,
    outcome: str = "reward",
    propensity: str = "propensity",
    treatment: str | None = None,
    batch_size: int = 100_000,
) -> Certificate:
    """Certify MSM bounds on ``E[Y(1)]`` from a streamed log under confounding budget ``gamma``.

    Each decision supplies an ``outcome`` cell and a nominal ``propensity`` cell; when ``treatment``
    is given, only decisions with that indicator cell ``> 0.5`` (the treated units) enter the bound.
    Returns a ``kind=BOUNDED`` :class:`Certificate` whose ``value`` is the Tan interval: it holds
    ``E[Y(1)]`` whenever the true confounding odds ratio is at most ``gamma``, collapses to the IPW
    point at ``gamma = 1``, and widens monotonically. ``source`` streams in row batches of
    ``batch_size`` without materialising the log (only the two needed columns accumulate).
    """
    if gamma < 1.0:
        raise ValueError("gamma must be >= 1")
    names = (outcome, propensity) if treatment is None else (outcome, propensity, treatment)
    joiner = KeyJoiner(names)
    y_chunks: list[FloatArray] = []
    e_chunks: list[FloatArray] = []
    n_batches = 0
    for log in iter_log_batches(source, batch_size):
        n_batches += 1
        cols = joiner.drain(log)
        y = cols[outcome]
        e = cols[propensity]
        if treatment is not None:
            mask = cols[treatment] > 0.5
            y = y[mask]
            e = e[mask]
        if y.shape[0]:
            y_chunks.append(y)
            e_chunks.append(e)

    y = np.concatenate(y_chunks) if y_chunks else np.empty(0, dtype=np.float64)
    e = np.concatenate(e_chunks) if e_chunks else np.empty(0, dtype=np.float64)
    n = int(y.shape[0])
    fingerprint = _fingerprint(float(n), gamma)
    if n == 0:
        return Certificate(
            claim="MSM bounds refused: no treated units in stream",
            estimand=EstimandSpec(query="do", target="mean"),
            kind=Kind.BOUNDED,
            value=None,
            alpha=None,
            assumptions=(),
            method="refused",
            witness=None,
            hedge=Hedge("no-treated-units", {"outcome": outcome, "propensity": propensity}),
            provenance=Provenance.create(data_fingerprint=fingerprint),
            ci=None,
        )

    interval = ipw_sensitivity_bounds(y.tolist(), e.tolist(), gamma=gamma, return_certificate=False)
    return Certificate(
        claim=f"E[Y(1)] ∈ [{interval.lower:.4g}, {interval.upper:.4g}] under Γ={gamma:g}",
        estimand=EstimandSpec(query="do", target="mean"),
        kind=Kind.BOUNDED,
        value=interval,
        alpha=None,
        assumptions=(
            Assumption(name="MSM", params={"gamma": gamma}, checkable=False),
            Assumption(name="logged-propensities", params={}, checkable=False),
        ),
        method=(
            f"Tan MSM closed form (streamed, {n_batches} batches, n={n}, dropped={joiner.dropped})"
        ),
        witness=None,
        hedge=None,
        provenance=Provenance.create(data_fingerprint=fingerprint),
        ci=None,
    )

causalrl.backends.streaming.StreamingMoments

Running count / mean / (sample) variance over streamed batches.

Combines batches with Chan et al.'s parallel merge — carrying count, mean and M2 (the running sum of squared deviations) — which is numerically stable and order-independent, so update-ing batches or merge-ing partial accumulators yields the same moments (up to floating point) as a single pass over the concatenated data.

Source code in src/causalrl/backends/streaming.py

class StreamingMoments:
    """Running count / mean / (sample) variance over streamed batches.

    Combines batches with Chan et al.'s parallel merge — carrying ``count``, ``mean`` and ``M2``
    (the running sum of squared deviations) — which is numerically stable and order-independent, so
    ``update``-ing batches or ``merge``-ing partial accumulators yields the same moments (up to
    floating point) as a single pass over the concatenated data.
    """

    def __init__(self) -> None:
        self._n: int = 0
        self._mean: float = 0.0
        self._m2: float = 0.0  # running sum of squared deviations from the mean

    @property
    def count(self) -> int:
        return self._n

    @property
    def mean(self) -> float:
        return self._mean if self._n > 0 else float("nan")

    @property
    def variance(self) -> float:
        """Sample variance (ddof=1); ``nan`` for fewer than two observations."""
        return self._m2 / (self._n - 1) if self._n > 1 else float("nan")

    @property
    def std(self) -> float:
        v = self.variance
        return float("nan") if math.isnan(v) else math.sqrt(v)

    @property
    def se(self) -> float:
        """Standard error of the mean; ``nan`` for fewer than two observations."""
        return self.std / math.sqrt(self._n) if self._n > 1 else float("nan")

    def update(self, values: FloatArray | Sequence[float]) -> StreamingMoments:
        """Absorb a batch of observations; returns ``self`` for chaining."""
        x = np.asarray(values, dtype=np.float64).reshape(-1)
        n_b = int(x.shape[0])
        if n_b == 0:
            return self
        mean_b = float(x.mean())
        m2_b = float(((x - mean_b) ** 2).sum())
        self._absorb(n_b, mean_b, m2_b)
        return self

    def merge(self, other: StreamingMoments) -> StreamingMoments:
        """Fold another accumulator into this one (associative); returns ``self``."""
        if other._n > 0:
            self._absorb(other._n, other._mean, other._m2)
        return self

    def _absorb(self, n_b: int, mean_b: float, m2_b: float) -> None:
        n_a, mean_a, m2_a = self._n, self._mean, self._m2
        n = n_a + n_b
        delta = mean_b - mean_a
        self._n = n
        self._mean = mean_a + delta * n_b / n
        self._m2 = m2_a + m2_b + delta * delta * n_a * n_b / n

    def ci(self, z: float) -> Interval:
        """Wald confidence interval ``mean ± z·se`` (the caller supplies the ``z`` multiplier)."""
        h = z * self.se
        return Interval(self._mean - h, self._mean + h)

variance property

Sample variance (ddof=1); nan for fewer than two observations.

se property

Standard error of the mean; nan for fewer than two observations.

update(values)

Absorb a batch of observations; returns self for chaining.

Source code in src/causalrl/backends/streaming.py

def update(self, values: FloatArray | Sequence[float]) -> StreamingMoments:
    """Absorb a batch of observations; returns ``self`` for chaining."""
    x = np.asarray(values, dtype=np.float64).reshape(-1)
    n_b = int(x.shape[0])
    if n_b == 0:
        return self
    mean_b = float(x.mean())
    m2_b = float(((x - mean_b) ** 2).sum())
    self._absorb(n_b, mean_b, m2_b)
    return self

merge(other)

Fold another accumulator into this one (associative); returns self.

Source code in src/causalrl/backends/streaming.py

def merge(self, other: StreamingMoments) -> StreamingMoments:
    """Fold another accumulator into this one (associative); returns ``self``."""
    if other._n > 0:
        self._absorb(other._n, other._mean, other._m2)
    return self

ci(z)

Wald confidence interval mean ± z·se (the caller supplies the z multiplier).

Source code in src/causalrl/backends/streaming.py

def ci(self, z: float) -> Interval:
    """Wald confidence interval ``mean ± z·se`` (the caller supplies the ``z`` multiplier)."""
    h = z * self.se
    return Interval(self._mean - h, self._mean + h)

causalrl.backends.streaming.WeightedStreamingRatio

Streaming self-normalised (Hájek) weighted mean Σ w y / Σ w with a one-pass SE.

Accumulates the six sufficient statistics n, Σw, Σwy, Σw², Σw²y, Σw²y², from which the self-normalised estimate V = Σwy/Σw and its standard error follow in a single pass: the Hájek influence contribution is ψ_i = w_i (y_i - V) / w̄ (mean zero by construction), so SE = √(Σ w_i² (y_i - V)²) / Σ w — expandable from the stored moments without a second pass.

With importance weights w_i = π_target(a_i|s_i)/π_behavior(a_i|s_i) and y_i the reward, V is the self-normalised IS off-policy value; :meth:effective_sample_size is the standard (Σw)²/Σw² overlap diagnostic used to hedge when positivity is too weak (invariant I3).

Source code in src/causalrl/backends/streaming.py

class WeightedStreamingRatio:
    """Streaming self-normalised (Hájek) weighted mean ``Σ w y / Σ w`` with a one-pass SE.

    Accumulates the six sufficient statistics ``n, Σw, Σwy, Σw², Σw²y, Σw²y²``, from which the
    self-normalised estimate ``V = Σwy/Σw`` and its standard error follow in a single pass: the
    Hájek influence contribution is ``ψ_i = w_i (y_i - V) / w̄`` (mean zero by construction), so
    ``SE = √(Σ w_i² (y_i - V)²) / Σ w`` — expandable from the stored moments without a second pass.

    With importance weights ``w_i = π_target(a_i|s_i)/π_behavior(a_i|s_i)`` and ``y_i`` the reward,
    ``V`` is the self-normalised IS off-policy value; :meth:`effective_sample_size` is the standard
    ``(Σw)²/Σw²`` overlap diagnostic used to hedge when positivity is too weak (invariant I3).
    """

    def __init__(self) -> None:
        self._n: int = 0
        self._sw: float = 0.0
        self._swy: float = 0.0
        self._sw2: float = 0.0
        self._sw2y: float = 0.0
        self._sw2y2: float = 0.0

    @property
    def count(self) -> int:
        return self._n

    @property
    def sum_weights(self) -> float:
        return self._sw

    @property
    def value(self) -> float:
        return self._swy / self._sw if self._sw != 0.0 else float("nan")

    @property
    def effective_sample_size(self) -> float:
        """Kish effective sample size ``(Σw)² / Σw²`` (0 when no weight has accumulated)."""
        return self._sw * self._sw / self._sw2 if self._sw2 > 0.0 else 0.0

    @property
    def se(self) -> float:
        """Influence-function standard error of :attr:`value`; ``nan`` below two observations."""
        if self._sw == 0.0 or self._n < 2:
            return float("nan")
        v = self.value
        s = self._sw2y2 - 2.0 * v * self._sw2y + v * v * self._sw2
        return math.sqrt(max(s, 0.0)) / self._sw

    def update(
        self, weights: FloatArray | Sequence[float], values: FloatArray | Sequence[float]
    ) -> WeightedStreamingRatio:
        """Absorb a batch of ``(weight, value)`` pairs; returns ``self`` for chaining."""
        w = np.asarray(weights, dtype=np.float64).reshape(-1)
        y = np.asarray(values, dtype=np.float64).reshape(-1)
        if w.shape != y.shape:
            raise ValueError(f"weights {w.shape} and values {y.shape} must have equal length")
        if w.shape[0] == 0:
            return self
        w2 = w * w
        self._n += int(w.shape[0])
        self._sw += float(w.sum())
        self._swy += float((w * y).sum())
        self._sw2 += float(w2.sum())
        self._sw2y += float((w2 * y).sum())
        self._sw2y2 += float((w2 * y * y).sum())
        return self

    def merge(self, other: WeightedStreamingRatio) -> WeightedStreamingRatio:
        """Fold another accumulator into this one (associative); returns ``self``."""
        self._n += other._n
        self._sw += other._sw
        self._swy += other._swy
        self._sw2 += other._sw2
        self._sw2y += other._sw2y
        self._sw2y2 += other._sw2y2
        return self

    def ci(self, z: float) -> Interval:
        """Wald confidence interval ``value ± z·se`` (the caller supplies the ``z`` multiplier)."""
        h = z * self.se
        v = self.value
        return Interval(v - h, v + h)

effective_sample_size property

Kish effective sample size (Σw)² / Σw² (0 when no weight has accumulated).

se property

Influence-function standard error of :attr:value; nan below two observations.

update(weights, values)

Absorb a batch of (weight, value) pairs; returns self for chaining.

Source code in src/causalrl/backends/streaming.py

def update(
    self, weights: FloatArray | Sequence[float], values: FloatArray | Sequence[float]
) -> WeightedStreamingRatio:
    """Absorb a batch of ``(weight, value)`` pairs; returns ``self`` for chaining."""
    w = np.asarray(weights, dtype=np.float64).reshape(-1)
    y = np.asarray(values, dtype=np.float64).reshape(-1)
    if w.shape != y.shape:
        raise ValueError(f"weights {w.shape} and values {y.shape} must have equal length")
    if w.shape[0] == 0:
        return self
    w2 = w * w
    self._n += int(w.shape[0])
    self._sw += float(w.sum())
    self._swy += float((w * y).sum())
    self._sw2 += float(w2.sum())
    self._sw2y += float((w2 * y).sum())
    self._sw2y2 += float((w2 * y * y).sum())
    return self

merge(other)

Fold another accumulator into this one (associative); returns self.

Source code in src/causalrl/backends/streaming.py

def merge(self, other: WeightedStreamingRatio) -> WeightedStreamingRatio:
    """Fold another accumulator into this one (associative); returns ``self``."""
    self._n += other._n
    self._sw += other._sw
    self._swy += other._swy
    self._sw2 += other._sw2
    self._sw2y += other._sw2y
    self._sw2y2 += other._sw2y2
    return self

ci(z)

Wald confidence interval value ± z·se (the caller supplies the z multiplier).

Source code in src/causalrl/backends/streaming.py

def ci(self, z: float) -> Interval:
    """Wald confidence interval ``value ± z·se`` (the caller supplies the ``z`` multiplier)."""
    h = z * self.se
    v = self.value
    return Interval(v - h, v + h)

Interop & Scale (Phase 4)

causalrl.interop.sbi_numpyro.regimes_from_posterior(samples, *, name_prefix='posterior', selection=(), max_regimes=None, seed=None)

Turn posterior samples over environment parameters into one :class:Regime per sample.

samples is either a mapping param_name -> samples or a (matrix, names) pair with matrix of shape (n_samples, n_params). Each regime carries that sample's values and marks selection as the mechanism variables that differ across the population/config. max_regimes subsamples without replacement (deterministic given seed) to keep a large posterior tractable.

Source code in src/causalrl/interop/sbi_numpyro.py

def regimes_from_posterior(
    samples: PosteriorSamples,
    *,
    name_prefix: str = "posterior",
    selection: Iterable[str] = (),
    max_regimes: int | None = None,
    seed: int | None = None,
) -> list[Regime]:
    """Turn posterior samples over environment parameters into one :class:`Regime` per sample.

    ``samples`` is either a mapping ``param_name -> samples`` or a ``(matrix, names)`` pair with
    ``matrix`` of shape ``(n_samples, n_params)``. Each regime carries that sample's values
    and marks ``selection`` as the mechanism variables that differ across the population/config.
    ``max_regimes`` subsamples without replacement (deterministic given ``seed``) to keep a large
    posterior tractable.
    """
    names, matrix = _as_named_matrix(samples)
    n = int(matrix.shape[0])
    idx = np.arange(n)
    if max_regimes is not None and max_regimes < n:
        rng = np.random.default_rng(seed)
        idx = np.sort(rng.choice(n, size=max_regimes, replace=False))
    sel = frozenset(selection)
    regimes: list[Regime] = []
    for j in idx.tolist():
        params = {names[k]: float(matrix[j, k]) for k in range(len(names))}
        regimes.append(Regime.create(f"{name_prefix}[{j}]", selection=sel, parameters=params))
    return regimes

causalrl.interop.sbi_numpyro.across_regimes(regimes, fn)

Quantify a functional across calibrated regimes as its [min, max] :class:Interval.

fn maps each regime (its parameter configuration) to a scalar — e.g. a transported effect or a partial-identification bound evaluated under that configuration. The returned interval is the range over the posterior ensemble: a worst-case envelope across calibrated configurations.

Source code in src/causalrl/interop/sbi_numpyro.py

def across_regimes(regimes: Iterable[Regime], fn: Callable[[Regime], float]) -> Interval:
    """Quantify a functional across calibrated regimes as its ``[min, max]`` :class:`Interval`.

    ``fn`` maps each regime (its parameter configuration) to a scalar — e.g. a transported effect or
    a partial-identification bound evaluated under that configuration. The returned interval is the
    range over the posterior ensemble: a worst-case envelope across calibrated configurations.
    """
    values = [float(fn(r)) for r in regimes]
    if not values:
        raise ValueError("no regimes to evaluate")
    return Interval(min(values), max(values))

causalrl.interop.columnar_sim.ColumnarSimulator

Reference :class:~causalrl.protocols.CausalEnvProtocol adapter over row-emitting callables.

sample_fn(n, seed) returns an iterable of row mappings (the TrajectoryLog schema); the adapter wraps them into a :class:TrajectoryLog. do_fn(interventions, n, seed) is optional: without it, :meth:do raises :class:CausalInterfaceUnavailableError (honest: the simulator exposes only observational sampling). ledger optionally supplies a white-box :class:NoiseLedger for counterfactuals.

Source code in src/causalrl/interop/columnar_sim.py

class ColumnarSimulator:
    """Reference :class:`~causalrl.protocols.CausalEnvProtocol` adapter over row-emitting callables.

    ``sample_fn(n, seed)`` returns an iterable of row mappings (the ``TrajectoryLog`` schema); the
    adapter wraps them into a :class:`TrajectoryLog`. ``do_fn(interventions, n, seed)`` is optional:
    without it, :meth:`do` raises :class:`CausalInterfaceUnavailableError` (honest: the simulator
    exposes only observational sampling). ``ledger`` optionally supplies a white-box
    :class:`NoiseLedger` for counterfactuals.
    """

    def __init__(
        self,
        sample_fn: SampleFn,
        do_fn: DoFn | None = None,
        *,
        metadata: Mapping[str, Any] | None = None,
        ledger: NoiseLedger | None = None,
    ) -> None:
        self._sample_fn = sample_fn
        self._do_fn = do_fn
        self._metadata: dict[str, Any] = dict(metadata or {})
        self._ledger = ledger

    def sample(
        self, n: int, *, seed: int | None = None, regime: Regime | None = None
    ) -> TrajectoryLog:
        return TrajectoryLog.from_rows(list(self._sample_fn(n, seed)), self._metadata)

    def do(
        self,
        interventions: Mapping[str, Any],
        n: int,
        *,
        seed: int | None = None,
        regime: Regime | None = None,
    ) -> TrajectoryLog:
        if self._do_fn is None:
            raise CausalInterfaceUnavailableError(
                "this simulator implements only observational sample(); provide a do_fn for do()"
            )
        return TrajectoryLog.from_rows(list(self._do_fn(interventions, n, seed)), self._metadata)

    def noise_ledger(self) -> NoiseLedger | None:
        return self._ledger

causalrl.interop.columnar_sim.simulator_from_callables(sample_fn, do_fn=None, *, metadata=None, ledger=None)

Build a :class:ColumnarSimulator from a sample callable (and optional do/ledger).

Source code in src/causalrl/interop/columnar_sim.py

def simulator_from_callables(
    sample_fn: SampleFn,
    do_fn: DoFn | None = None,
    *,
    metadata: Mapping[str, Any] | None = None,
    ledger: NoiseLedger | None = None,
) -> ColumnarSimulator:
    """Build a :class:`ColumnarSimulator` from a sample callable (and optional do/ledger)."""
    return ColumnarSimulator(sample_fn, do_fn, metadata=metadata, ledger=ledger)

causalrl.ope.certify.certify_policy(dataset, target_actions, *, gamma_max=10.0, alpha=None)

Certify whether a learned policy's value improvement over the behaviour policy is robust to hidden confounding.

target_actions[i] is the action the learned policy would take at the i-th logged decision's state (e.g. from a d3rlpy policy's greedy prediction, using the same observation encoding it was trained on). dataset is any :class:~causalrl.data.logged.LoggedDecisions: the tabular :class:~causalrl.ConfoundedTrajectoryDataset whose actions are arm indices, or a :class:~causalrl.data.logged.FeatureDecisionLog whose states are feature vectors and whose actions are :data:~causalrl.Intervention assignments -- so an :class:~causalrl.agents.interventional.InterventionalAgent can be certified directly, with no arm codebook in between. The action type is inferred from the log, so a mismatched pair is a type error rather than a silent miscount.

The contrast is the off-policy value V(pi) - V(behaviour) under Tan's marginal sensitivity model: the nominal logging propensities are the dataset's empirical behavior_propensity, pi's per-unit target probability is 1 where it matches the logged action (greedy), and the behaviour arm's target probability is the logging propensity, so its self-normalised value is the logged empirical return. Returns the same :class:~causalrl.DecisionCertificate as :func:causalrl.certify_decision.

Passing alpha additionally runs the finite-sample lower-confidence-bound gate for safe policy improvement: :func:causalrl.conformal_action_value calibrates a distribution-free lower bound on a fresh return under pi and under the logging policy from the same logs, and certified then requires the confounding verdict and lcb(pi) >= lcb(behaviour). Too few effectively-weighted samples for the level leave lcb(pi) at -inf, which refuses rather than passes — no evidence is not evidence of safety. The gate is on the return of a single decision, not on V(pi): a policy with the better mean but a worse downside fails it.

Honest scope: the MSM sensitivity is on the logging propensities; the behaviour value is the logged empirical return (on-policy for the behaviour policy, hence unconfounded). This is the one-step / terminal-return contrast; the per-step cumulative-reward extension uses :func:causalrl.msm_per_step_bounds.

Source code in src/causalrl/ope/certify.py

def certify_policy(
    dataset: LoggedDecisions[ActionT],
    target_actions: Sequence[ActionT],
    *,
    gamma_max: float = 10.0,
    alpha: float | None = None,
) -> DecisionCertificate:
    """Certify whether a learned policy's value improvement over the behaviour policy is robust
    to hidden confounding.

    ``target_actions[i]`` is the action the learned policy would take at the ``i``-th logged
    decision's state (e.g. from a d3rlpy policy's greedy prediction, using the same observation
    encoding it was trained on). ``dataset`` is any
    :class:`~causalrl.data.logged.LoggedDecisions`: the tabular
    :class:`~causalrl.ConfoundedTrajectoryDataset` whose actions are arm indices, or a
    :class:`~causalrl.data.logged.FeatureDecisionLog` whose states are feature vectors and whose
    actions are :data:`~causalrl.Intervention` assignments -- so an
    :class:`~causalrl.agents.interventional.InterventionalAgent` can be certified directly, with no
    arm codebook in between. The action type is inferred from the log, so a mismatched pair is a
    type error rather than a silent miscount.

    The contrast is the off-policy value ``V(pi) - V(behaviour)`` under
    Tan's marginal sensitivity model: the nominal logging propensities are the dataset's empirical
    ``behavior_propensity``, ``pi``'s per-unit target probability is ``1`` where it matches the
    logged action (greedy), and the behaviour arm's target probability is the logging propensity, so
    its self-normalised value is the logged empirical return. Returns the same
    :class:`~causalrl.DecisionCertificate` as :func:`causalrl.certify_decision`.

    Passing ``alpha`` additionally runs the finite-sample **lower-confidence-bound gate** for safe
    policy improvement: :func:`causalrl.conformal_action_value` calibrates a distribution-free
    lower bound on a fresh return under ``pi`` and under the logging policy from the same logs, and
    ``certified`` then requires the confounding verdict *and* ``lcb(pi) >= lcb(behaviour)``. Too
    few effectively-weighted samples for the level leave ``lcb(pi)`` at ``-inf``, which refuses
    rather than passes — no evidence is not evidence of safety. The gate is on the return of a
    single decision, not on ``V(pi)``: a policy with the better mean but a worse downside fails it.

    Honest scope: the MSM sensitivity is on the logging propensities; the behaviour value is the
    logged empirical return (on-policy for the behaviour policy, hence unconfounded). This is the
    one-step / terminal-return contrast; the per-step cumulative-reward extension uses
    :func:`causalrl.msm_per_step_bounds`.
    """
    outcomes = list(dataset.outcomes())
    e0 = list(dataset.logging_propensities())
    target_on = [1.0 if m else 0.0 for m in dataset.matches(target_actions)]
    contrast = PolicyValueContrast(
        outcomes=outcomes,
        logging_propensities=e0,
        target_on=target_on,
        target_off=list(e0),
    )
    cert = certify_estimate(contrast, gamma_max=gamma_max, labels=("learned policy", "behavior"))
    if alpha is None:
        return cert
    band = conformal_action_value(dataset, target_actions, alpha=alpha).ci
    reference = conformal_action_value(dataset, None, alpha=alpha).ci
    assert band is not None and reference is not None  # conformal_action_value always sets ci
    passed = math.isfinite(band.lower) and band.lower >= reference.lower
    summary = (
        f"{cert.summary} Finite-sample downside gate (conformal, alpha={alpha:g}): "
        f"{'PASS' if passed else 'REFUSE'} — one decision under the learned policy returns at "
        f"least {band.lower:.3f}, vs {reference.lower:.3f} under behavior."
    )
    return cert._replace(
        certified=cert.certified and passed, conformal_lcb=band.lower, summary=summary
    )

causalrl.scale.d3rlpy.certify_fqe(estimated_value, *, policy=None, method='FQE')

Wrap a fitted-Q-evaluation off-policy value estimate in an honest EMPIRICAL certificate.

Fitted-Q evaluation (d3rlpy's FQE) is a model-based off-policy value estimate with no identification guarantee under hidden confounding, so the certificate is kind=EMPIRICAL (I2/I3) — it reports the estimate and its provenance, never a causal point claim. Upgrade to an identified certify_effect (§7.2) or bound it with :func:causalrl.certify_policy when the structure licenses it. estimated_value is the scalar FQE value (e.g. the initial-state value a trained FQE predicts for the evaluated policy).

Source code in src/causalrl/scale/d3rlpy.py

def certify_fqe(
    estimated_value: float, *, policy: str | None = None, method: str = "FQE"
) -> Certificate:
    """Wrap a fitted-Q-evaluation off-policy value estimate in an honest ``EMPIRICAL`` certificate.

    Fitted-Q evaluation (d3rlpy's ``FQE``) is a model-based off-policy value estimate with no
    identification guarantee under hidden confounding, so the certificate is ``kind=EMPIRICAL``
    (I2/I3) — it reports the estimate and its provenance, never a causal point claim. Upgrade to an
    identified ``certify_effect`` (§7.2) or bound it with :func:`causalrl.certify_policy` when the
    structure licenses it. ``estimated_value`` is the scalar FQE value (e.g. the initial-state value
    a trained ``FQE`` predicts for the evaluated policy).
    """
    value = float(estimated_value)
    return Certificate(
        claim=f"FQE off-policy value estimate V(pi) ~= {value:.4g}",
        estimand=EstimandSpec(query="policy_value", target="mean", policy=policy),
        kind=Kind.EMPIRICAL,
        value=value,
        alpha=None,
        assumptions=(Assumption(name="fqe-model", params={"method": method}, checkable=False),),
        method=f"fitted-Q evaluation ({method})",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
        ci=None,
    )

causalrl.scale.d3rlpy.policy_actions(policy, observations)

Greedy actions a trained (duck-typed) policy takes on observations — a do() handle.

Calls the policy's predict (d3rlpy's algorithm interface); returns integer actions ready to pass as target_actions to :func:causalrl.certify_policy, or as a policy intervention. Never imports d3rlpy; any object with a predict(observations) method works.

Source code in src/causalrl/scale/d3rlpy.py

def policy_actions(policy: Any, observations: Any) -> list[int]:
    """Greedy actions a trained (duck-typed) policy takes on ``observations`` — a ``do()`` handle.

    Calls the policy's ``predict`` (d3rlpy's algorithm interface); returns integer actions ready
    to pass as ``target_actions`` to :func:`causalrl.certify_policy`, or as a policy intervention.
    Never imports d3rlpy; any object with a ``predict(observations)`` method works.
    """
    predictions = np.asarray(policy.predict(np.asarray(observations)))
    return [int(a) for a in predictions.reshape(-1).tolist()]

Set-Valued Interventions

causalrl.intervention.InterventionSpace dataclass

Which variables may be intervened on right now, and the values each may take.

Hashable and composable, following :class:causalrl.Regime's idiom: the domains are stored as a normalised tuple of pairs so that two spaces built from equal mappings compare equal regardless of insertion order. Build one with :meth:create.

A space is a context-dependent object. The intended pattern is to construct a fresh one per decision from the current state -- dropping the levers this state does not permit and narrowing the values the others may take -- and hand it to the agent for that decision. What stays fixed across decisions is the graph; what varies is this.

Source code in src/causalrl/intervention.py

@dataclass(frozen=True)
class InterventionSpace:
    """Which variables may be intervened on right now, and the values each may take.

    Hashable and composable, following :class:`causalrl.Regime`'s idiom: the domains are stored as
    a normalised tuple of pairs so that two spaces built from equal mappings compare equal
    regardless of insertion order. Build one with :meth:`create`.

    A space is a *context-dependent* object. The intended pattern is to construct a fresh one per
    decision from the current state -- dropping the levers this state does not permit and
    narrowing the values the others may take -- and hand it to the agent for that decision. What
    stays fixed across decisions is the graph; what varies is this.
    """

    domains: tuple[tuple[str, InterventionDomain], ...] = ()

    def __post_init__(self) -> None:
        # Raw value tuples are normalised to Discrete, so direct construction from the pre-3.1
        # ``(name, values)`` shape keeps working unchanged.
        normalised = tuple(
            sorted(((n, _as_domain(d)) for n, d in self.domains), key=lambda kv: kv[0])
        )
        object.__setattr__(self, "domains", normalised)

    @classmethod
    def create(cls, domains: Mapping[str, Iterable[Any]]) -> InterventionSpace:
        """Build a space from ``{variable: allowed values}``.

        A variable mapped to an empty domain is refused rather than silently retained: it would be
        a variable that is nominally manipulable but has nothing it can be set to, which makes
        every intervention set containing it unsatisfiable while still counting as admissible.
        Drop the variable instead -- that is what "not manipulable here" means.
        """
        pairs: list[tuple[str, InterventionDomain]] = []
        for name in sorted(domains):
            try:
                pairs.append((name, _as_domain(domains[name])))
            except ValueError as exc:
                raise ValueError(f"variable {name!r}: {exc}") from exc
        return cls(tuple(pairs))

    @property
    def variables(self) -> frozenset[str]:
        """The manipulable variables — what :func:`causalrl.pomis` takes as ``manipulable``."""
        return frozenset(name for name, _ in self.domains)

    def domain(self, variable: str) -> InterventionDomain:
        """The :data:`InterventionDomain` of ``variable``.

        Raises ``KeyError`` if it is not manipulable here.
        """
        for name, domain in self.domains:
            if name == variable:
                return domain
        raise KeyError(
            f"{variable!r} is not manipulable in this InterventionSpace "
            f"(manipulable: {sorted(self.variables)})"
        )

    def values(self, variable: str) -> tuple[Any, ...]:
        """The admissible values of a DISCRETE ``variable``.

        Raises ``KeyError`` if ``variable`` is not manipulable, and ``TypeError`` if its domain is
        continuous -- an interval has no value list, and returning a sampled or gridded stand-in
        would quietly turn an exact domain into an approximation of one.
        """
        domain = self.domain(variable)
        if isinstance(domain, Continuous):
            raise TypeError(
                f"{variable!r} has the continuous domain [{domain.low}, {domain.high}], which has "
                "no enumerable value list. Use permits/project/sample to work with it, or pass a "
                "Discrete domain if you want arms."
            )
        return domain.values

    def permits(self, intervention: Intervention) -> bool:
        """Whether every variable in ``intervention`` is manipulable and set to an allowed value."""
        lookup = dict(self.domains)
        return all(
            name in lookup and lookup[name].contains(value) for name, value in intervention.items()
        )

    def project(self, intervention: Intervention) -> Intervention:
        """``intervention`` moved into this space: each value clipped or snapped to its domain.

        The counterpart of :meth:`permits` for a search that proposes freely and needs its
        proposals made admissible -- clipping a dose to its safe range rather than discarding the
        candidate. Raises ``KeyError`` naming any variable that is not manipulable here, since no
        amount of clipping makes an unavailable lever available.
        """
        return {name: self.domain(name).project(value) for name, value in intervention.items()}

    def sample(self, variables: Iterable[str], rng: np.random.Generator) -> Intervention:
        """One admissible assignment of ``variables``, drawn uniformly from each domain.

        The continuous counterpart of :meth:`assignments`: where a finite space can be enumerated,
        an interval must be searched, and a search needs proposals.
        """
        return {name: self.domain(name).sample(rng) for name in sorted(variables)}

    def restrict(self, variables: Iterable[str]) -> InterventionSpace:
        """The sub-space over ``variables``, ignoring any that are not manipulable here."""
        keep = set(variables)
        return InterventionSpace(tuple(kv for kv in self.domains if kv[0] in keep))

    def __and__(self, other: InterventionSpace) -> InterventionSpace:
        """Intersect two spaces: variables in both, each with the values both allow.

        A variable whose domains do not overlap drops out entirely, since retaining it with an
        empty domain is exactly what :meth:`create` refuses.
        """
        theirs = dict(other.domains)
        pairs: list[tuple[str, InterventionDomain]] = []
        for name, mine in self.domains:
            if name not in theirs:
                continue
            merged = _intersect(mine, theirs[name])
            if merged is not None:
                pairs.append((name, merged))
        return InterventionSpace(tuple(pairs))

    def assignments(self, variables: Iterable[str]) -> Iterator[Intervention]:
        """Every assignment of ``variables`` admissible here — the arms of an intervention set.

        This is the step from "which variables to intervene on" (what POMIS answers) to "which
        interventions to choose between" (what an agent needs). The empty set yields exactly one
        assignment, the empty intervention, which is the observational regime rather than nothing.

        Raises ``KeyError`` naming any variable that is not manipulable here, and ``ValueError``
        when the product exceeds :data:`_MAX_ASSIGNMENTS`.
        """
        names = sorted(variables)
        # KeyError names a non-manipulable variable; TypeError names a continuous one, which has
        # no arms to enumerate -- the case that needs a search rather than an arm list.
        domains = [self.values(name) for name in names]
        total = 1
        for values in domains:
            total *= len(values)
        if total > _MAX_ASSIGNMENTS:
            sizes = ", ".join(f"{n}={len(v)}" for n, v in zip(names, domains, strict=True))
            raise ValueError(
                f"enumerating assignments of {names} would produce {total} arms -- the product of "
                f"their domain sizes ({sizes}) -- above the _MAX_ASSIGNMENTS={_MAX_ASSIGNMENTS} "
                "limit. A decision over that many arms needs structure (a factored value model or "
                "a search over assignments), not an enumerated arm list."
            )
        for combination in product(*domains):
            yield dict(zip(names, combination, strict=True))

variables property

The manipulable variables — what :func:causalrl.pomis takes as manipulable.

create(domains) classmethod

Build a space from {variable: allowed values}.

A variable mapped to an empty domain is refused rather than silently retained: it would be a variable that is nominally manipulable but has nothing it can be set to, which makes every intervention set containing it unsatisfiable while still counting as admissible. Drop the variable instead -- that is what "not manipulable here" means.

Source code in src/causalrl/intervention.py

@classmethod
def create(cls, domains: Mapping[str, Iterable[Any]]) -> InterventionSpace:
    """Build a space from ``{variable: allowed values}``.

    A variable mapped to an empty domain is refused rather than silently retained: it would be
    a variable that is nominally manipulable but has nothing it can be set to, which makes
    every intervention set containing it unsatisfiable while still counting as admissible.
    Drop the variable instead -- that is what "not manipulable here" means.
    """
    pairs: list[tuple[str, InterventionDomain]] = []
    for name in sorted(domains):
        try:
            pairs.append((name, _as_domain(domains[name])))
        except ValueError as exc:
            raise ValueError(f"variable {name!r}: {exc}") from exc
    return cls(tuple(pairs))

domain(variable)

The :data:InterventionDomain of variable.

Raises KeyError if it is not manipulable here.

Source code in src/causalrl/intervention.py

def domain(self, variable: str) -> InterventionDomain:
    """The :data:`InterventionDomain` of ``variable``.

    Raises ``KeyError`` if it is not manipulable here.
    """
    for name, domain in self.domains:
        if name == variable:
            return domain
    raise KeyError(
        f"{variable!r} is not manipulable in this InterventionSpace "
        f"(manipulable: {sorted(self.variables)})"
    )

values(variable)

The admissible values of a DISCRETE variable.

Raises KeyError if variable is not manipulable, and TypeError if its domain is continuous -- an interval has no value list, and returning a sampled or gridded stand-in would quietly turn an exact domain into an approximation of one.

Source code in src/causalrl/intervention.py

def values(self, variable: str) -> tuple[Any, ...]:
    """The admissible values of a DISCRETE ``variable``.

    Raises ``KeyError`` if ``variable`` is not manipulable, and ``TypeError`` if its domain is
    continuous -- an interval has no value list, and returning a sampled or gridded stand-in
    would quietly turn an exact domain into an approximation of one.
    """
    domain = self.domain(variable)
    if isinstance(domain, Continuous):
        raise TypeError(
            f"{variable!r} has the continuous domain [{domain.low}, {domain.high}], which has "
            "no enumerable value list. Use permits/project/sample to work with it, or pass a "
            "Discrete domain if you want arms."
        )
    return domain.values

permits(intervention)

Whether every variable in intervention is manipulable and set to an allowed value.

Source code in src/causalrl/intervention.py

def permits(self, intervention: Intervention) -> bool:
    """Whether every variable in ``intervention`` is manipulable and set to an allowed value."""
    lookup = dict(self.domains)
    return all(
        name in lookup and lookup[name].contains(value) for name, value in intervention.items()
    )

project(intervention)

intervention moved into this space: each value clipped or snapped to its domain.

The counterpart of :meth:permits for a search that proposes freely and needs its proposals made admissible -- clipping a dose to its safe range rather than discarding the candidate. Raises KeyError naming any variable that is not manipulable here, since no amount of clipping makes an unavailable lever available.

Source code in src/causalrl/intervention.py

def project(self, intervention: Intervention) -> Intervention:
    """``intervention`` moved into this space: each value clipped or snapped to its domain.

    The counterpart of :meth:`permits` for a search that proposes freely and needs its
    proposals made admissible -- clipping a dose to its safe range rather than discarding the
    candidate. Raises ``KeyError`` naming any variable that is not manipulable here, since no
    amount of clipping makes an unavailable lever available.
    """
    return {name: self.domain(name).project(value) for name, value in intervention.items()}

sample(variables, rng)

One admissible assignment of variables, drawn uniformly from each domain.

The continuous counterpart of :meth:assignments: where a finite space can be enumerated, an interval must be searched, and a search needs proposals.

Source code in src/causalrl/intervention.py

def sample(self, variables: Iterable[str], rng: np.random.Generator) -> Intervention:
    """One admissible assignment of ``variables``, drawn uniformly from each domain.

    The continuous counterpart of :meth:`assignments`: where a finite space can be enumerated,
    an interval must be searched, and a search needs proposals.
    """
    return {name: self.domain(name).sample(rng) for name in sorted(variables)}

restrict(variables)

The sub-space over variables, ignoring any that are not manipulable here.

Source code in src/causalrl/intervention.py

def restrict(self, variables: Iterable[str]) -> InterventionSpace:
    """The sub-space over ``variables``, ignoring any that are not manipulable here."""
    keep = set(variables)
    return InterventionSpace(tuple(kv for kv in self.domains if kv[0] in keep))

__and__(other)

Intersect two spaces: variables in both, each with the values both allow.

A variable whose domains do not overlap drops out entirely, since retaining it with an empty domain is exactly what :meth:create refuses.

Source code in src/causalrl/intervention.py

def __and__(self, other: InterventionSpace) -> InterventionSpace:
    """Intersect two spaces: variables in both, each with the values both allow.

    A variable whose domains do not overlap drops out entirely, since retaining it with an
    empty domain is exactly what :meth:`create` refuses.
    """
    theirs = dict(other.domains)
    pairs: list[tuple[str, InterventionDomain]] = []
    for name, mine in self.domains:
        if name not in theirs:
            continue
        merged = _intersect(mine, theirs[name])
        if merged is not None:
            pairs.append((name, merged))
    return InterventionSpace(tuple(pairs))

assignments(variables)

Every assignment of variables admissible here — the arms of an intervention set.

This is the step from "which variables to intervene on" (what POMIS answers) to "which interventions to choose between" (what an agent needs). The empty set yields exactly one assignment, the empty intervention, which is the observational regime rather than nothing.

Raises KeyError naming any variable that is not manipulable here, and ValueError when the product exceeds :data:_MAX_ASSIGNMENTS.

Source code in src/causalrl/intervention.py

def assignments(self, variables: Iterable[str]) -> Iterator[Intervention]:
    """Every assignment of ``variables`` admissible here — the arms of an intervention set.

    This is the step from "which variables to intervene on" (what POMIS answers) to "which
    interventions to choose between" (what an agent needs). The empty set yields exactly one
    assignment, the empty intervention, which is the observational regime rather than nothing.

    Raises ``KeyError`` naming any variable that is not manipulable here, and ``ValueError``
    when the product exceeds :data:`_MAX_ASSIGNMENTS`.
    """
    names = sorted(variables)
    # KeyError names a non-manipulable variable; TypeError names a continuous one, which has
    # no arms to enumerate -- the case that needs a search rather than an arm list.
    domains = [self.values(name) for name in names]
    total = 1
    for values in domains:
        total *= len(values)
    if total > _MAX_ASSIGNMENTS:
        sizes = ", ".join(f"{n}={len(v)}" for n, v in zip(names, domains, strict=True))
        raise ValueError(
            f"enumerating assignments of {names} would produce {total} arms -- the product of "
            f"their domain sizes ({sizes}) -- above the _MAX_ASSIGNMENTS={_MAX_ASSIGNMENTS} "
            "limit. A decision over that many arms needs structure (a factored value model or "
            "a search over assignments), not an enumerated arm list."
        )
    for combination in product(*domains):
        yield dict(zip(names, combination, strict=True))

causalrl.intervention.canonical(intervention)

A hashable, order-independent key for an intervention.

Intervention is a Mapping and therefore unhashable, which makes the obvious things -- deduplicating arms, using an intervention as a dict key, comparing two of them for identity -- awkward at every call site. This is the one canonical form, sorted by variable name so {"A": 1, "B": 0} and {"B": 0, "A": 1} produce the same key.

Source code in src/causalrl/intervention.py

def canonical(intervention: Intervention) -> tuple[tuple[str, Any], ...]:
    """A hashable, order-independent key for an intervention.

    ``Intervention`` is a ``Mapping`` and therefore unhashable, which makes the obvious things --
    deduplicating arms, using an intervention as a dict key, comparing two of them for identity --
    awkward at every call site. This is the one canonical form, sorted by variable name so
    ``{"A": 1, "B": 0}`` and ``{"B": 0, "A": 1}`` produce the same key.
    """
    return tuple(sorted(intervention.items()))

causalrl.agents.interventional.InterventionalAgent

Bases: ABC

An agent whose action is an intervention — an assignment, not an arm index.

Implementations must return an intervention that space permits. The base class does not enforce that (a subclass may have a cheaper way to stay in bounds than re-checking), but :meth:check_permitted is provided for the ones that want it.

deadline is advisory and cooperative: nothing here interrupts a running computation, so an implementation that ignores it will simply overrun. An agent that does honour it should keep a usable incumbent answer at all times and return it once :meth:~causalrl.deadline.Deadline.expired reports the budget is gone. None means no budget, which is the default so that offline and test call sites stay free of timing concerns.

Source code in src/causalrl/agents/interventional.py

class InterventionalAgent(ABC):
    """An agent whose action is an intervention — an assignment, not an arm index.

    Implementations must return an intervention that ``space`` permits. The base class does not
    enforce that (a subclass may have a cheaper way to stay in bounds than re-checking), but
    :meth:`check_permitted` is provided for the ones that want it.

    ``deadline`` is advisory and cooperative: nothing here interrupts a running computation, so an
    implementation that ignores it will simply overrun. An agent that *does* honour it should keep
    a usable incumbent answer at all times and return it once
    :meth:`~causalrl.deadline.Deadline.expired` reports the budget is gone. ``None`` means no
    budget, which is the default so that offline and test call sites stay free of timing concerns.
    """

    @abstractmethod
    def act(
        self,
        observation: Mapping[str, Any],
        *,
        space: InterventionSpace,
        deadline: Deadline | None = None,
    ) -> Intervention:
        """Choose an intervention admissible in ``space``, ideally within ``deadline``."""

    @abstractmethod
    def update(
        self, observation: Mapping[str, Any], intervention: Intervention, reward: float
    ) -> None:
        """Learn from the reward that followed ``intervention``."""

    @staticmethod
    def check_permitted(intervention: Intervention, space: InterventionSpace) -> Intervention:
        """Return ``intervention`` if ``space`` permits it, else raise ``ValueError`` saying why.

        A helper rather than an enforced invariant: an agent that constructs its candidates from
        ``space`` in the first place cannot violate it and should not pay to re-check.
        """
        if not space.permits(intervention):
            offending = {
                name: value
                for name, value in intervention.items()
                if name not in space.variables or value not in space.values(name)
            }
            raise ValueError(
                f"intervention {dict(intervention)!r} is not admissible in this "
                f"InterventionSpace: {offending!r} names a variable that is not manipulable here, "
                f"or a value outside its domain (manipulable: {sorted(space.variables)})"
            )
        return intervention

act(observation, *, space, deadline=None) abstractmethod

Choose an intervention admissible in space, ideally within deadline.

Source code in src/causalrl/agents/interventional.py

@abstractmethod
def act(
    self,
    observation: Mapping[str, Any],
    *,
    space: InterventionSpace,
    deadline: Deadline | None = None,
) -> Intervention:
    """Choose an intervention admissible in ``space``, ideally within ``deadline``."""

update(observation, intervention, reward) abstractmethod

Learn from the reward that followed intervention.

Source code in src/causalrl/agents/interventional.py

@abstractmethod
def update(
    self, observation: Mapping[str, Any], intervention: Intervention, reward: float
) -> None:
    """Learn from the reward that followed ``intervention``."""

check_permitted(intervention, space) staticmethod

Return intervention if space permits it, else raise ValueError saying why.

A helper rather than an enforced invariant: an agent that constructs its candidates from space in the first place cannot violate it and should not pay to re-check.

Source code in src/causalrl/agents/interventional.py

@staticmethod
def check_permitted(intervention: Intervention, space: InterventionSpace) -> Intervention:
    """Return ``intervention`` if ``space`` permits it, else raise ``ValueError`` saying why.

    A helper rather than an enforced invariant: an agent that constructs its candidates from
    ``space`` in the first place cannot violate it and should not pay to re-check.
    """
    if not space.permits(intervention):
        offending = {
            name: value
            for name, value in intervention.items()
            if name not in space.variables or value not in space.values(name)
        }
        raise ValueError(
            f"intervention {dict(intervention)!r} is not admissible in this "
            f"InterventionSpace: {offending!r} names a variable that is not manipulable here, "
            f"or a value outside its domain (manipulable: {sorted(space.variables)})"
        )
    return intervention

causalrl.agents.interventional.ScalarAgentAdapter

Bases: InterventionalAgent

Present an arm-indexed :class:causalrl.Agent as an :class:InterventionalAgent.

arms[i] is the intervention the wrapped agent's action i denotes, so the adapter is just a codebook in both directions: :meth:act maps the chosen index to its intervention, and :meth:update maps an intervention back to its index before forwarding the reward.

The wrapped agent's arm list is fixed, so space can only be checked, not honoured. A scalar agent has no way to be told that arm 3 is unavailable this turn -- its policy is defined over all its arms. :meth:act therefore verifies that the arm it returns is admissible and raises if it is not, rather than silently emitting an inadmissible action or quietly substituting a different one. Where the arm list genuinely varies per decision, implement :class:InterventionalAgent directly instead of adapting a scalar agent.

deadline is accepted and ignored for the same reason: :meth:causalrl.Agent.act takes no budget, so there is nothing to pass it to.

Source code in src/causalrl/agents/interventional.py

class ScalarAgentAdapter(InterventionalAgent):
    """Present an arm-indexed :class:`causalrl.Agent` as an :class:`InterventionalAgent`.

    ``arms[i]`` is the intervention the wrapped agent's action ``i`` denotes, so the adapter is
    just a codebook in both directions: :meth:`act` maps the chosen index to its intervention, and
    :meth:`update` maps an intervention back to its index before forwarding the reward.

    **The wrapped agent's arm list is fixed, so ``space`` can only be checked, not honoured.** A
    scalar agent has no way to be told that arm 3 is unavailable this turn -- its policy is defined
    over all its arms. :meth:`act` therefore verifies that the arm it returns is admissible and
    raises if it is not, rather than silently emitting an inadmissible action or quietly
    substituting a different one. Where the arm list genuinely varies per decision, implement
    :class:`InterventionalAgent` directly instead of adapting a scalar agent.

    ``deadline`` is accepted and ignored for the same reason: :meth:`causalrl.Agent.act` takes no
    budget, so there is nothing to pass it to.
    """

    def __init__(self, agent: Agent, arms: Sequence[Intervention]) -> None:
        if not arms:
            raise ValueError("arms must be non-empty: an agent with no arms has nothing to choose")
        keys = [canonical(arm) for arm in arms]
        duplicates = sorted({key for key in keys if keys.count(key) > 1})
        if duplicates:
            raise ValueError(
                f"arms contains duplicate intervention(s) {[dict(k) for k in duplicates]!r}: the "
                "codebook must be one-to-one, or update() cannot tell which index a reward "
                "belongs to."
            )
        self._agent = agent
        self._arms = [dict(arm) for arm in arms]
        self._index_of = {key: i for i, key in enumerate(keys)}

    @property
    def arms(self) -> tuple[Intervention, ...]:
        """The codebook: ``arms[i]`` is the intervention denoted by the wrapped agent's action."""
        return tuple(dict(arm) for arm in self._arms)

    def act(
        self,
        observation: Mapping[str, Any],
        *,
        space: InterventionSpace,
        deadline: Deadline | None = None,
    ) -> Intervention:
        """Map the wrapped agent's chosen index to its intervention, checked against ``space``."""
        index = self._agent.act(dict(observation))
        if not 0 <= index < len(self._arms):
            raise IndexError(
                f"wrapped agent chose action {index}, outside the codebook of {len(self._arms)} "
                "arms: the agent's action space and the arms passed to ScalarAgentAdapter disagree"
            )
        return self.check_permitted(self._arms[index], space)

    def update(
        self, observation: Mapping[str, Any], intervention: Intervention, reward: float
    ) -> None:
        """Forward ``reward`` to the wrapped agent under the index that ``intervention`` encodes."""
        key = canonical(intervention)
        if key not in self._index_of:
            raise KeyError(
                f"intervention {dict(intervention)!r} is not in this adapter's codebook, so there "
                "is no arm index to credit the reward to. update() must be called with an "
                "intervention that act() returned."
            )
        self._agent.update(dict(observation), self._index_of[key], reward)

arms property

The codebook: arms[i] is the intervention denoted by the wrapped agent's action.

act(observation, *, space, deadline=None)

Map the wrapped agent's chosen index to its intervention, checked against space.

Source code in src/causalrl/agents/interventional.py

def act(
    self,
    observation: Mapping[str, Any],
    *,
    space: InterventionSpace,
    deadline: Deadline | None = None,
) -> Intervention:
    """Map the wrapped agent's chosen index to its intervention, checked against ``space``."""
    index = self._agent.act(dict(observation))
    if not 0 <= index < len(self._arms):
        raise IndexError(
            f"wrapped agent chose action {index}, outside the codebook of {len(self._arms)} "
            "arms: the agent's action space and the arms passed to ScalarAgentAdapter disagree"
        )
    return self.check_permitted(self._arms[index], space)

update(observation, intervention, reward)

Forward reward to the wrapped agent under the index that intervention encodes.

Source code in src/causalrl/agents/interventional.py

def update(
    self, observation: Mapping[str, Any], intervention: Intervention, reward: float
) -> None:
    """Forward ``reward`` to the wrapped agent under the index that ``intervention`` encodes."""
    key = canonical(intervention)
    if key not in self._index_of:
        raise KeyError(
            f"intervention {dict(intervention)!r} is not in this adapter's codebook, so there "
            "is no arm index to credit the reward to. update() must be called with an "
            "intervention that act() returned."
        )
    self._agent.update(dict(observation), self._index_of[key], reward)

causalrl.deadline.Deadline dataclass

A monotonic instant by which a decision must be returned.

Construct with :meth:after (a duration from now) rather than the raw constructor, which takes an absolute :func:time.monotonic reading and is mainly useful in tests where the clock is supplied explicitly.

Source code in src/causalrl/deadline.py

@dataclass(frozen=True)
class Deadline:
    """A monotonic instant by which a decision must be returned.

    Construct with :meth:`after` (a duration from now) rather than the raw constructor, which
    takes an absolute :func:`time.monotonic` reading and is mainly useful in tests where the clock
    is supplied explicitly.
    """

    at: float
    """Absolute :func:`time.monotonic` value at which the budget expires."""

    @classmethod
    def after(cls, seconds: float) -> Deadline:
        """A deadline ``seconds`` from now.

        A non-positive duration is allowed and yields an already-expired deadline: callers that
        compute a remaining budget by subtraction should get "no time left", not an exception.
        """
        return cls(time.monotonic() + seconds)

    def remaining(self) -> float:
        """Seconds left, clamped at ``0.0`` once the deadline has passed."""
        return max(0.0, self.at - time.monotonic())

    def expired(self) -> bool:
        """Whether the budget is exhausted."""
        return time.monotonic() >= self.at

    def fraction_of(self, share: float) -> float:
        """``share`` of the remaining budget, in seconds — for splitting time across sub-steps.

        ``share`` must lie in ``[0, 1]``: it is a fraction of what is left, not a multiplier, and
        a value above 1 would hand out a slice larger than the remaining budget under a name that
        says otherwise.
        """
        if not 0.0 <= share <= 1.0:
            raise ValueError(
                f"share={share} must lie in [0, 1]: fraction_of hands out that fraction of the "
                "remaining budget, so a value above 1 would return more time than the deadline "
                "allows while claiming to be a share of it."
            )
        return self.remaining() * share

at instance-attribute

Absolute :func:time.monotonic value at which the budget expires.

after(seconds) classmethod

A deadline seconds from now.

A non-positive duration is allowed and yields an already-expired deadline: callers that compute a remaining budget by subtraction should get "no time left", not an exception.

Source code in src/causalrl/deadline.py

@classmethod
def after(cls, seconds: float) -> Deadline:
    """A deadline ``seconds`` from now.

    A non-positive duration is allowed and yields an already-expired deadline: callers that
    compute a remaining budget by subtraction should get "no time left", not an exception.
    """
    return cls(time.monotonic() + seconds)

remaining()

Seconds left, clamped at 0.0 once the deadline has passed.

Source code in src/causalrl/deadline.py

def remaining(self) -> float:
    """Seconds left, clamped at ``0.0`` once the deadline has passed."""
    return max(0.0, self.at - time.monotonic())

expired()

Whether the budget is exhausted.

Source code in src/causalrl/deadline.py

def expired(self) -> bool:
    """Whether the budget is exhausted."""
    return time.monotonic() >= self.at

fraction_of(share)

share of the remaining budget, in seconds — for splitting time across sub-steps.

share must lie in [0, 1]: it is a fraction of what is left, not a multiplier, and a value above 1 would hand out a slice larger than the remaining budget under a name that says otherwise.

Source code in src/causalrl/deadline.py

def fraction_of(self, share: float) -> float:
    """``share`` of the remaining budget, in seconds — for splitting time across sub-steps.

    ``share`` must lie in ``[0, 1]``: it is a fraction of what is left, not a multiplier, and
    a value above 1 would hand out a slice larger than the remaining budget under a name that
    says otherwise.
    """
    if not 0.0 <= share <= 1.0:
        raise ValueError(
            f"share={share} must lie in [0, 1]: fraction_of hands out that fraction of the "
            "remaining budget, so a value above 1 would return more time than the deadline "
            "allows while claiming to be a share of it."
        )
    return self.remaining() * share

Known Mechanisms

causalrl.scm.fitters.PinnedMechanism

Deploy a known structural equation at this node instead of fitting one.

Real systems are rarely all-known or all-unknown: a documented rule, a tariff, a physical law or a contractual formula sits next to behaviour nobody has a closed form for. This fitter is the "known" half. Pass it to fit_scm(..., families={"V": PinnedMechanism(g)}) and node V deploys V = g(parents) + U exactly as given, while every other node is learned from the data as usual.

mean is the deterministic part g, taking the same dict[str, Tensor] of parent columns every mechanism in this module receives and returning one column. The noise enters additively, so the mechanism is invertible and counterfactuals at this node stay identified — which is part of the point: pinning a mechanism can restore an identification that a fitted node would have lost.

noise is the exogenous distribution. Leaving it None estimates it by resampling the residuals V - g(parents) on the training rows, so the shape of the equation is pinned while the spread around it still comes from data; supply a distribution to pin that too.

A pinned node is still scored, and that score is the assumption's test. fit_scm reports :attr:~causalrl.scm.fit.NodeFit.holdout_score for this node like any other — but here it is not a measure of how well a flexible family fitted, it is a measure of whether the equation you asserted actually describes held-out data. A pinned node with a poor holdout score is telling you the assertion is wrong, and nothing else in the pipeline will.

Source code in src/causalrl/scm/fitters.py

class PinnedMechanism:
    """Deploy a **known** structural equation at this node instead of fitting one.

    Real systems are rarely all-known or all-unknown: a documented rule, a tariff, a physical law
    or a contractual formula sits next to behaviour nobody has a closed form for. This fitter is
    the "known" half. Pass it to ``fit_scm(..., families={"V": PinnedMechanism(g)})`` and node
    ``V`` deploys ``V = g(parents) + U`` exactly as given, while every other node is learned from
    the data as usual.

    ``mean`` is the deterministic part ``g``, taking the same ``dict[str, Tensor]`` of parent
    columns every mechanism in this module receives and returning one column. The noise enters
    additively, so the mechanism is invertible and counterfactuals at this node stay identified —
    which is part of the point: pinning a mechanism can restore an identification that a fitted
    node would have lost.

    ``noise`` is the exogenous distribution. Leaving it ``None`` estimates it by resampling the
    residuals ``V - g(parents)`` on the training rows, so the *shape* of the equation is pinned
    while the spread around it still comes from data; supply a distribution to pin that too.

    **A pinned node is still scored, and that score is the assumption's test.** ``fit_scm`` reports
    :attr:`~causalrl.scm.fit.NodeFit.holdout_score` for this node like any other — but here it is
    not a measure of how well a flexible family fitted, it is a measure of whether the equation you
    asserted actually describes held-out data. A pinned node with a poor holdout score is telling
    you the assertion is wrong, and nothing else in the pipeline will.
    """

    def __init__(
        self,
        mean: Callable[[dict[str, Tensor]], Tensor],
        *,
        noise: Distribution | None = None,
    ) -> None:
        self._mean = mean
        self._noise = noise

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        y = np.asarray(child, dtype=float)
        names = sorted(parents)
        parent_tensors = {
            name: torch.tensor(  # type: ignore[reportPrivateImportUsage]
                np.asarray(parents[name], dtype=float),
                dtype=torch.float32,  # type: ignore[reportPrivateImportUsage]
            )
            for name in names
        }
        with torch.no_grad():
            predicted_tensor = self._mean(parent_tensors).reshape(-1)
        if predicted_tensor.shape[0] == 1 and len(y) != 1:
            # A parentless (or otherwise constant) equation has nothing to vary over and may
            # legitimately return a single value; broadcast it to the column it stands for.
            predicted_tensor = predicted_tensor.expand(len(y))
        elif predicted_tensor.shape[0] != len(y):
            raise ValueError(
                f"PinnedMechanism's mean returned {predicted_tensor.shape[0]} value(s) for "
                f"{len(y)} training row(s): it must map the parent columns it is given to one "
                "column of the same length (or a single value, for a constant equation), so that "
                "the residual V - g(parents) is well defined."
            )
        predicted = predicted_tensor.numpy().astype(float)
        residual = torch.tensor(  # type: ignore[reportPrivateImportUsage]
            y - predicted,
            dtype=torch.float32,  # type: ignore[reportPrivateImportUsage]
        )

        mean_fn = self._mean

        def mechanism(parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
            return mean_fn(parent_values).reshape(-1) + noise.reshape(-1)

        def residual_closure(parent_values: dict[str, Tensor], value: Tensor) -> Tensor:
            """Exact abduction: the noise is recovered as ``V - g(parents)``, nothing fitted."""
            return value.reshape(-1) - mean_fn(parent_values).reshape(-1)

        fitted_mechanism = FunctionalMechanism(names, mechanism)
        fitted_mechanism.residual = residual_closure  # type: ignore[attr-defined]
        return FittedMechanism(
            mechanism=fitted_mechanism,
            noise=self._noise if self._noise is not None else _Empirical(residual),
            invertible=True,
            score=_r2(y, predicted),
        )

Continuous States

causalrl.state.StateEncoder

Bases: Protocol

Maps an observation to a fixed-length feature vector.

dim is the length of every vector :meth:encode returns; downstream learners size their design matrices from it, so it must not vary between calls.

Source code in src/causalrl/state.py

@runtime_checkable
class StateEncoder(Protocol):
    """Maps an observation to a fixed-length feature vector.

    ``dim`` is the length of every vector :meth:`encode` returns; downstream learners size their
    design matrices from it, so it must not vary between calls.
    """

    dim: int

    def encode(self, observation: Mapping[str, Any]) -> FloatArray: ...

causalrl.state.OneHotEncoder

Encode a discrete state index as an indicator vector — the tabular case, contained.

This is the bridge between the two worlds. Feature-space learning with these features and an exactly-solving linear learner reproduces the tabular backup, because the indicator basis spans every function on a finite state set. That makes tabular behaviour a special case of the fitted path rather than a separate code path, and it is what the faithfulness tests check.

key names the observation entry holding the state index, matching the "state" key the shipped tabular agents already read.

Source code in src/causalrl/state.py

class OneHotEncoder:
    """Encode a discrete state index as an indicator vector — the tabular case, contained.

    This is the bridge between the two worlds. Feature-space learning with these features and an
    exactly-solving linear learner reproduces the tabular backup, because the indicator basis spans
    every function on a finite state set. That makes tabular behaviour a *special case* of the
    fitted path rather than a separate code path, and it is what the faithfulness tests check.

    ``key`` names the observation entry holding the state index, matching the ``"state"`` key the
    shipped tabular agents already read.
    """

    def __init__(self, n_states: int, *, key: str = "state") -> None:
        if n_states < 1:
            raise ValueError(f"n_states={n_states} must be at least 1")
        self.dim = n_states
        self.n_states = n_states
        self.key = key

    def encode(self, observation: Mapping[str, Any]) -> FloatArray:
        index = int(observation[self.key])
        if not 0 <= index < self.n_states:
            raise IndexError(
                f"state index {index} is outside [0, {self.n_states}) -- OneHotEncoder was built "
                f"for {self.n_states} states, so this observation comes from a different space."
            )
        vector = np.zeros(self.n_states, dtype=np.float64)
        vector[index] = 1.0
        return vector

causalrl.state.IdentityEncoder

Read named continuous entries out of the observation, in the order given.

The minimal continuous encoder: no basis expansion, so a linear learner over these features fits a linear value function. keys fixes the column order, which must stay stable for the features to mean the same thing across calls.

Source code in src/causalrl/state.py

class IdentityEncoder:
    """Read named continuous entries out of the observation, in the order given.

    The minimal continuous encoder: no basis expansion, so a linear learner over these features
    fits a linear value function. ``keys`` fixes the column order, which must stay stable for the
    features to mean the same thing across calls.
    """

    def __init__(self, keys: Sequence[str]) -> None:
        if not keys:
            raise ValueError("keys must name at least one observation entry")
        self.keys = tuple(keys)
        self.dim = len(self.keys)

    def encode(self, observation: Mapping[str, Any]) -> FloatArray:
        return np.array([float(observation[key]) for key in self.keys], dtype=np.float64)

causalrl.state.RBFEncoder

Gaussian radial basis features over an inner encoder's output.

The nonlinear counterpart to :class:IdentityEncoder, and the same basis :class:causalrl.FunctionApproxBackdoorAgent already uses for a continuous confounder — the point being that a continuous state and a continuous confounder want the same machinery, which is precisely why the tabular state interface was the odd one out.

centers has one row per basis function, each of the inner encoder's dimension. A constant column is prepended so a learner can represent an offset without a separate intercept term.

Source code in src/causalrl/state.py

class RBFEncoder:
    """Gaussian radial basis features over an inner encoder's output.

    The nonlinear counterpart to :class:`IdentityEncoder`, and the same basis
    :class:`causalrl.FunctionApproxBackdoorAgent` already uses for a continuous confounder — the
    point being that a continuous *state* and a continuous *confounder* want the same machinery,
    which is precisely why the tabular state interface was the odd one out.

    ``centers`` has one row per basis function, each of the inner encoder's dimension. A constant
    column is prepended so a learner can represent an offset without a separate intercept term.
    """

    def __init__(self, inner: StateEncoder, centers: FloatArray, *, bandwidth: float = 1.0) -> None:
        grid = np.asarray(centers, dtype=np.float64)
        if grid.ndim != 2 or grid.shape[1] != inner.dim:
            raise ValueError(
                f"centers must have shape (n_centers, {inner.dim}) to match the inner encoder's "
                f"dimension; got {grid.shape}"
            )
        if bandwidth <= 0.0:
            raise ValueError(f"bandwidth={bandwidth} must be positive")
        self._inner = inner
        self._centers = grid
        self._bandwidth = bandwidth
        self.dim = grid.shape[0] + 1

    def encode(self, observation: Mapping[str, Any]) -> FloatArray:
        raw = self._inner.encode(observation)
        squared = ((raw[None, :] - self._centers) ** 2).sum(axis=1)
        return np.concatenate(
            [np.ones(1, dtype=np.float64), np.exp(-0.5 * squared / self._bandwidth**2)]
        )

causalrl.state.FeatureTransition dataclass

A logged transition whose endpoints are feature vectors rather than state indices.

The feature-space counterpart of :class:causalrl.data.dataset.Transition, and the argument type of the fitted backup. done marks an absorbing terminal, whose successor carries no future value — the same convention the tabular agents use, kept so that one-hot features reproduce their behaviour exactly.

Source code in src/causalrl/state.py

@dataclass(frozen=True)
class FeatureTransition:
    """A logged transition whose endpoints are feature vectors rather than state indices.

    The feature-space counterpart of :class:`causalrl.data.dataset.Transition`, and the argument
    type of the fitted backup. ``done`` marks an absorbing terminal, whose successor carries no
    future value — the same convention the tabular agents use, kept so that one-hot features
    reproduce their behaviour exactly.
    """

    state: FloatArray
    action: int
    reward: float
    next_state: FloatArray
    done: bool

    def __post_init__(self) -> None:
        state = np.asarray(self.state, dtype=np.float64).reshape(-1)
        next_state = np.asarray(self.next_state, dtype=np.float64).reshape(-1)
        if state.shape != next_state.shape:
            raise ValueError(
                f"state and next_state must share a feature dimension; got {state.shape} and "
                f"{next_state.shape}. Both endpoints must come from the same encoder."
            )
        object.__setattr__(self, "state", state)
        object.__setattr__(self, "next_state", next_state)

causalrl.state.encode_batch(encoder, observations)

Stack encoder.encode over observations into an (n, dim) design matrix.

Returns a correctly-shaped empty array for an empty sequence, so a caller assembling a design matrix from a possibly-empty slice does not have to special-case it.

Source code in src/causalrl/state.py

def encode_batch(encoder: StateEncoder, observations: Sequence[Mapping[str, Any]]) -> FloatArray:
    """Stack ``encoder.encode`` over ``observations`` into an ``(n, dim)`` design matrix.

    Returns a correctly-shaped empty array for an empty sequence, so a caller assembling a design
    matrix from a possibly-empty slice does not have to special-case it.
    """
    if not observations:
        return np.zeros((0, encoder.dim), dtype=np.float64)
    return np.stack([np.asarray(encoder.encode(obs), dtype=np.float64) for obs in observations])

causalrl.agents.fitted.FittedQIteration

Bases: BatchAgent

Finite-horizon fitted Q iteration over encoded state features.

Backward induction, identical in shape to :class:causalrl.DOVI's but with the table replaced by one fitted regressor per (step, action)::

V_{H+1}(s)  = 0
Q_h(·, a)   = fit( features(s) -> r + [not done] · V_{h+1}(s') )   over rows with A = a
Q_h         = min(Q_h, reward_max · (H - h + 1))      # the global optimism cap
V_h(s)      = max_a Q_h(s, a)

encoder decides what a state is. With :class:~causalrl.state.OneHotEncoder and a least-squares regressor the fitted backup reproduces the tabular one, because the indicator basis spans every function on a finite state set — the tabular agent is a special case of this one, not a separate algorithm.

An action never taken at some step has no data to fit, so its Q is set to the cap. That keeps the backup optimistic in the unexplored direction, matching how the tabular agent treats an unvisited cell.

reward_max must be a genuine per-step upper bound on the reward; the cap is only as valid as that number. Rewards are not otherwise assumed bounded below.

Source code in src/causalrl/agents/fitted.py

class FittedQIteration(BatchAgent):
    """Finite-horizon fitted Q iteration over encoded state features.

    Backward induction, identical in shape to :class:`causalrl.DOVI`'s but with the table replaced
    by one fitted regressor per ``(step, action)``::

        V_{H+1}(s)  = 0
        Q_h(·, a)   = fit( features(s) -> r + [not done] · V_{h+1}(s') )   over rows with A = a
        Q_h         = min(Q_h, reward_max · (H - h + 1))      # the global optimism cap
        V_h(s)      = max_a Q_h(s, a)

    ``encoder`` decides what a state *is*. With :class:`~causalrl.state.OneHotEncoder` and a
    least-squares regressor the fitted backup reproduces the tabular one, because the indicator
    basis spans every function on a finite state set — the tabular agent is a special case of this
    one, not a separate algorithm.

    An action never taken at some step has no data to fit, so its Q is set to the cap. That keeps
    the backup optimistic in the unexplored direction, matching how the tabular agent treats an
    unvisited cell.

    ``reward_max`` must be a genuine per-step upper bound on the reward; the cap is only as valid
    as that number. Rewards are not otherwise assumed bounded below.
    """

    def __init__(
        self,
        n_actions: int,
        horizon: int,
        encoder: StateEncoder,
        *,
        regressor: Callable[[], Regressor] | None = None,
        reward_max: float = 1.0,
        seed: int | None = None,
    ) -> None:
        if n_actions < 1:
            raise ValueError(f"n_actions={n_actions} must be at least 1")
        if horizon < 1:
            raise ValueError(f"horizon={horizon} must be at least 1")
        self.n_actions = n_actions
        self.horizon = horizon
        self.encoder = encoder
        self.reward_max = float(reward_max)
        self._make_regressor = regressor if regressor is not None else RidgeRegressor
        self._buffer: list[FeatureTransition] = []
        self._models: dict[tuple[int, int], Regressor] = {}
        self._fitted = False
        self._rng = np.random.default_rng(seed)

    @property
    def is_certified(self) -> bool:
        """Always ``False``: a fitted backup carries no per-cell causal bound.

        Named to match :attr:`causalrl.DOVI.is_certified` so the two are directly comparable. The
        tabular agent can answer ``True`` because its optimism is capped by a Manski bound computed
        per ``(state, action)``; this one caps optimism globally by horizon and reward range, which
        is valid but is not an identification guarantee.
        """
        return False

    def cap(self, step: int) -> float:
        """The global optimism cap at step ``step``: ``reward_max * (horizon - step + 1)``.

        The largest return achievable from ``step`` onward when each of the remaining steps pays at
        most ``reward_max``. Valid with no assumption on the function class, and much weaker than a
        per-cell bound.
        """
        if not 1 <= step <= self.horizon:
            raise ValueError(f"step={step} must lie in [1, {self.horizon}]")
        return self.reward_max * (self.horizon - step + 1)

    def observe(self, transition: FeatureTransition) -> None:
        """Add one feature-typed transition to the buffer and invalidate the current plan."""
        self._buffer.append(transition)
        self._fitted = False

    def buffered_transitions(self) -> tuple[FeatureTransition, ...]:
        """The transitions observed so far and not yet discarded by a ``fit(transitions=...)``.

        The online buffer is otherwise write-only: a caller driving this agent through
        :meth:`observe_step` had no way to see what it had actually collected, which makes
        "did my driver wire the hook up correctly" unanswerable without reaching into a private
        attribute.
        """
        return tuple(self._buffer)

    def observe_transition(self, state: int, action: int, next_state: int, done: bool) -> None:
        """Refused: this agent's states are feature vectors, not indices.

        The tabular hook is overridden rather than left as the inherited no-op, which would accept
        the call and silently discard the transition. Encode the endpoints and call
        :meth:`observe` with a :class:`~causalrl.state.FeatureTransition` instead.
        """
        raise NotImplementedError(
            "FittedQIteration takes feature-typed transitions: the inherited "
            "observe_transition(state: int, ...) hook cannot express a continuous state, and "
            "accepting it here would silently drop the transition. Encode both endpoints with "
            "this agent's encoder and call observe(FeatureTransition(...))."
        )

    def observe_step(
        self,
        observation: dict[str, Any],
        action: int,
        reward: float,
        next_observation: dict[str, Any],
        done: bool,
    ) -> None:
        """Encode both endpoints with this agent's encoder and buffer the transition.

        The feature-space answer to the tabular hook above: a driver hands over raw observations
        and the agent encodes them, so :func:`causalrl.eval.harness.run_episodes` can drive this
        agent without knowing anything about its state representation.
        """
        self.observe(
            FeatureTransition(
                state=self.encoder.encode(observation),
                action=action,
                reward=float(reward),
                next_state=self.encoder.encode(next_observation),
                done=done,
            )
        )

    def fit(self, transitions: Sequence[FeatureTransition] | None = None) -> FittedQIteration:
        """Run the backward induction over ``transitions`` (or the observed buffer) and plan.

        Passing ``transitions`` replaces the buffer, so a caller can refit from a fresh batch
        without carrying stale online data. Raises ``ValueError`` on an empty dataset — there is no
        such thing as a fitted value function with nothing to fit it on.
        """
        if transitions is not None:
            self._buffer = list(transitions)

        batch = unpack_transitions(self._buffer, encoder=self.encoder, n_actions=self.n_actions)
        states, next_states = batch.states, batch.next_states
        actions, rewards, not_done = batch.actions, batch.rewards, batch.not_done

        self._models = {}
        value_at_next = np.zeros(len(self._buffer), dtype=np.float64)
        for step in range(self.horizon, 0, -1):
            targets = np.minimum(rewards + not_done * value_at_next, self.cap(step))
            for action in range(self.n_actions):
                rows = actions == action
                if not rows.any():
                    continue  # no data: q_values() falls back to the cap for this (step, action)
                model = self._make_regressor()
                model.fit(states[rows], targets[rows])
                self._models[(step, action)] = model
            value_at_next = self._value(next_states, step)
        self._fitted = True
        return self

    def _q_matrix(self, features: FloatArray, step: int) -> FloatArray:
        """``(n_rows, n_actions)`` capped Q values at ``step`` for a batch of feature rows."""
        cap = self.cap(step)
        out = np.full((features.shape[0], self.n_actions), cap, dtype=np.float64)
        for action in range(self.n_actions):
            model = self._models.get((step, action))
            if model is not None:
                out[:, action] = np.minimum(np.asarray(model.predict(features)).reshape(-1), cap)
        return out

    def _value(self, features: FloatArray, step: int) -> FloatArray:
        return self._q_matrix(features, step).max(axis=1)

    def q_values(self, observation: Mapping[str, Any], step: int = 1) -> FloatArray:
        """Capped Q values for every action at ``observation``, as a ``(n_actions,)`` array."""
        self._require_fit()
        features = np.asarray(self.encoder.encode(observation), dtype=np.float64).reshape(1, -1)
        return self._q_matrix(features, step)[0]

    def value(self, observation: Mapping[str, Any], step: int = 1) -> float:
        """``max_a Q_step(observation, a)`` — the fitted value at this observation."""
        return float(self.q_values(observation, step).max())

    def act(self, observation: dict[str, Any]) -> int:
        """Greedy action at the observation's step, ties broken uniformly at random.

        The step is read from ``observation["t"]`` (0-indexed, as the shipped environments emit it)
        and clamped to the horizon, matching :meth:`causalrl.DOVI.act`.
        """
        step = min(int(observation.get("t", 0)) + 1, self.horizon)
        scores = self.q_values(observation, step)
        best = float(scores.max())
        winners = [a for a in range(self.n_actions) if scores[a] >= best - 1e-12]
        return int(self._rng.choice(winners))

    def _require_fit(self) -> None:
        if not self._fitted:
            raise RuntimeError(
                "FittedQIteration has no plan yet: call fit() after supplying transitions. "
                "Acting before fitting would return the optimism cap for every action, which is "
                "an arbitrary tie rather than a decision."
            )

    def certificate(self) -> Certificate:
        """An ``EMPIRICAL`` certificate recording exactly which guarantee was given up.

        The hedge carries ``downgraded_from="bounded"``: the tabular path is partially identified
        by a per-cell Manski bound, and replacing the table with a function class forfeits that in
        exchange for a global horizon cap. Reporting this as ``BOUNDED`` would claim an
        identification status the fitted backup does not have.
        """
        return Certificate(
            claim=(
                f"fitted Q iteration over {self.encoder.dim}-dimensional state features "
                f"(horizon={self.horizon}, actions={self.n_actions})"
            ),
            estimand=EstimandSpec(query="policy_value", target="mean"),
            kind=Kind.EMPIRICAL,
            value=None,
            alpha=None,
            assumptions=(
                Assumption(
                    name="reward-ceiling",
                    params={"reward_max": self.reward_max, "horizon": self.horizon},
                    checkable=False,
                ),
                Assumption(
                    name="function-class",
                    params={
                        "encoder": type(self.encoder).__name__,
                        "dim": self.encoder.dim,
                        "regressor": getattr(self._make_regressor, "__name__", "callable"),
                    },
                    checkable=False,
                ),
            ),
            method="fitted_q_iteration",
            witness=None,
            hedge=Hedge(
                reason=(
                    "optimism is capped globally by reward_max * steps remaining, not by a "
                    "per-(state, action) Manski bound: in feature space there are no cells to "
                    "bound, so the tabular partial-identification guarantee does not carry over"
                ),
                detail={"tabular_equivalent": "DOVI", "global_cap_at_step_1": self.cap(1)},
                downgraded_from="bounded",
            ),
            provenance=Provenance.create(),
        )

is_certified property

Always False: a fitted backup carries no per-cell causal bound.

Named to match :attr:causalrl.DOVI.is_certified so the two are directly comparable. The tabular agent can answer True because its optimism is capped by a Manski bound computed per (state, action); this one caps optimism globally by horizon and reward range, which is valid but is not an identification guarantee.

cap(step)

The global optimism cap at step step: reward_max * (horizon - step + 1).

The largest return achievable from step onward when each of the remaining steps pays at most reward_max. Valid with no assumption on the function class, and much weaker than a per-cell bound.

Source code in src/causalrl/agents/fitted.py

def cap(self, step: int) -> float:
    """The global optimism cap at step ``step``: ``reward_max * (horizon - step + 1)``.

    The largest return achievable from ``step`` onward when each of the remaining steps pays at
    most ``reward_max``. Valid with no assumption on the function class, and much weaker than a
    per-cell bound.
    """
    if not 1 <= step <= self.horizon:
        raise ValueError(f"step={step} must lie in [1, {self.horizon}]")
    return self.reward_max * (self.horizon - step + 1)

observe(transition)

Add one feature-typed transition to the buffer and invalidate the current plan.

Source code in src/causalrl/agents/fitted.py

def observe(self, transition: FeatureTransition) -> None:
    """Add one feature-typed transition to the buffer and invalidate the current plan."""
    self._buffer.append(transition)
    self._fitted = False

buffered_transitions()

The transitions observed so far and not yet discarded by a fit(transitions=...).

The online buffer is otherwise write-only: a caller driving this agent through :meth:observe_step had no way to see what it had actually collected, which makes "did my driver wire the hook up correctly" unanswerable without reaching into a private attribute.

Source code in src/causalrl/agents/fitted.py

def buffered_transitions(self) -> tuple[FeatureTransition, ...]:
    """The transitions observed so far and not yet discarded by a ``fit(transitions=...)``.

    The online buffer is otherwise write-only: a caller driving this agent through
    :meth:`observe_step` had no way to see what it had actually collected, which makes
    "did my driver wire the hook up correctly" unanswerable without reaching into a private
    attribute.
    """
    return tuple(self._buffer)

observe_transition(state, action, next_state, done)

Refused: this agent's states are feature vectors, not indices.

The tabular hook is overridden rather than left as the inherited no-op, which would accept the call and silently discard the transition. Encode the endpoints and call :meth:observe with a :class:~causalrl.state.FeatureTransition instead.

Source code in src/causalrl/agents/fitted.py

def observe_transition(self, state: int, action: int, next_state: int, done: bool) -> None:
    """Refused: this agent's states are feature vectors, not indices.

    The tabular hook is overridden rather than left as the inherited no-op, which would accept
    the call and silently discard the transition. Encode the endpoints and call
    :meth:`observe` with a :class:`~causalrl.state.FeatureTransition` instead.
    """
    raise NotImplementedError(
        "FittedQIteration takes feature-typed transitions: the inherited "
        "observe_transition(state: int, ...) hook cannot express a continuous state, and "
        "accepting it here would silently drop the transition. Encode both endpoints with "
        "this agent's encoder and call observe(FeatureTransition(...))."
    )

observe_step(observation, action, reward, next_observation, done)

Encode both endpoints with this agent's encoder and buffer the transition.

The feature-space answer to the tabular hook above: a driver hands over raw observations and the agent encodes them, so :func:causalrl.eval.harness.run_episodes can drive this agent without knowing anything about its state representation.

Source code in src/causalrl/agents/fitted.py

def observe_step(
    self,
    observation: dict[str, Any],
    action: int,
    reward: float,
    next_observation: dict[str, Any],
    done: bool,
) -> None:
    """Encode both endpoints with this agent's encoder and buffer the transition.

    The feature-space answer to the tabular hook above: a driver hands over raw observations
    and the agent encodes them, so :func:`causalrl.eval.harness.run_episodes` can drive this
    agent without knowing anything about its state representation.
    """
    self.observe(
        FeatureTransition(
            state=self.encoder.encode(observation),
            action=action,
            reward=float(reward),
            next_state=self.encoder.encode(next_observation),
            done=done,
        )
    )

fit(transitions=None)

Run the backward induction over transitions (or the observed buffer) and plan.

Passing transitions replaces the buffer, so a caller can refit from a fresh batch without carrying stale online data. Raises ValueError on an empty dataset — there is no such thing as a fitted value function with nothing to fit it on.

Source code in src/causalrl/agents/fitted.py

def fit(self, transitions: Sequence[FeatureTransition] | None = None) -> FittedQIteration:
    """Run the backward induction over ``transitions`` (or the observed buffer) and plan.

    Passing ``transitions`` replaces the buffer, so a caller can refit from a fresh batch
    without carrying stale online data. Raises ``ValueError`` on an empty dataset — there is no
    such thing as a fitted value function with nothing to fit it on.
    """
    if transitions is not None:
        self._buffer = list(transitions)

    batch = unpack_transitions(self._buffer, encoder=self.encoder, n_actions=self.n_actions)
    states, next_states = batch.states, batch.next_states
    actions, rewards, not_done = batch.actions, batch.rewards, batch.not_done

    self._models = {}
    value_at_next = np.zeros(len(self._buffer), dtype=np.float64)
    for step in range(self.horizon, 0, -1):
        targets = np.minimum(rewards + not_done * value_at_next, self.cap(step))
        for action in range(self.n_actions):
            rows = actions == action
            if not rows.any():
                continue  # no data: q_values() falls back to the cap for this (step, action)
            model = self._make_regressor()
            model.fit(states[rows], targets[rows])
            self._models[(step, action)] = model
        value_at_next = self._value(next_states, step)
    self._fitted = True
    return self

q_values(observation, step=1)

Capped Q values for every action at observation, as a (n_actions,) array.

Source code in src/causalrl/agents/fitted.py

def q_values(self, observation: Mapping[str, Any], step: int = 1) -> FloatArray:
    """Capped Q values for every action at ``observation``, as a ``(n_actions,)`` array."""
    self._require_fit()
    features = np.asarray(self.encoder.encode(observation), dtype=np.float64).reshape(1, -1)
    return self._q_matrix(features, step)[0]

value(observation, step=1)

max_a Q_step(observation, a) — the fitted value at this observation.

Source code in src/causalrl/agents/fitted.py

def value(self, observation: Mapping[str, Any], step: int = 1) -> float:
    """``max_a Q_step(observation, a)`` — the fitted value at this observation."""
    return float(self.q_values(observation, step).max())

act(observation)

Greedy action at the observation's step, ties broken uniformly at random.

The step is read from observation["t"] (0-indexed, as the shipped environments emit it) and clamped to the horizon, matching :meth:causalrl.DOVI.act.

Source code in src/causalrl/agents/fitted.py

def act(self, observation: dict[str, Any]) -> int:
    """Greedy action at the observation's step, ties broken uniformly at random.

    The step is read from ``observation["t"]`` (0-indexed, as the shipped environments emit it)
    and clamped to the horizon, matching :meth:`causalrl.DOVI.act`.
    """
    step = min(int(observation.get("t", 0)) + 1, self.horizon)
    scores = self.q_values(observation, step)
    best = float(scores.max())
    winners = [a for a in range(self.n_actions) if scores[a] >= best - 1e-12]
    return int(self._rng.choice(winners))

certificate()

An EMPIRICAL certificate recording exactly which guarantee was given up.

The hedge carries downgraded_from="bounded": the tabular path is partially identified by a per-cell Manski bound, and replacing the table with a function class forfeits that in exchange for a global horizon cap. Reporting this as BOUNDED would claim an identification status the fitted backup does not have.

Source code in src/causalrl/agents/fitted.py

def certificate(self) -> Certificate:
    """An ``EMPIRICAL`` certificate recording exactly which guarantee was given up.

    The hedge carries ``downgraded_from="bounded"``: the tabular path is partially identified
    by a per-cell Manski bound, and replacing the table with a function class forfeits that in
    exchange for a global horizon cap. Reporting this as ``BOUNDED`` would claim an
    identification status the fitted backup does not have.
    """
    return Certificate(
        claim=(
            f"fitted Q iteration over {self.encoder.dim}-dimensional state features "
            f"(horizon={self.horizon}, actions={self.n_actions})"
        ),
        estimand=EstimandSpec(query="policy_value", target="mean"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=None,
        assumptions=(
            Assumption(
                name="reward-ceiling",
                params={"reward_max": self.reward_max, "horizon": self.horizon},
                checkable=False,
            ),
            Assumption(
                name="function-class",
                params={
                    "encoder": type(self.encoder).__name__,
                    "dim": self.encoder.dim,
                    "regressor": getattr(self._make_regressor, "__name__", "callable"),
                },
                checkable=False,
            ),
        ),
        method="fitted_q_iteration",
        witness=None,
        hedge=Hedge(
            reason=(
                "optimism is capped globally by reward_max * steps remaining, not by a "
                "per-(state, action) Manski bound: in feature space there are no cells to "
                "bound, so the tabular partial-identification guarantee does not carry over"
            ),
            detail={"tabular_equivalent": "DOVI", "global_cap_at_step_1": self.cap(1)},
            downgraded_from="bounded",
        ),
        provenance=Provenance.create(),
    )

Function-Valued Causal Bounds

causalrl.bounds.functional.FunctionalManskiBounds

Cross-fitted pointwise Manski bounds on E[R | do(a), X = x].

outcome_model and propensity_model are factories returning a fresh :class:~causalrl.estimate.nuisance.Regressor / :class:~causalrl.estimate.nuisance.Classifier — the same duck-typed protocols the DML estimators use, so any scikit-learn estimator works and the defaults keep the core dependency-free. The propensity is fitted one-vs-rest and normalised across actions.

reward_range must genuinely contain the reward; it is what the unlogged fraction is bounded by, and the bound is only valid if it holds. It defaults to (0, 1), matching :func:causalrl.causal_q_bounds.

Cross-fitting produces two different predictors, deliberately kept apart: :attr:in_sample are out-of-fold values for the training rows (each from a model that never saw that row), and :meth:bounds evaluates new points by averaging the fold models. Use the former to summarise the data you fitted on and the latter to bound anywhere else.

Source code in src/causalrl/bounds/functional.py

class FunctionalManskiBounds:
    """Cross-fitted pointwise Manski bounds on ``E[R | do(a), X = x]``.

    ``outcome_model`` and ``propensity_model`` are factories returning a fresh
    :class:`~causalrl.estimate.nuisance.Regressor` / :class:`~causalrl.estimate.nuisance.Classifier`
    — the same duck-typed protocols the DML estimators use, so any scikit-learn estimator works and
    the defaults keep the core dependency-free. The propensity is fitted one-vs-rest and normalised
    across actions.

    ``reward_range`` must genuinely contain the reward; it is what the unlogged fraction is bounded
    by, and the bound is only valid if it holds. It defaults to ``(0, 1)``, matching
    :func:`causalrl.causal_q_bounds`.

    Cross-fitting produces two different predictors, deliberately kept apart:
    :attr:`in_sample` are out-of-fold values for the training rows (each from a model that never saw
    that row), and :meth:`bounds` evaluates new points by averaging the fold models. Use the former
    to summarise the data you fitted on and the latter to bound anywhere else.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        outcome_model: Callable[[], Regressor] | None = None,
        propensity_model: Callable[[], Classifier] | None = None,
        reward_range: tuple[float, float] = (0.0, 1.0),
        n_folds: int = 5,
        seed: int = 0,
    ) -> None:
        if n_actions < 1:
            raise ValueError(f"n_actions={n_actions} must be at least 1")
        if n_folds < 2:
            raise ValueError(
                f"n_folds={n_folds} must be at least 2: cross-fitting needs a held-out part, and "
                "a single fold would fit and predict on the same rows -- the in-sample plug-in "
                "this class exists to avoid, since it makes the interval optimistically tight."
            )
        low, high = float(reward_range[0]), float(reward_range[1])
        if not low < high:
            raise ValueError(f"reward_range={reward_range} must satisfy low < high")
        self.n_actions = n_actions
        self.reward_range = (low, high)
        self.n_folds = n_folds
        self.seed = seed
        self._outcome = outcome_model if outcome_model is not None else RidgeRegressor
        self._propensity = propensity_model if propensity_model is not None else LogisticRegressor
        self._fold_models: list[tuple[dict[int, Regressor], dict[int, Classifier]]] = []
        self._in_sample: tuple[FloatArray, FloatArray] | None = None
        self._diagnostic: OverlapDiagnostic | None = None

    def fit(
        self, features: FloatArray, actions: NDArray[np.int_], rewards: FloatArray
    ) -> FunctionalManskiBounds:
        """Cross-fit the outcome and propensity models and record the out-of-fold bounds."""
        x = np.asarray(features, dtype=np.float64)
        a = np.asarray(actions, dtype=int).reshape(-1)
        r = np.asarray(rewards, dtype=np.float64).reshape(-1)
        if x.ndim != 2:
            raise ValueError(f"features must be a 2-D (n, dim) design matrix; got shape {x.shape}")
        if not (len(x) == len(a) == len(r)):
            raise ValueError(
                f"features, actions and rewards must agree in length; got {len(x)}, {len(a)}, "
                f"{len(r)}"
            )
        if len(x) < self.n_folds:
            raise ValueError(
                f"{len(x)} row(s) cannot be split into n_folds={self.n_folds}: every fold must "
                "hold at least one row for the out-of-fold bound to be defined."
            )
        if a.min() < 0 or a.max() >= self.n_actions:
            raise ValueError(
                f"actions must lie in [0, {self.n_actions}); observed [{a.min()}, {a.max()}]"
            )

        n = len(x)
        index = np.arange(n)
        np.random.default_rng(self.seed).shuffle(index)
        folds = np.array_split(index, self.n_folds)

        self._fold_models = []
        mu_oof = np.zeros((n, self.n_actions))
        e_oof = np.zeros((n, self.n_actions))
        for test in folds:
            train = np.setdiff1d(index, test, assume_unique=True)
            outcome, propensity = self._fit_fold(x[train], a[train], r[train])
            self._fold_models.append((outcome, propensity))
            mu_oof[test], e_oof[test] = self._predict(x[test], [(outcome, propensity)])

        low, high = self.reward_range
        logged = mu_oof * e_oof
        unlogged = 1.0 - e_oof
        self._in_sample = (logged + low * unlogged, logged + high * unlogged)
        taken = e_oof[np.arange(n), a]
        self._diagnostic = OverlapDiagnostic(
            min_propensity=float(taken.min()),
            mean_propensity=float(taken.mean()),
            vacuous_fraction=float(np.mean(e_oof < _VACUOUS_PROPENSITY)),
        )
        return self

    def _fit_fold(
        self, x: FloatArray, a: NDArray[np.int_], r: FloatArray
    ) -> tuple[dict[int, Regressor], dict[int, Classifier]]:
        outcome: dict[int, Regressor] = {}
        propensity: dict[int, Classifier] = {}
        for action in range(self.n_actions):
            rows = a == action
            if not rows.any():
                continue  # never logged in this fold: e stays 0, so the bound stays vacuous
            outcome[action] = self._outcome().fit(x[rows], r[rows])
            if rows.all():
                continue  # only action present: normalisation below gives it probability 1
            propensity[action] = self._propensity().fit(x, rows.astype(np.float64))
        return outcome, propensity

    def _predict(
        self,
        x: FloatArray,
        models: list[tuple[dict[int, Regressor], dict[int, Classifier]]],
    ) -> tuple[FloatArray, FloatArray]:
        """Average ``mu`` and (one-vs-rest, normalised) ``e`` over the supplied fold models."""
        mu = np.zeros((len(x), self.n_actions))
        e = np.zeros((len(x), self.n_actions))
        for outcome, propensity in models:
            raw = np.zeros((len(x), self.n_actions))
            for action in range(self.n_actions):
                if action in outcome:
                    mu[:, action] += np.asarray(outcome[action].predict(x)).reshape(-1)
                if action in propensity:
                    raw[:, action] = np.clip(
                        np.asarray(propensity[action].predict_proba(x)).reshape(-1), 0.0, 1.0
                    )
                elif action in outcome:
                    raw[:, action] = 1.0  # sole logged action in this fold
            total = raw.sum(axis=1, keepdims=True)
            e += np.divide(raw, total, out=np.zeros_like(raw), where=total > 0.0)
        return mu / len(models), e / len(models)

    @property
    def in_sample(self) -> tuple[FloatArray, FloatArray]:
        """Out-of-fold ``(lower, upper)`` for the training rows, each ``(n, n_actions)``."""
        if self._in_sample is None:
            raise RuntimeError("FunctionalManskiBounds has not been fitted yet; call fit() first")
        return self._in_sample

    def bounds(self, features: FloatArray) -> tuple[FloatArray, FloatArray]:
        """``(lower, upper)`` at arbitrary feature rows, from the averaged fold models.

        Each is ``(n_rows, n_actions)``. Evaluating the training rows through this path is
        *not* the same as :attr:`in_sample`: the average includes the fold model that saw them, so
        the interval is optimistically tight there.
        """
        if not self._fold_models:
            raise RuntimeError("FunctionalManskiBounds has not been fitted yet; call fit() first")
        x = np.asarray(features, dtype=np.float64)
        if x.ndim != 2:
            raise ValueError(f"features must be a 2-D (n, dim) design matrix; got shape {x.shape}")
        mu, e = self._predict(x, self._fold_models)
        low, high = self.reward_range
        return mu * e + low * (1.0 - e), mu * e + high * (1.0 - e)

    def diagnostic(self) -> OverlapDiagnostic:
        """Overlap summary from the out-of-fold propensities — see :class:`OverlapDiagnostic`."""
        if self._diagnostic is None:
            raise RuntimeError("FunctionalManskiBounds has not been fitted yet; call fit() first")
        return self._diagnostic

in_sample property

Out-of-fold (lower, upper) for the training rows, each (n, n_actions).

fit(features, actions, rewards)

Cross-fit the outcome and propensity models and record the out-of-fold bounds.

Source code in src/causalrl/bounds/functional.py

def fit(
    self, features: FloatArray, actions: NDArray[np.int_], rewards: FloatArray
) -> FunctionalManskiBounds:
    """Cross-fit the outcome and propensity models and record the out-of-fold bounds."""
    x = np.asarray(features, dtype=np.float64)
    a = np.asarray(actions, dtype=int).reshape(-1)
    r = np.asarray(rewards, dtype=np.float64).reshape(-1)
    if x.ndim != 2:
        raise ValueError(f"features must be a 2-D (n, dim) design matrix; got shape {x.shape}")
    if not (len(x) == len(a) == len(r)):
        raise ValueError(
            f"features, actions and rewards must agree in length; got {len(x)}, {len(a)}, "
            f"{len(r)}"
        )
    if len(x) < self.n_folds:
        raise ValueError(
            f"{len(x)} row(s) cannot be split into n_folds={self.n_folds}: every fold must "
            "hold at least one row for the out-of-fold bound to be defined."
        )
    if a.min() < 0 or a.max() >= self.n_actions:
        raise ValueError(
            f"actions must lie in [0, {self.n_actions}); observed [{a.min()}, {a.max()}]"
        )

    n = len(x)
    index = np.arange(n)
    np.random.default_rng(self.seed).shuffle(index)
    folds = np.array_split(index, self.n_folds)

    self._fold_models = []
    mu_oof = np.zeros((n, self.n_actions))
    e_oof = np.zeros((n, self.n_actions))
    for test in folds:
        train = np.setdiff1d(index, test, assume_unique=True)
        outcome, propensity = self._fit_fold(x[train], a[train], r[train])
        self._fold_models.append((outcome, propensity))
        mu_oof[test], e_oof[test] = self._predict(x[test], [(outcome, propensity)])

    low, high = self.reward_range
    logged = mu_oof * e_oof
    unlogged = 1.0 - e_oof
    self._in_sample = (logged + low * unlogged, logged + high * unlogged)
    taken = e_oof[np.arange(n), a]
    self._diagnostic = OverlapDiagnostic(
        min_propensity=float(taken.min()),
        mean_propensity=float(taken.mean()),
        vacuous_fraction=float(np.mean(e_oof < _VACUOUS_PROPENSITY)),
    )
    return self

bounds(features)

(lower, upper) at arbitrary feature rows, from the averaged fold models.

Each is (n_rows, n_actions). Evaluating the training rows through this path is not the same as :attr:in_sample: the average includes the fold model that saw them, so the interval is optimistically tight there.

Source code in src/causalrl/bounds/functional.py

def bounds(self, features: FloatArray) -> tuple[FloatArray, FloatArray]:
    """``(lower, upper)`` at arbitrary feature rows, from the averaged fold models.

    Each is ``(n_rows, n_actions)``. Evaluating the training rows through this path is
    *not* the same as :attr:`in_sample`: the average includes the fold model that saw them, so
    the interval is optimistically tight there.
    """
    if not self._fold_models:
        raise RuntimeError("FunctionalManskiBounds has not been fitted yet; call fit() first")
    x = np.asarray(features, dtype=np.float64)
    if x.ndim != 2:
        raise ValueError(f"features must be a 2-D (n, dim) design matrix; got shape {x.shape}")
    mu, e = self._predict(x, self._fold_models)
    low, high = self.reward_range
    return mu * e + low * (1.0 - e), mu * e + high * (1.0 - e)

diagnostic()

Overlap summary from the out-of-fold propensities — see :class:OverlapDiagnostic.

Source code in src/causalrl/bounds/functional.py

def diagnostic(self) -> OverlapDiagnostic:
    """Overlap summary from the out-of-fold propensities — see :class:`OverlapDiagnostic`."""
    if self._diagnostic is None:
        raise RuntimeError("FunctionalManskiBounds has not been fitted yet; call fit() first")
    return self._diagnostic

causalrl.bounds.functional.OverlapDiagnostic dataclass

How much support the fitted propensity actually found, per action.

min_propensity and mean_propensity are taken over the estimated e(x, a) at the training rows for the action actually taken. vacuous_fraction is the share of (row, action) pairs whose estimated propensity falls below :data:_VACUOUS_PROPENSITY, i.e. where the interval is close to the whole reward range.

A high vacuous_fraction is not a failure — it is the bound correctly reporting that the logs do not pin the effect down. The case to be suspicious of is the opposite: a very high mean_propensity from a flexible model, which yields tight intervals that are only as trustworthy as the propensity fit.

Source code in src/causalrl/bounds/functional.py

@dataclass(frozen=True)
class OverlapDiagnostic:
    """How much support the fitted propensity actually found, per action.

    ``min_propensity`` and ``mean_propensity`` are taken over the estimated ``e(x, a)`` at the
    training rows for the action actually taken. ``vacuous_fraction`` is the share of ``(row,
    action)`` pairs whose estimated propensity falls below :data:`_VACUOUS_PROPENSITY`, i.e. where
    the interval is close to the whole reward range.

    A high ``vacuous_fraction`` is not a failure — it is the bound correctly reporting that the logs
    do not pin the effect down. The case to be suspicious of is the opposite: a very high
    ``mean_propensity`` from a flexible model, which yields tight intervals that are only as
    trustworthy as the propensity fit.
    """

    min_propensity: float
    mean_propensity: float
    vacuous_fraction: float

    def summary(self) -> str:
        return (
            f"overlap: min={self.min_propensity:.3f} mean={self.mean_propensity:.3f} "
            f"vacuous={self.vacuous_fraction:.1%}"
        )

causalrl.agents.bounded_fitted.BoundedFittedQIteration

Bases: BatchAgent

Finite-horizon envelope propagation over encoded features, under Manski reward bounds.

encoder decides what a state is; bounds_model supplies the immediate-reward interval at each feature point (default :class:~causalrl.bounds.functional.FunctionalManskiBounds with the dependency-free nuisances); regressor estimates the continuation expectation.

As with :class:~causalrl.agents.fitted.FittedQIteration, the tabular case is contained rather than replaced: with :class:~causalrl.state.OneHotEncoder the reward bounds reduce to :func:causalrl.causal_q_bounds per cell, so the envelope reduces to the interval DOVI's ceiling caps.

transition_assumption mirrors :class:causalrl.DOVI. At "unknown" with a horizon above 1, :meth:fit raises :class:~causalrl.exceptions.UnverifiedAssumptionError unless allow_heuristic=True — propagating an interval through a confounded successor distribution produces a number that is not a bound, and running it silently would be the one failure mode this class exists to prevent.

Source code in src/causalrl/agents/bounded_fitted.py

class BoundedFittedQIteration(BatchAgent):
    """Finite-horizon envelope propagation over encoded features, under Manski reward bounds.

    ``encoder`` decides what a state is; ``bounds_model`` supplies the immediate-reward interval at
    each feature point (default :class:`~causalrl.bounds.functional.FunctionalManskiBounds` with the
    dependency-free nuisances); ``regressor`` estimates the continuation expectation.

    As with :class:`~causalrl.agents.fitted.FittedQIteration`, the tabular case is contained rather
    than replaced: with :class:`~causalrl.state.OneHotEncoder` the reward bounds reduce to
    :func:`causalrl.causal_q_bounds` per cell, so the envelope reduces to the interval DOVI's
    ceiling caps.

    ``transition_assumption`` mirrors :class:`causalrl.DOVI`. At ``"unknown"`` with a horizon above
    1, :meth:`fit` raises :class:`~causalrl.exceptions.UnverifiedAssumptionError` unless
    ``allow_heuristic=True`` — propagating an interval through a confounded successor distribution
    produces a number that is not a bound, and running it silently would be the one failure mode
    this class exists to prevent.
    """

    def __init__(
        self,
        n_actions: int,
        horizon: int,
        encoder: StateEncoder,
        *,
        bounds_model: FunctionalManskiBounds | None = None,
        regressor: Callable[[], Regressor] | None = None,
        reward_range: tuple[float, float] = (0.0, 1.0),
        transition_assumption: TransitionAssumption = "unknown",
        allow_heuristic: bool = False,
        seed: int | None = None,
    ) -> None:
        if n_actions < 1:
            raise ValueError(f"n_actions={n_actions} must be at least 1")
        if horizon < 1:
            raise ValueError(f"horizon={horizon} must be at least 1")
        if transition_assumption not in ("unknown", "unconfounded"):
            raise ValueError(
                "transition_assumption must be 'unknown' or 'unconfounded', "
                f"got {transition_assumption!r}"
            )
        self.n_actions = n_actions
        self.horizon = horizon
        self.encoder = encoder
        self.reward_range = (float(reward_range[0]), float(reward_range[1]))
        self.transition_assumption = transition_assumption
        self.allow_heuristic = allow_heuristic
        self._bounds_model = (
            bounds_model
            if bounds_model is not None
            else FunctionalManskiBounds(n_actions, reward_range=self.reward_range)
        )
        self._make_regressor = regressor if regressor is not None else RidgeRegressor
        self._buffer: list[FeatureTransition] = []
        self._lower_cont: dict[tuple[int, int], Regressor] = {}
        self._upper_cont: dict[tuple[int, int], Regressor] = {}
        self._fitted = False
        self._rng = np.random.default_rng(seed)

    @property
    def is_certified(self) -> bool:
        """Whether the propagated interval is a bound rather than heuristic value propagation.

        True at horizon 1 (nothing is propagated, so only the reward bound is used) or when the
        caller asserts unconfounded transitions. Mirrors :attr:`causalrl.DOVI.is_certified`.

        Note this is a statement about the *transition* assumption only. Even when it is ``True``,
        the interval still rests on the outcome and propensity models being correctly specified —
        an assumption the tabular bound does not need, and one every certificate here records.
        """
        return self.horizon <= 1 or self.transition_assumption == "unconfounded"

    def fit(
        self, transitions: Sequence[FeatureTransition] | None = None
    ) -> BoundedFittedQIteration:
        """Fit the reward bounds and propagate the envelope back from the horizon.

        Passing ``transitions`` replaces the observed buffer, matching
        :meth:`causalrl.FittedQIteration.fit`; omitting it refits from whatever
        :meth:`observe_step` has accumulated.
        """
        if not self.is_certified and not self.allow_heuristic:
            raise UnverifiedAssumptionError(
                "multi-step envelope propagation requires transition_assumption='unconfounded': "
                "the continuation expectation is taken over the LOGGED successor distribution, so "
                "if the behaviour policy's hidden confounder also drives the dynamics the "
                "propagated interval bounds nothing. Set allow_heuristic=True to run it anyway as "
                "value propagation (the certificate then reports EMPIRICAL, not BOUNDED)."
            )
        if transitions is not None:
            self._buffer = list(transitions)
        if not self._buffer:
            raise ValueError(
                "BoundedFittedQIteration.fit needs transitions: pass them explicitly or feed "
                "them in through observe_step first. There is no envelope over an empty log."
            )
        batch = unpack_transitions(self._buffer, encoder=self.encoder, n_actions=self.n_actions)
        states, next_states = batch.states, batch.next_states
        actions, rewards, not_done = batch.actions, batch.rewards, batch.not_done

        self._bounds_model.fit(states, actions, rewards)
        reward_low_next, reward_high_next = self._bounds_model.bounds(next_states)

        self._lower_cont = {}
        self._upper_cont = {}
        # Continuation value of the step above, evaluated at each row's successor. Zero at H+1.
        lower_at_successor = np.zeros(len(self._buffer))
        upper_at_successor = np.zeros(len(self._buffer))
        for step in range(self.horizon, 0, -1):
            lower_target = not_done * lower_at_successor
            upper_target = not_done * upper_at_successor
            for action in range(self.n_actions):
                rows = actions == action
                if not rows.any():
                    continue  # no data: the continuation defaults to 0 for this (step, action)
                self._lower_cont[(step, action)] = self._make_regressor().fit(
                    states[rows], lower_target[rows]
                )
                self._upper_cont[(step, action)] = self._make_regressor().fit(
                    states[rows], upper_target[rows]
                )
            low, high = self._envelope(next_states, step, reward_low_next, reward_high_next)
            lower_at_successor = low.max(axis=1)
            upper_at_successor = high.max(axis=1)
        self._fitted = True
        return self

    def _envelope(
        self,
        features: FloatArray,
        step: int,
        reward_low: FloatArray,
        reward_high: FloatArray,
    ) -> tuple[FloatArray, FloatArray]:
        """``(lower, upper)`` envelopes at ``step``, each ``(n_rows, n_actions)``, range-clipped."""
        remaining = self.horizon - step + 1
        floor = self.reward_range[0] * remaining
        cap = self.reward_range[1] * remaining
        low = np.array(reward_low, dtype=np.float64, copy=True)
        high = np.array(reward_high, dtype=np.float64, copy=True)
        for action in range(self.n_actions):
            lower_model = self._lower_cont.get((step, action))
            if lower_model is not None:
                low[:, action] += np.asarray(lower_model.predict(features)).reshape(-1)
            upper_model = self._upper_cont.get((step, action))
            if upper_model is not None:
                high[:, action] += np.asarray(upper_model.predict(features)).reshape(-1)
        # A return over `remaining` steps cannot leave the scaled reward range, whatever the
        # regressions extrapolate to; and the envelope must not invert.
        low = np.clip(low, floor, cap)
        high = np.clip(high, floor, cap)
        return low, np.maximum(high, low)

    def envelope(
        self, observation: Mapping[str, Any], step: int = 1
    ) -> tuple[FloatArray, FloatArray]:
        """``(lower, upper)`` arrays over actions at ``observation`` and ``step``."""
        self._require_fit()
        if not 1 <= step <= self.horizon:
            raise ValueError(f"step={step} must lie in [1, {self.horizon}]")
        features = np.asarray(self.encoder.encode(observation), dtype=np.float64).reshape(1, -1)
        reward_low, reward_high = self._bounds_model.bounds(features)
        low, high = self._envelope(features, step, reward_low, reward_high)
        return low[0], high[0]

    def interval(self, observation: Mapping[str, Any], action: int, step: int = 1) -> Interval:
        """The bound on ``Q_step(observation, action)`` as an :class:`Interval`."""
        low, high = self.envelope(observation, step)
        if not 0 <= action < self.n_actions:
            raise ValueError(f"action={action} must lie in [0, {self.n_actions})")
        return Interval(float(low[action]), float(high[action]))

    def non_dominated(self, observation: Mapping[str, Any], step: int = 1) -> list[int]:
        """Actions whose upper bound reaches the best lower bound — the survivors.

        The feature-space counterpart of
        :func:`causalrl.agents.primitives.non_dominated_actions`. An action whose whole interval
        lies below another's floor cannot be optimal at this state and can be dropped before any
        exploration. As in the tabular case, natural Manski bounds are wide and this is often a
        no-op returning every action, which is correct rather than a defect.
        """
        low, high = self.envelope(observation, step)
        best_lower = float(low.max())
        return [a for a in range(self.n_actions) if float(high[a]) >= best_lower]

    def act(self, observation: dict[str, Any]) -> int:
        """Optimism in the face of confounding: the action with the highest upper bound."""
        step = min(int(observation.get("t", 0)) + 1, self.horizon)
        _, high = self.envelope(observation, step)
        best = float(high.max())
        winners = [a for a in range(self.n_actions) if float(high[a]) >= best - 1e-12]
        return int(self._rng.choice(winners))

    def observe_transition(self, state: int, action: int, next_state: int, done: bool) -> None:
        """Refused: states here are feature vectors, not indices."""
        raise NotImplementedError(
            "BoundedFittedQIteration takes feature-typed transitions; the inherited int-typed "
            "hook cannot express a continuous state. Encode both endpoints and pass "
            "FeatureTransition objects to fit()."
        )

    def observe_step(
        self,
        observation: dict[str, Any],
        action: int,
        reward: float,
        next_observation: dict[str, Any],
        done: bool,
    ) -> None:
        """Encode both endpoints and buffer the transition.

        See :meth:`causalrl.FittedQIteration.observe_step`.
        """
        self._buffer.append(
            FeatureTransition(
                state=self.encoder.encode(observation),
                action=action,
                reward=float(reward),
                next_state=self.encoder.encode(next_observation),
                done=done,
            )
        )
        self._fitted = False

    def buffered_transitions(self) -> tuple[FeatureTransition, ...]:
        """The transitions observed so far and not yet discarded by a ``fit(transitions=...)``.

        The online buffer is otherwise write-only: a caller driving this agent through
        :meth:`observe_step` had no way to see what it had actually collected, which makes
        "did my driver wire the hook up correctly" unanswerable without reaching into a private
        attribute.
        """
        return tuple(self._buffer)

    def _require_fit(self) -> None:
        if not self._fitted:
            raise RuntimeError(
                "BoundedFittedQIteration has no envelope yet; call fit() with transitions first."
            )

    def certificate(self) -> Certificate:
        """A ``BOUNDED`` certificate when the transition assumption holds, else ``EMPIRICAL``.

        Either way the assumptions are recorded explicitly, including the two the tabular path does
        not need — that the outcome regression and the propensity model are correctly specified.
        The overlap diagnostic rides along on the propensity assumption, since it is the one
        checkable thing about it: where estimated propensities are near zero the interval is wide
        and honest, and where a flexible model has driven them near one the interval is tight and
        only as good as that model.
        """
        self._require_fit()
        diagnostic = self._bounds_model.diagnostic()
        assumptions = (
            Assumption(
                name="reward-range",
                params={"low": self.reward_range[0], "high": self.reward_range[1]},
                checkable=False,
            ),
            Assumption(name="outcome-model-specification", checkable=False),
            Assumption(
                name="propensity-model-specification",
                checkable=True,
                diagnostic={
                    "min_propensity": diagnostic.min_propensity,
                    "mean_propensity": diagnostic.mean_propensity,
                    "vacuous_fraction": diagnostic.vacuous_fraction,
                },
            ),
            Assumption(
                name="unconfounded-transitions",
                params={"transition_assumption": self.transition_assumption},
                checkable=False,
            ),
        )
        claim = (
            f"Manski envelope on Q over {self.encoder.dim}-dimensional state features "
            f"(horizon={self.horizon}, actions={self.n_actions})"
        )
        if self.is_certified:
            return Certificate(
                claim=claim,
                estimand=EstimandSpec(query="policy_value", target="mean"),
                kind=Kind.BOUNDED,
                value=None,
                alpha=None,
                assumptions=assumptions,
                method="bounded_fitted_q_iteration",
                witness=None,
                hedge=Hedge(
                    reason=(
                        "the interval is a bound only up to correct specification of the outcome "
                        "and propensity models: unlike the tabular per-cell bound, whose inputs "
                        "are counts, these are fitted functions and their misspecification "
                        "narrows the interval"
                    ),
                    detail={"overlap": diagnostic.summary()},
                ),
                provenance=Provenance.create(),
            )
        return Certificate(
            claim=claim,
            estimand=EstimandSpec(query="policy_value", target="mean"),
            kind=Kind.EMPIRICAL,
            value=None,
            alpha=None,
            assumptions=assumptions,
            method="bounded_fitted_q_iteration",
            witness=None,
            hedge=Hedge(
                reason=(
                    "transition_assumption='unknown' with horizon > 1: the continuation "
                    "expectation is taken over the logged successor distribution, which a "
                    "confounder driving the dynamics would make non-interventional, so this is "
                    "value propagation rather than a bound"
                ),
                detail={"overlap": diagnostic.summary(), "allow_heuristic": self.allow_heuristic},
                downgraded_from="bounded",
            ),
            provenance=Provenance.create(),
        )

is_certified property

Whether the propagated interval is a bound rather than heuristic value propagation.

True at horizon 1 (nothing is propagated, so only the reward bound is used) or when the caller asserts unconfounded transitions. Mirrors :attr:causalrl.DOVI.is_certified.

Note this is a statement about the transition assumption only. Even when it is True, the interval still rests on the outcome and propensity models being correctly specified — an assumption the tabular bound does not need, and one every certificate here records.

fit(transitions=None)

Fit the reward bounds and propagate the envelope back from the horizon.

Passing transitions replaces the observed buffer, matching :meth:causalrl.FittedQIteration.fit; omitting it refits from whatever :meth:observe_step has accumulated.

Source code in src/causalrl/agents/bounded_fitted.py

def fit(
    self, transitions: Sequence[FeatureTransition] | None = None
) -> BoundedFittedQIteration:
    """Fit the reward bounds and propagate the envelope back from the horizon.

    Passing ``transitions`` replaces the observed buffer, matching
    :meth:`causalrl.FittedQIteration.fit`; omitting it refits from whatever
    :meth:`observe_step` has accumulated.
    """
    if not self.is_certified and not self.allow_heuristic:
        raise UnverifiedAssumptionError(
            "multi-step envelope propagation requires transition_assumption='unconfounded': "
            "the continuation expectation is taken over the LOGGED successor distribution, so "
            "if the behaviour policy's hidden confounder also drives the dynamics the "
            "propagated interval bounds nothing. Set allow_heuristic=True to run it anyway as "
            "value propagation (the certificate then reports EMPIRICAL, not BOUNDED)."
        )
    if transitions is not None:
        self._buffer = list(transitions)
    if not self._buffer:
        raise ValueError(
            "BoundedFittedQIteration.fit needs transitions: pass them explicitly or feed "
            "them in through observe_step first. There is no envelope over an empty log."
        )
    batch = unpack_transitions(self._buffer, encoder=self.encoder, n_actions=self.n_actions)
    states, next_states = batch.states, batch.next_states
    actions, rewards, not_done = batch.actions, batch.rewards, batch.not_done

    self._bounds_model.fit(states, actions, rewards)
    reward_low_next, reward_high_next = self._bounds_model.bounds(next_states)

    self._lower_cont = {}
    self._upper_cont = {}
    # Continuation value of the step above, evaluated at each row's successor. Zero at H+1.
    lower_at_successor = np.zeros(len(self._buffer))
    upper_at_successor = np.zeros(len(self._buffer))
    for step in range(self.horizon, 0, -1):
        lower_target = not_done * lower_at_successor
        upper_target = not_done * upper_at_successor
        for action in range(self.n_actions):
            rows = actions == action
            if not rows.any():
                continue  # no data: the continuation defaults to 0 for this (step, action)
            self._lower_cont[(step, action)] = self._make_regressor().fit(
                states[rows], lower_target[rows]
            )
            self._upper_cont[(step, action)] = self._make_regressor().fit(
                states[rows], upper_target[rows]
            )
        low, high = self._envelope(next_states, step, reward_low_next, reward_high_next)
        lower_at_successor = low.max(axis=1)
        upper_at_successor = high.max(axis=1)
    self._fitted = True
    return self

envelope(observation, step=1)

(lower, upper) arrays over actions at observation and step.

Source code in src/causalrl/agents/bounded_fitted.py

def envelope(
    self, observation: Mapping[str, Any], step: int = 1
) -> tuple[FloatArray, FloatArray]:
    """``(lower, upper)`` arrays over actions at ``observation`` and ``step``."""
    self._require_fit()
    if not 1 <= step <= self.horizon:
        raise ValueError(f"step={step} must lie in [1, {self.horizon}]")
    features = np.asarray(self.encoder.encode(observation), dtype=np.float64).reshape(1, -1)
    reward_low, reward_high = self._bounds_model.bounds(features)
    low, high = self._envelope(features, step, reward_low, reward_high)
    return low[0], high[0]

interval(observation, action, step=1)

The bound on Q_step(observation, action) as an :class:Interval.

Source code in src/causalrl/agents/bounded_fitted.py

def interval(self, observation: Mapping[str, Any], action: int, step: int = 1) -> Interval:
    """The bound on ``Q_step(observation, action)`` as an :class:`Interval`."""
    low, high = self.envelope(observation, step)
    if not 0 <= action < self.n_actions:
        raise ValueError(f"action={action} must lie in [0, {self.n_actions})")
    return Interval(float(low[action]), float(high[action]))

non_dominated(observation, step=1)

Actions whose upper bound reaches the best lower bound — the survivors.

The feature-space counterpart of :func:causalrl.agents.primitives.non_dominated_actions. An action whose whole interval lies below another's floor cannot be optimal at this state and can be dropped before any exploration. As in the tabular case, natural Manski bounds are wide and this is often a no-op returning every action, which is correct rather than a defect.

Source code in src/causalrl/agents/bounded_fitted.py

def non_dominated(self, observation: Mapping[str, Any], step: int = 1) -> list[int]:
    """Actions whose upper bound reaches the best lower bound — the survivors.

    The feature-space counterpart of
    :func:`causalrl.agents.primitives.non_dominated_actions`. An action whose whole interval
    lies below another's floor cannot be optimal at this state and can be dropped before any
    exploration. As in the tabular case, natural Manski bounds are wide and this is often a
    no-op returning every action, which is correct rather than a defect.
    """
    low, high = self.envelope(observation, step)
    best_lower = float(low.max())
    return [a for a in range(self.n_actions) if float(high[a]) >= best_lower]

act(observation)

Optimism in the face of confounding: the action with the highest upper bound.

Source code in src/causalrl/agents/bounded_fitted.py

def act(self, observation: dict[str, Any]) -> int:
    """Optimism in the face of confounding: the action with the highest upper bound."""
    step = min(int(observation.get("t", 0)) + 1, self.horizon)
    _, high = self.envelope(observation, step)
    best = float(high.max())
    winners = [a for a in range(self.n_actions) if float(high[a]) >= best - 1e-12]
    return int(self._rng.choice(winners))

observe_transition(state, action, next_state, done)

Refused: states here are feature vectors, not indices.

Source code in src/causalrl/agents/bounded_fitted.py

def observe_transition(self, state: int, action: int, next_state: int, done: bool) -> None:
    """Refused: states here are feature vectors, not indices."""
    raise NotImplementedError(
        "BoundedFittedQIteration takes feature-typed transitions; the inherited int-typed "
        "hook cannot express a continuous state. Encode both endpoints and pass "
        "FeatureTransition objects to fit()."
    )

observe_step(observation, action, reward, next_observation, done)

Encode both endpoints and buffer the transition.

See :meth:causalrl.FittedQIteration.observe_step.

Source code in src/causalrl/agents/bounded_fitted.py

def observe_step(
    self,
    observation: dict[str, Any],
    action: int,
    reward: float,
    next_observation: dict[str, Any],
    done: bool,
) -> None:
    """Encode both endpoints and buffer the transition.

    See :meth:`causalrl.FittedQIteration.observe_step`.
    """
    self._buffer.append(
        FeatureTransition(
            state=self.encoder.encode(observation),
            action=action,
            reward=float(reward),
            next_state=self.encoder.encode(next_observation),
            done=done,
        )
    )
    self._fitted = False

buffered_transitions()

The transitions observed so far and not yet discarded by a fit(transitions=...).

The online buffer is otherwise write-only: a caller driving this agent through :meth:observe_step had no way to see what it had actually collected, which makes "did my driver wire the hook up correctly" unanswerable without reaching into a private attribute.

Source code in src/causalrl/agents/bounded_fitted.py

def buffered_transitions(self) -> tuple[FeatureTransition, ...]:
    """The transitions observed so far and not yet discarded by a ``fit(transitions=...)``.

    The online buffer is otherwise write-only: a caller driving this agent through
    :meth:`observe_step` had no way to see what it had actually collected, which makes
    "did my driver wire the hook up correctly" unanswerable without reaching into a private
    attribute.
    """
    return tuple(self._buffer)

certificate()

A BOUNDED certificate when the transition assumption holds, else EMPIRICAL.

Either way the assumptions are recorded explicitly, including the two the tabular path does not need — that the outcome regression and the propensity model are correctly specified. The overlap diagnostic rides along on the propensity assumption, since it is the one checkable thing about it: where estimated propensities are near zero the interval is wide and honest, and where a flexible model has driven them near one the interval is tight and only as good as that model.

Source code in src/causalrl/agents/bounded_fitted.py

def certificate(self) -> Certificate:
    """A ``BOUNDED`` certificate when the transition assumption holds, else ``EMPIRICAL``.

    Either way the assumptions are recorded explicitly, including the two the tabular path does
    not need — that the outcome regression and the propensity model are correctly specified.
    The overlap diagnostic rides along on the propensity assumption, since it is the one
    checkable thing about it: where estimated propensities are near zero the interval is wide
    and honest, and where a flexible model has driven them near one the interval is tight and
    only as good as that model.
    """
    self._require_fit()
    diagnostic = self._bounds_model.diagnostic()
    assumptions = (
        Assumption(
            name="reward-range",
            params={"low": self.reward_range[0], "high": self.reward_range[1]},
            checkable=False,
        ),
        Assumption(name="outcome-model-specification", checkable=False),
        Assumption(
            name="propensity-model-specification",
            checkable=True,
            diagnostic={
                "min_propensity": diagnostic.min_propensity,
                "mean_propensity": diagnostic.mean_propensity,
                "vacuous_fraction": diagnostic.vacuous_fraction,
            },
        ),
        Assumption(
            name="unconfounded-transitions",
            params={"transition_assumption": self.transition_assumption},
            checkable=False,
        ),
    )
    claim = (
        f"Manski envelope on Q over {self.encoder.dim}-dimensional state features "
        f"(horizon={self.horizon}, actions={self.n_actions})"
    )
    if self.is_certified:
        return Certificate(
            claim=claim,
            estimand=EstimandSpec(query="policy_value", target="mean"),
            kind=Kind.BOUNDED,
            value=None,
            alpha=None,
            assumptions=assumptions,
            method="bounded_fitted_q_iteration",
            witness=None,
            hedge=Hedge(
                reason=(
                    "the interval is a bound only up to correct specification of the outcome "
                    "and propensity models: unlike the tabular per-cell bound, whose inputs "
                    "are counts, these are fitted functions and their misspecification "
                    "narrows the interval"
                ),
                detail={"overlap": diagnostic.summary()},
            ),
            provenance=Provenance.create(),
        )
    return Certificate(
        claim=claim,
        estimand=EstimandSpec(query="policy_value", target="mean"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=None,
        assumptions=assumptions,
        method="bounded_fitted_q_iteration",
        witness=None,
        hedge=Hedge(
            reason=(
                "transition_assumption='unknown' with horizon > 1: the continuation "
                "expectation is taken over the logged successor distribution, which a "
                "confounder driving the dynamics would make non-interventional, so this is "
                "value propagation rather than a bound"
            ),
            detail={"overlap": diagnostic.summary(), "allow_heuristic": self.allow_heuristic},
            downgraded_from="bounded",
        ),
        provenance=Provenance.create(),
    )

Core — Complete Reference

causalrl.agents.causal_mbrl.CausalMBRLAgent

Route an offline confounded decision problem to the right causal planner, behind fit/act.

Routing (first match wins):

============================ ============================================================ You provide Strategy (planner) ============================ ============================================================ horizon=... "sequential" — deconfounded value iteration (DOVI) transport=(...) "transport" — deconfound + transport by the target P continuous_confounder "function_approx" — ridge-RBF back-door on a continuous Z covariates=[...] "g_formula" — multivariate standardization (real datasets) graph=None "discovery" — skeleton discovery + temporal tiers, adjust otherwise (graph given) "backdoor" — back-door adjust an observed confounder ============================ ============================================================

Then fit on offline data and act on an observation. explain() reports the chosen strategy and the identification basis (adjustment set / transportability / confounder) — the trust artifact for the identifiable regime this agent targets. It beats correlational offline RL only there (confounded / offline / transfer); it is not a clean-benchmark performer.

Source code in src/causalrl/agents/causal_mbrl.py

class CausalMBRLAgent:
    """Route an offline confounded decision problem to the right causal planner, behind fit/act.

    Routing (first match wins):

    ============================  ============================================================
    You provide                    Strategy (planner)
    ============================  ============================================================
    ``horizon=...``                ``"sequential"`` — deconfounded value iteration (``DOVI``)
    ``transport=(...)``            ``"transport"`` — deconfound + transport by the target ``P``
    ``continuous_confounder``      ``"function_approx"`` — ridge-RBF back-door on a continuous Z
    ``covariates=[...]``           ``"g_formula"`` — multivariate standardization (real datasets)
    ``graph=None``                 ``"discovery"`` — skeleton discovery + temporal tiers, adjust
    otherwise (``graph`` given)    ``"backdoor"`` — back-door adjust an observed confounder
    ============================  ============================================================

    Then ``fit`` on offline data and ``act`` on an observation. ``explain()`` reports the chosen
    strategy and the identification basis (adjustment set / transportability / confounder) — the
    trust artifact for the identifiable regime this agent targets. It beats correlational offline
    RL only there (confounded / offline / transfer); it is not a clean-benchmark performer.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        graph: CausalGraph | None = None,
        variables: Sequence[str] | None = None,
        tiers: Sequence[Sequence[str]] | None = None,
        treatment: str = "A",
        outcome: str = "Y",
        transport: Sequence[str] | None = None,
        continuous_confounder: bool = False,
        covariates: Sequence[str] | None = None,
        outcome_model: Callable[[], Any] | None = None,
        horizon: int | None = None,
        n_states: int | None = None,
        seed: int = 0,
    ) -> None:
        self.n_actions = n_actions
        self._tiers = tiers
        self._planner: Any
        if horizon is not None:
            if n_states is None:
                raise ValueError("sequential mode (horizon set) requires n_states")
            self.strategy = "sequential"
            from causalrl.agents.dovi import DOVI  # lazy: keep the front door torch-light otherwise

            self._planner = DOVI(
                n_states=n_states,
                n_actions=n_actions,
                horizon=horizon,
                seed=seed,
                transition_assumption="unconfounded",
            )
        elif transport is not None:
            if graph is None:
                raise ValueError("transport mode requires a graph")
            self.strategy = "transport"
            self._planner = TransportBackdoorAgent(
                n_actions, graph=graph, treatment=treatment, outcome=outcome, transport=transport
            )
        elif continuous_confounder:
            if graph is None:
                raise ValueError("continuous_confounder mode requires a graph")
            self.strategy = "function_approx"
            self._planner = FunctionApproxBackdoorAgent(
                n_actions, graph=graph, treatment=treatment, outcome=outcome
            )
        elif covariates is not None:
            self.strategy = "g_formula"
            self._planner = GFormulaBackdoorAgent(
                n_actions,
                covariates=covariates,
                treatment=treatment,
                outcome=outcome,
                outcome_model=outcome_model,
            )
        elif graph is None:
            if variables is None or tiers is None:
                raise ValueError("discovery mode (no graph) requires variables and tiers")
            self.strategy = "discovery"
            self._planner = DiscoveryBackdoorAgent(
                n_actions, variables=variables, treatment=treatment, outcome=outcome
            )
        else:
            self.strategy = "backdoor"
            self._planner = BackdoorAdjustedAgent(
                n_actions, graph=graph, treatment=treatment, outcome=outcome
            )

    def fit(
        self,
        data: Any,
        *,
        target_covariates: Mapping[str, np.ndarray] | None = None,
    ) -> CausalMBRLAgent:
        """Fit the routed planner on offline ``data`` and return ``self``.

        ``data`` is columnar ``{treatment, outcome, *covariates}`` for every strategy except
        ``"sequential"``, which takes a ``ConfoundedTrajectoryDataset`` of trajectory logs.
        ``target_covariates`` (target-domain draws of the transport variables) is required for the
        ``"transport"`` strategy.
        """
        if self.strategy == "sequential":
            self._planner.ingest_offline(data)
        elif self.strategy == "transport":
            if target_covariates is None:
                raise ValueError("transport strategy requires target_covariates")
            self._planner.fit(data, target_covariates=target_covariates)
        elif self.strategy == "discovery":
            if self._tiers is None:
                raise ValueError("discovery strategy requires tiers")
            self._planner.discover_and_fit(data, tiers=self._tiers)
        else:
            self._planner.fit(data)
        return self

    def act(self, observation: dict[str, Any]) -> int:
        """Return the planner's action for ``observation``.

        Supply the variables the planner conditions on **by name** — the adjustment set for
        ``"backdoor"``/``"discovery"``, those plus the transport variables for ``"transport"``, the
        continuous confounder for ``"function_approx"``, the covariates for ``"g_formula"`` (e.g.
        ``{"Z": 1}`` or ``{"age": 42, "smoke": 1}``) — and the decision is contextual. The
        ``"sequential"`` planner takes ``{"state": s}`` instead. An observation carrying none of
        them falls back to the marginal ``argmax_a E[Y | do(a)]``, which is the right answer only
        when there is genuinely no context to condition on.
        """
        return int(self._planner.act(observation))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        """The fitted outcome model is fixed after ``fit``; the policy varies with the observation,
        not with online experience."""

    @property
    def planner(self) -> Any:
        """The underlying specialized planner, for full control or inspection."""
        return self._planner

    @property
    def adjustment_set(self) -> tuple[str, ...] | None:
        """The back-door adjustment set the planner used, when the strategy exposes one."""
        adjustment = getattr(self._planner, "adjustment", None)
        return tuple(adjustment) if adjustment is not None else None

    def explain(self) -> str:
        """Human-readable provenance: the strategy and the identification basis it relied on."""
        parts = [f"strategy={self.strategy}"]
        adjustment = self.adjustment_set
        if adjustment is not None:
            parts.append("adjustment_set={" + ", ".join(adjustment) + "}")
        if self.strategy == "transport":
            parts.append(f"transportable={getattr(self._planner, 'transportable', None)}")
        if self.strategy == "function_approx":
            parts.append(f"confounder={getattr(self._planner, 'confounder', None)}")
        if self.strategy == "g_formula":
            parts.append(f"covariates={list(getattr(self._planner, 'covariates', ()))}")
        return "CausalMBRLAgent(" + ", ".join(parts) + ")"

planner property

The underlying specialized planner, for full control or inspection.

adjustment_set property

The back-door adjustment set the planner used, when the strategy exposes one.

fit(data, *, target_covariates=None)

Fit the routed planner on offline data and return self.

data is columnar {treatment, outcome, *covariates} for every strategy except "sequential", which takes a ConfoundedTrajectoryDataset of trajectory logs. target_covariates (target-domain draws of the transport variables) is required for the "transport" strategy.

Source code in src/causalrl/agents/causal_mbrl.py

def fit(
    self,
    data: Any,
    *,
    target_covariates: Mapping[str, np.ndarray] | None = None,
) -> CausalMBRLAgent:
    """Fit the routed planner on offline ``data`` and return ``self``.

    ``data`` is columnar ``{treatment, outcome, *covariates}`` for every strategy except
    ``"sequential"``, which takes a ``ConfoundedTrajectoryDataset`` of trajectory logs.
    ``target_covariates`` (target-domain draws of the transport variables) is required for the
    ``"transport"`` strategy.
    """
    if self.strategy == "sequential":
        self._planner.ingest_offline(data)
    elif self.strategy == "transport":
        if target_covariates is None:
            raise ValueError("transport strategy requires target_covariates")
        self._planner.fit(data, target_covariates=target_covariates)
    elif self.strategy == "discovery":
        if self._tiers is None:
            raise ValueError("discovery strategy requires tiers")
        self._planner.discover_and_fit(data, tiers=self._tiers)
    else:
        self._planner.fit(data)
    return self

act(observation)

Return the planner's action for observation.

Supply the variables the planner conditions on by name — the adjustment set for "backdoor"/"discovery", those plus the transport variables for "transport", the continuous confounder for "function_approx", the covariates for "g_formula" (e.g. {"Z": 1} or {"age": 42, "smoke": 1}) — and the decision is contextual. The "sequential" planner takes {"state": s} instead. An observation carrying none of them falls back to the marginal argmax_a E[Y | do(a)], which is the right answer only when there is genuinely no context to condition on.

Source code in src/causalrl/agents/causal_mbrl.py

def act(self, observation: dict[str, Any]) -> int:
    """Return the planner's action for ``observation``.

    Supply the variables the planner conditions on **by name** — the adjustment set for
    ``"backdoor"``/``"discovery"``, those plus the transport variables for ``"transport"``, the
    continuous confounder for ``"function_approx"``, the covariates for ``"g_formula"`` (e.g.
    ``{"Z": 1}`` or ``{"age": 42, "smoke": 1}``) — and the decision is contextual. The
    ``"sequential"`` planner takes ``{"state": s}`` instead. An observation carrying none of
    them falls back to the marginal ``argmax_a E[Y | do(a)]``, which is the right answer only
    when there is genuinely no context to condition on.
    """
    return int(self._planner.act(observation))

update(observation, action, reward)

The fitted outcome model is fixed after fit; the policy varies with the observation, not with online experience.

Source code in src/causalrl/agents/causal_mbrl.py

def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
    """The fitted outcome model is fixed after ``fit``; the policy varies with the observation,
    not with online experience."""

explain()

Human-readable provenance: the strategy and the identification basis it relied on.

Source code in src/causalrl/agents/causal_mbrl.py

def explain(self) -> str:
    """Human-readable provenance: the strategy and the identification basis it relied on."""
    parts = [f"strategy={self.strategy}"]
    adjustment = self.adjustment_set
    if adjustment is not None:
        parts.append("adjustment_set={" + ", ".join(adjustment) + "}")
    if self.strategy == "transport":
        parts.append(f"transportable={getattr(self._planner, 'transportable', None)}")
    if self.strategy == "function_approx":
        parts.append(f"confounder={getattr(self._planner, 'confounder', None)}")
    if self.strategy == "g_formula":
        parts.append(f"covariates={list(getattr(self._planner, 'covariates', ()))}")
    return "CausalMBRLAgent(" + ", ".join(parts) + ")"

causalrl.agents.bandits.CausalThompsonSampling

Bases: Agent

Thompson sampling with one Beta posterior per (intuition, arm) cell.

Conditioning on the observed confounder proxy intuition de-confounds the choice, letting the agent learn the lucky arm for each intuition value.

Sampling uses a per-instance numpy.random.Generator seeded from seed (see :class:NaiveThompsonSampling for the reproducibility contract).

Source code in src/causalrl/agents/bandits.py

class CausalThompsonSampling(Agent):
    """Thompson sampling with one Beta posterior per (intuition, arm) cell.

    Conditioning on the observed confounder proxy `intuition` de-confounds the choice,
    letting the agent learn the lucky arm for each intuition value.

    Sampling uses a per-instance ``numpy.random.Generator`` seeded from ``seed`` (see
    :class:`NaiveThompsonSampling` for the reproducibility contract).
    """

    def __init__(self, n_arms: int, n_contexts: int, seed: int | None = None) -> None:
        self.n_arms = n_arms
        self.n_contexts = n_contexts
        self._alpha = np.ones((n_contexts, n_arms))
        self._beta = np.ones((n_contexts, n_arms))
        self._rng = np.random.default_rng(seed)

    def act(self, observation: dict[str, Any]) -> int:
        ctx = int(observation["intuition"])
        samples = np.asarray(self._rng.beta(self._alpha[ctx], self._beta[ctx]), dtype=np.float64)
        return int(np.argmax(samples))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        ctx = int(observation["intuition"])
        if reward > 0:
            self._alpha[ctx, action] += 1.0
        else:
            self._beta[ctx, action] += 1.0

causalrl.scm.fit.fit_scm(data, *, graph, families=None, holdout=0.2, seed=0)

Fit every mechanism of graph from data; return a learned SCM.

Each node is fitted by families[node] when given, else by dtype: integer-valued columns with at most 20 levels get :class:TabularCPT, everything else :class:ANMFit. The result is an ordinary :class:StructuralCausalModel marked provenance="fitted".

Mechanisms may be pinned instead of learned. Passing :class:~causalrl.scm.fitters.PinnedMechanism as a node's family deploys the equation you supply at that node while the rest of the graph is still fitted from data — the ordinary case where a documented rule sits next to behaviour that has no closed form. A model with both kinds of node carries provenance="mixed"; one whose every node is pinned carries provenance="specified", since nothing about its equations came from the data. Pinned nodes are listed by :attr:FitReport.pinned_nodes and still receive a holdout_score, which for them tests the asserted equation rather than measuring a fit.

The deployed mechanisms are fitted on a fraction of the data. data is split once, by a seed-controlled permutation shared by every node: 1 - holdout of the rows train the mechanisms the returned SCM actually carries, and the remaining holdout fraction is scored but never trained on. Nothing is refitted on the full sample afterwards, so at the default holdout=0.2 every number this model reports — every do(), every contrast — comes from an 80%-data fit. Pass a smaller holdout to trade the out-of-sample score for training rows.

holdout must lie in (0, 1). 0.0 is refused rather than silently reporting an in-sample number under :attr:NodeFit.holdout_score, whose contract is data the fit never saw; >= 1.0 would leave a one-row training set.

Raises :class:NotIdentifiableError on a graph with bidirected edges: under latent confounding a node's mechanism is not recoverable by regression on its observed parents, so this function has no single SCM to return. :func:causalrl.fit_scm_bounded takes that graph instead, fitting the nodes whose own parents are unconfounded and bounding the rest.

Source code in src/causalrl/scm/fit.py

def fit_scm(
    data: Mapping[str, np.ndarray],
    *,
    graph: CausalGraph,
    families: Mapping[str, MechanismFitter] | None = None,
    holdout: float = 0.2,
    seed: int = 0,
) -> StructuralCausalModel:
    """Fit every mechanism of ``graph`` from ``data``; return a learned SCM.

    Each node is fitted by ``families[node]`` when given, else by dtype: integer-valued columns
    with at most 20 levels get :class:`TabularCPT`, everything else :class:`ANMFit`. The result is
    an ordinary :class:`StructuralCausalModel` marked ``provenance="fitted"``.

    **Mechanisms may be pinned instead of learned.** Passing
    :class:`~causalrl.scm.fitters.PinnedMechanism` as a node's family deploys the equation you
    supply at that node while the rest of the graph is still fitted from ``data`` — the ordinary
    case where a documented rule sits next to behaviour that has no closed form. A model with both
    kinds of node carries ``provenance="mixed"``; one whose every node is pinned carries
    ``provenance="specified"``, since nothing about its equations came from the data. Pinned nodes
    are listed by :attr:`FitReport.pinned_nodes` and still receive a ``holdout_score``, which for
    them tests the asserted equation rather than measuring a fit.

    **The deployed mechanisms are fitted on a fraction of the data.** ``data`` is split once, by a
    ``seed``-controlled permutation shared by every node: ``1 - holdout`` of the rows train the
    mechanisms the returned SCM actually carries, and the remaining ``holdout`` fraction is scored
    but never trained on. Nothing is refitted on the full sample afterwards, so at the default
    ``holdout=0.2`` every number this model reports — every ``do()``, every contrast — comes from
    an 80%-data fit. Pass a smaller ``holdout`` to trade the out-of-sample score for training rows.

    ``holdout`` must lie in ``(0, 1)``. ``0.0`` is refused rather than silently reporting an
    in-sample number under :attr:`NodeFit.holdout_score`, whose contract is data the fit never
    saw; ``>= 1.0`` would leave a one-row training set.

    Raises :class:`NotIdentifiableError` on a graph with bidirected edges: under latent confounding
    a node's mechanism is not recoverable by regression on its observed parents, so this function
    has no single SCM to return. :func:`causalrl.fit_scm_bounded` takes that graph instead, fitting
    the nodes whose own parents are unconfounded and bounding the rest.
    """
    if not 0.0 < holdout < 1.0:
        raise ValueError(
            f"holdout={holdout} must lie in (0, 1): it is the fraction of rows held out of every "
            "mechanism's fit and scored as NodeFit.holdout_score. At 0.0 there is no such data, "
            "so the reported score would be an in-sample number under an out-of-sample name; at "
            "1.0 or above there is nothing left to fit on."
        )
    if graph.has_bidirected_edges():
        raise NotIdentifiableError(
            "fit_scm cannot fit a graph with bidirected edges: under latent confounding a "
            "mechanism is not identified by regression on observed parents, so there is no single "
            "SCM to return. Use causalrl.fit_scm_bounded, which fits every node whose own parents "
            "are unconfounded and reports a Manski interval at the rest.",
            witness=graph.bidirected_edges,
        )
    missing = [node for node in graph.nodes if node not in data]
    if missing:
        raise KeyError(f"data is missing column(s) for graph node(s): {sorted(missing)}")

    columns = {node: np.asarray(data[node]) for node in graph.nodes}
    n = len(next(iter(columns.values())))
    rng = np.random.default_rng(seed)
    permutation = rng.permutation(n)
    cut = max(1, int(n * (1.0 - holdout)))
    train_index, test_index = permutation[:cut], permutation[cut:]

    overrides = families or {}
    mechanisms: dict[str, Mechanism] = {}
    exogenous: dict[str, Distribution] = {}
    fits: list[NodeFit] = []
    for node in graph.topological_order():
        parents = tuple(sorted(graph.parents(node)))
        fitter = overrides.get(node) or (TabularCPT() if _is_discrete(columns[node]) else ANMFit())
        train_parents = {p: columns[p][train_index] for p in parents}
        fitted = fitter.fit(train_parents, columns[node][train_index])
        test_parents = {p: columns[p][test_index] for p in parents}
        holdout_score = (
            evaluate_holdout(fitted, test_parents, columns[node][test_index])
            # A validated holdout in (0, 1) leaves the test partition empty only at n == 1, where
            # the in-sample score is the only number there is.
            if len(test_index) > 0
            else fitted.score
        )
        mechanisms[node] = fitted.mechanism
        exogenous[node] = fitted.noise
        fits.append(
            NodeFit(
                node=node,
                family=_family_name(fitter),
                parents=parents,
                holdout_score=holdout_score,
                invertible=fitted.invertible,
                pinned=isinstance(fitter, PinnedMechanism),
            )
        )
        fitted.mechanism.invertible = fitted.invertible  # type: ignore[attr-defined]

    return StructuralCausalModel(
        graph,
        mechanisms,
        exogenous,
        provenance=_provenance(fits),
        fit_report=FitReport(nodes=tuple(fits), n_samples=n),
    )

Identification, Transport & Discovery — Complete Reference

causalrl.identification.counterfactual_bounds.CounterfactualBound

Bases: NamedTuple

Bounds on a counterfactual expectation, with whether they are tight.

A separate type from :class:~causalrl.identification.bounds.Interval, whose documented tuple-unpacking contract (lo, hi = interval) a third field would break. tight is False only for bounds this library knows it cannot close; phase-1 counterfactual_interval refuses such queries rather than returning a loose answer, so it always reports True.

Source code in src/causalrl/identification/counterfactual_bounds.py

class CounterfactualBound(NamedTuple):
    """Bounds on a counterfactual expectation, with whether they are tight.

    A separate type from :class:`~causalrl.identification.bounds.Interval`, whose documented
    tuple-unpacking contract (``lo, hi = interval``) a third field would break. ``tight`` is
    ``False`` only for bounds this library knows it cannot close; phase-1
    ``counterfactual_interval`` refuses such queries rather than returning a loose answer, so it
    always reports ``True``.
    """

    lower: float
    upper: float
    tight: bool

    @property
    def interval(self) -> Interval:
        """The bound as a plain :class:`Interval`, for the partial-identification surface."""
        return Interval(self.lower, self.upper)

interval property

The bound as a plain :class:Interval, for the partial-identification surface.

causalrl.identification.bounds.Interval

Bases: NamedTuple

A partial-identification interval [lower, upper].

Tuple-compatible: lo, hi = interval and interval[0] work, so existing tuple-consuming callers are unaffected; .lower / .upper read more clearly.

Source code in src/causalrl/identification/bounds.py

class Interval(NamedTuple):
    """A partial-identification interval ``[lower, upper]``.

    Tuple-compatible: ``lo, hi = interval`` and ``interval[0]`` work, so existing
    tuple-consuming callers are unaffected; ``.lower`` / ``.upper`` read more clearly.
    """

    lower: float
    upper: float

causalrl.identification.bounds.PivotalityCertificate

Bases: NamedTuple

Result of :func:pivotality_certificate (one-sided sign-robustness certificate).

Source code in src/causalrl/identification/bounds.py

class PivotalityCertificate(NamedTuple):
    """Result of :func:`pivotality_certificate` (one-sided sign-robustness certificate)."""

    certified: bool
    naive: float
    bias_bound: float
    mi_flip: float
    mi_measured: float | None

causalrl.identification.estimate.PolicyValueContrast dataclass

An off-policy value contrast V(pi_on) - V(pi_off) in the terms certify_estimate needs.

Supply MSM evidence (logging_propensities + target_on + target_off) and/or structural evidence (treated + confounder_bins or mi_cap); at least one layer must be runnable. See :func:causalrl.certify_estimate.

Source code in src/causalrl/identification/estimate.py

@dataclass(frozen=True)
class PolicyValueContrast:
    """An off-policy value contrast ``V(pi_on) - V(pi_off)`` in the terms certify_estimate needs.

    Supply MSM evidence (``logging_propensities`` + ``target_on`` + ``target_off``) and/or
    structural evidence (``treated`` + ``confounder_bins`` or ``mi_cap``); at least one layer
    must be runnable. See :func:`causalrl.certify_estimate`.
    """

    outcomes: Sequence[float]
    logging_propensities: Sequence[float] | None = None
    target_on: Sequence[float] | None = None
    target_off: Sequence[float] | None = None
    treated: Sequence[int] | None = None
    confounder_bins: Sequence[int] | None = None
    mi_cap: float | None = None

    def __post_init__(self) -> None:
        y = np.asarray(self.outcomes, dtype=float)
        if y.size == 0:
            raise ValueError("outcomes must be non-empty")
        n = int(y.size)
        if self.logging_propensities is not None:
            if self.target_on is None or self.target_off is None:
                raise ValueError("logging_propensities requires target_on and target_off")
            e0 = np.asarray(self.logging_propensities, dtype=float)
            on = np.asarray(self.target_on, dtype=float)
            off = np.asarray(self.target_off, dtype=float)
            if not (e0.size == on.size == off.size == n):
                raise ValueError(
                    "outcomes, logging_propensities, target_on, target_off must be equal length"
                )
            if not bool(np.all((e0 > 0.0) & (e0 <= 1.0))):
                raise ValueError("logging_propensities must lie in (0, 1]")
        if self.treated is not None:
            f = np.asarray(self.treated).astype(bool)
            if f.size != n:
                raise ValueError("treated must match outcomes length")
            if not (f.any() and (~f).any()):
                raise ValueError("both arms must be present in `treated`")
        if not (self.has_msm or self.has_pivotality):
            raise ValueError(
                "supply at least one evidence source: logging_propensities (+ target_on/off) for "
                "the MSM layer, or treated + confounder_bins / mi_cap for the pivotality layer"
            )

    @property
    def has_msm(self) -> bool:
        """Whether the MSM sensitivity layer can run (logging propensities supplied)."""
        return self.logging_propensities is not None

    @property
    def has_pivotality(self) -> bool:
        """Whether the structural sign-robustness layer can run (binary arms + a channel bound)."""
        return self.treated is not None and (
            self.confounder_bins is not None or self.mi_cap is not None
        )

    @classmethod
    def from_binary(
        cls,
        outcomes: Sequence[float],
        treated: Sequence[int],
        *,
        propensities: Sequence[float] | None = None,
        confounder_bins: Sequence[int] | None = None,
        mi_cap: float | None = None,
    ) -> PolicyValueContrast:
        """The raw-logs case: one-hot arms ``pi_on = 1{treated}``, ``pi_off = 1{control}``.

        Exactly the contrast :func:`causalrl.certify_decision` builds internally. The MSM layer runs
        only when ``propensities`` are given; the pivotality layer runs when ``confounder_bins`` or
        ``mi_cap`` is given.
        """
        f = np.asarray(treated).astype(float)
        on = f.tolist() if propensities is not None else None
        off = (1.0 - f).tolist() if propensities is not None else None
        return cls(
            outcomes=outcomes,
            logging_propensities=propensities,
            target_on=on,
            target_off=off,
            treated=list(treated),
            confounder_bins=confounder_bins,
            mi_cap=mi_cap,
        )

has_msm property

Whether the MSM sensitivity layer can run (logging propensities supplied).

has_pivotality property

Whether the structural sign-robustness layer can run (binary arms + a channel bound).

from_binary(outcomes, treated, *, propensities=None, confounder_bins=None, mi_cap=None) classmethod

The raw-logs case: one-hot arms pi_on = 1{treated}, pi_off = 1{control}.

Exactly the contrast :func:causalrl.certify_decision builds internally. The MSM layer runs only when propensities are given; the pivotality layer runs when confounder_bins or mi_cap is given.

Source code in src/causalrl/identification/estimate.py

@classmethod
def from_binary(
    cls,
    outcomes: Sequence[float],
    treated: Sequence[int],
    *,
    propensities: Sequence[float] | None = None,
    confounder_bins: Sequence[int] | None = None,
    mi_cap: float | None = None,
) -> PolicyValueContrast:
    """The raw-logs case: one-hot arms ``pi_on = 1{treated}``, ``pi_off = 1{control}``.

    Exactly the contrast :func:`causalrl.certify_decision` builds internally. The MSM layer runs
    only when ``propensities`` are given; the pivotality layer runs when ``confounder_bins`` or
    ``mi_cap`` is given.
    """
    f = np.asarray(treated).astype(float)
    on = f.tolist() if propensities is not None else None
    off = (1.0 - f).tolist() if propensities is not None else None
    return cls(
        outcomes=outcomes,
        logging_propensities=propensities,
        target_on=on,
        target_off=off,
        treated=list(treated),
        confounder_bins=confounder_bins,
        mi_cap=mi_cap,
    )

causalrl.identification.transport.TransportFormula dataclass

How to compute P*(y | do(x)) from source and target data.

Source code in src/causalrl/identification/transport.py

@dataclass(frozen=True)
class TransportFormula:
    """How to compute ``P*(y | do(x))`` from source and target data."""

    kind: Literal["direct", "adjustment"]
    adjustment_set: frozenset[str]
    expression: str

causalrl.identification.transport_regret.TransportRegretCertificate dataclass

The transport-regret certificate for P*(value | do(action)) and a specific policy.

Attributes

transportable: :func:is_transportable_effect on the selection diagram. formula: The readable :class:TransportFormula for the transportable part, or None. non_transportable_witness: The selection-marked variable(s) the policy must not trust off-distribution: the selection-marked variables appearing in the transport formula's adjustment set (the mechanism one would have to re-weight by the target marginal to transport), or the full selection set if the effect is outright non-identifiable. Empty for a direct formula with empty adjustment — the certified zero-regret case. reweight_required: Whether transporting the queried effect requires re-weighting by a selection-marked variable (non_transportable_witness non-empty). Can be True even when transportable is True: a policy that merely reads the witness without re-weighting still incurs the bounded regret. decision_dependence: mu — the fraction of decision units whose executed behaviour depends on the selection-marked mechanism, measured by the do()-sweep (:func:decision_flip_rate). A property of the specific policy being certified, supplied by the caller. value_range: (V_min, V_max) of the policy value, used to scale the bound. regret_bound: Interval(0, (V_max - V_min) * mu) — the computable transfer-regret upper bound. Collapses to Interval(0, 0) for a behaviour-invariant policy (mu = 0).

Source code in src/causalrl/identification/transport_regret.py

@dataclass(frozen=True)
class TransportRegretCertificate:
    """The transport-regret certificate for ``P*(value | do(action))`` and a specific policy.

    Attributes
    ----------
    transportable:
        :func:`is_transportable_effect` on the selection diagram.
    formula:
        The readable :class:`TransportFormula` for the transportable part, or ``None``.
    non_transportable_witness:
        The selection-marked variable(s) the policy must not trust off-distribution: the
        selection-marked variables appearing in the transport formula's adjustment set (the
        mechanism one would have to re-weight by the target marginal to transport), or the full
        selection set if the effect is outright non-identifiable. Empty for a ``direct`` formula
        with empty adjustment — the certified zero-regret case.
    reweight_required:
        Whether transporting the queried effect requires re-weighting by a selection-marked
        variable (``non_transportable_witness`` non-empty). Can be ``True`` even when
        ``transportable`` is ``True``: a policy that merely *reads* the witness without
        re-weighting still incurs the bounded regret.
    decision_dependence:
        ``mu`` — the fraction of decision units whose executed behaviour depends on the
        selection-marked mechanism, measured by the ``do()``-sweep (:func:`decision_flip_rate`).
        A property of the *specific policy being certified*, supplied by the caller.
    value_range:
        ``(V_min, V_max)`` of the policy value, used to scale the bound.
    regret_bound:
        ``Interval(0, (V_max - V_min) * mu)`` — the computable transfer-regret upper bound.
        Collapses to ``Interval(0, 0)`` for a behaviour-invariant policy (``mu = 0``).
    """

    transportable: bool
    formula: TransportFormula | None
    non_transportable_witness: frozenset[str]
    reweight_required: bool
    decision_dependence: float
    value_range: tuple[float, float]
    regret_bound: Interval

    def is_vacuous(self, *, frac: float = 0.9) -> bool:
        """Whether the bound is vacuous: its width >= ``frac`` of the trivial value span."""
        v_min, v_max = self.value_range
        span = v_max - v_min
        if span <= 0:
            return True
        width = self.regret_bound.upper - self.regret_bound.lower
        return width >= frac * span

is_vacuous(*, frac=0.9)

Whether the bound is vacuous: its width >= frac of the trivial value span.

Source code in src/causalrl/identification/transport_regret.py

def is_vacuous(self, *, frac: float = 0.9) -> bool:
    """Whether the bound is vacuous: its width >= ``frac`` of the trivial value span."""
    v_min, v_max = self.value_range
    span = v_max - v_min
    if span <= 0:
        return True
    width = self.regret_bound.upper - self.regret_bound.lower
    return width >= frac * span

causalrl.identification.criteria.backdoor_adjustment_set(graph, treatment, outcome)

Return the observed-parent back-door adjustment set for treatment.

SCOPE / PRECONDITION: this returns parents(treatment), which is a valid parent adjustment set when the treatment is not incident to latent confounding. It does not implement front-door identification or the general ID/sID/gID algorithms. Graphs outside that supported contract raise rather than returning an invalid adjustment set.

Source code in src/causalrl/identification/criteria.py

def backdoor_adjustment_set(graph: CausalGraph, treatment: str, outcome: str) -> set[str]:
    """Return the observed-parent back-door adjustment set for `treatment`.

    SCOPE / PRECONDITION: this returns ``parents(treatment)``, which is a valid parent
    adjustment set when the treatment is not incident to latent confounding. It does not
    implement front-door identification or the general ID/sID/gID algorithms. Graphs outside
    that supported contract raise rather than returning an invalid adjustment set.
    """
    parents = set(graph.parents(treatment))
    if graph.has_incident_bidirected_edges(treatment):
        raise NotIdentifiableError(
            "parent adjustment is unsupported when treatment has latent confounding; "
            "a front-door or general ID algorithm may be required"
        )
    return parents

causalrl.identification.decision.certify_estimate(estimate, *, gamma_max=10.0, labels=('pi_on', 'pi_off'))

Certify whether an off-policy value contrast's sign is robust to hidden confounding.

Runs the MSM tipping layer over the general V(pi_on) - V(pi_off) contribution bound when estimate carries logging propensities, and the structural pivotality layer when it carries a binary-arm reduction (treated + confounder_bins / mi_cap). Returns the same :class:DecisionCertificate as :func:certify_decision; labels names the two policies in the verdict (labels[0] when the contrast is positive).

Honest scope: the MSM sensitivity is on the logging propensities (sharp when the two target supports are disjoint, valid-but-conservative otherwise); the pivotality layer is defined only for the binary-arm contrast. No sensitivity claim is made for arbitrary outcome-model estimators.

Source code in src/causalrl/identification/decision.py

def certify_estimate(
    estimate: PolicyValueContrast,
    *,
    gamma_max: float = 10.0,
    labels: tuple[str, str] = ("pi_on", "pi_off"),
) -> DecisionCertificate:
    """Certify whether an off-policy value contrast's sign is robust to hidden confounding.

    Runs the MSM tipping layer over the general ``V(pi_on) - V(pi_off)`` contribution bound when
    ``estimate`` carries logging propensities, and the structural pivotality layer when it carries a
    binary-arm reduction (``treated`` + ``confounder_bins`` / ``mi_cap``). Returns the same
    :class:`DecisionCertificate` as :func:`certify_decision`; ``labels`` names the two policies in
    the verdict (``labels[0]`` when the contrast is positive).

    Honest scope: the MSM sensitivity is on the logging propensities (sharp when the two target
    supports are disjoint, valid-but-conservative otherwise); the pivotality layer is defined only
    for the binary-arm contrast. No sensitivity claim is made for arbitrary outcome-model
    estimators.
    """
    y = np.asarray(estimate.outcomes, dtype=float)

    treated = estimate.treated
    pivot: PivotalityCertificate | None = None
    if treated is not None and (
        estimate.confounder_bins is not None or estimate.mi_cap is not None
    ):
        pivot = pivotality_certificate(
            estimate.outcomes, treated, estimate.confounder_bins, mi_cap=estimate.mi_cap
        )

    if treated is not None:
        fb = np.asarray(treated).astype(bool)
        naive = float(y[fb].mean() - y[~fb].mean())
    else:
        naive = 0.0  # replaced from the MSM point below (a target-only contrast always has MSM)

    g_tip: float | None = None
    msm_certified: bool | None = None
    logging_propensities = estimate.logging_propensities
    if logging_propensities is not None:
        on, off = estimate.target_on, estimate.target_off
        assert on is not None and off is not None  # guaranteed by PolicyValueContrast.__post_init__
        outcomes = estimate.outcomes
        e0_l, on_l, off_l = list(logging_propensities), list(on), list(off)

        def _band(g: float) -> Interval:
            return msm_contribution_bounds(outcomes, e0_l, on_l, off_l, gamma=g)

        if treated is None:
            naive = float(_band(1.0).lower)
        g_tip = tipping_gamma(_band, reference=0.0, gamma_max=gamma_max)
        msm_certified = g_tip is None

    decision = (
        f"prefer {labels[0]}"
        if naive > 0
        else f"prefer {labels[1]}"
        if naive < 0
        else "indifferent"
    )
    certified = pivot.certified if pivot is not None else bool(msm_certified)
    summary = _summarise(decision, naive, pivot, g_tip, msm_certified, gamma_max)
    return DecisionCertificate(
        decision=decision,
        naive_contrast=naive,
        certified=certified,
        pivotality=pivot,
        tipping_gamma=g_tip,
        msm_certified=msm_certified,
        summary=summary,
    )

causalrl.transport.estimate.certify_transported_policy_value(diagram, source_data, target_data, *, stages, outcome, target_actions, alpha=0.05, outcome_model=None, policy='pi')

Certify a deterministic policy's transported finite-horizon value (hedge-first; §7.5).

stages[t] is {"history": (names...), "treatment": name} for stage t; target_actions gives a constant target action per stage. Sequential/policy-value transport is research-grade, so this is identified only in one subcase — when the selection difference is confined to the baseline distribution H_1 (a population shift, downstream mechanisms shared): the source sequential g-computation is reweighted to the target baseline marginal (discrete baseline strata, as in :func:certify_transported_effect). Any selection node on a time-varying covariate, treatment, or the outcome yields a hedge (I3), with transport_regret_certificate the shipped floor. Sequential ignorability is recorded as a non-checkable assumption.

Source code in src/causalrl/transport/estimate.py

def certify_transported_policy_value(
    diagram: SelectionDiagram,
    source_data: Mapping[str, Any],
    target_data: Mapping[str, Any],
    *,
    stages: Sequence[Mapping[str, Any]],
    outcome: str,
    target_actions: Sequence[float],
    alpha: float = 0.05,
    outcome_model: OutcomeFactory | None = None,
    policy: str = "pi",
) -> Certificate:
    """Certify a deterministic policy's transported finite-horizon value (hedge-first; §7.5).

    ``stages[t]`` is ``{"history": (names...), "treatment": name}`` for stage ``t``;
    ``target_actions`` gives a constant target action per stage. Sequential/policy-value transport
    is research-grade, so this is **identified only in one subcase** — when the selection difference
    is confined to the baseline distribution ``H_1`` (a population shift, downstream mechanisms
    shared): the source sequential g-computation is reweighted to the target baseline marginal
    (discrete baseline strata, as in :func:`certify_transported_effect`). Any selection node on a
    time-varying covariate, treatment, or the outcome yields a **hedge** (I3), with
    ``transport_regret_certificate`` the shipped floor. Sequential ignorability is recorded as a
    non-checkable assumption.
    """
    if not stages:
        raise ValueError("need at least one stage")
    if len(target_actions) != len(stages):
        raise ValueError("target_actions must have one action per stage")
    baseline_vars = [str(v) for v in stages[0]["history"]]
    selection = frozenset(diagram.selection_variables)
    horizon = len(stages)

    if not selection <= set(baseline_vars):
        return _hedged_sequential(
            diagram, sorted(selection), horizon, reason="non-transportable-sequential"
        )

    y_src = np.asarray(source_data[outcome], dtype=np.float64)
    n_src = len(y_src)
    histories = [_stack_named(source_data, list(s["history"])) for s in stages]
    treatments = [np.asarray(source_data[s["treatment"]], dtype=np.float64) for s in stages]
    tgt_actions = [np.full(n_src, float(a), dtype=np.float64) for a in target_actions]
    q1 = sequential_ice_values(
        histories, treatments, tgt_actions, y_src, outcome_model=outcome_model
    )

    if baseline_vars:
        baseline_src = _strata_keys(source_data, baseline_vars)
        baseline_tgt = _strata_keys(target_data, baseline_vars)
    else:
        n_tgt = len(np.asarray(target_data[outcome]))
        baseline_src = [() for _ in range(n_src)]
        baseline_tgt = [() for _ in range(n_tgt)]
    value = transport_sequential_gcomp(q1.tolist(), baseline_src, baseline_tgt)

    return Certificate(
        claim=f"V*({policy}) over horizon {horizon} = {value:.4g}",
        estimand=EstimandSpec(
            query="policy_value", target="mean", policy=policy, domains=("source", "target")
        ),
        kind=Kind.IDENTIFIED,
        value=value,
        alpha=alpha,
        assumptions=(
            Assumption("sequential-ignorability", {"horizon": horizon}, checkable=False),
            Assumption(
                "selection-diagram",
                {"selection": sorted(selection), "scope": "baseline-only"},
                checkable=True,
            ),
        ),
        method="sequential-transport-gcomp",
        witness=Witness(
            kind="sequential-baseline-transport",
            detail={"baseline": baseline_vars, "selection": sorted(selection), "horizon": horizon},
        ),
        hedge=None,
        provenance=Provenance.create(graph_hash=graph_hash(diagram.graph)),
    )

causalrl.identification.bounds.confounding_bias_per_step_bounds(outcomes, treated_by_step, confounder_bins, *, form='tv')

Per-step channel-cap bounds for sequential per-factor credits (THEORY Theorem-seq).

One row per unit (e.g. per episode); treated_by_step[t] is step t's binary action factor for every unit, confounder_bins the shared episode-level hidden variable. Each step's omitted-Z credit bias |naive_t - adjusted_t| is bounded by :func:confounding_bias_bound at that step's arm split; a cumulative additive credit's bias is bounded by the sum. This is the average-budget (information) sibling of the uniform-odds :func:msm_per_step_bounds: there each step gets an MSM gamma budget, here each step gets the channel budget MI(F_t;Z) <= MI(I_t;Z) of ITS OWN information set — in games, later information sets typically leak more about Z, so later credits are systematically less protected (the growing-channel prediction).

Source code in src/causalrl/identification/bounds.py

def confounding_bias_per_step_bounds(
    outcomes: Sequence[float],
    treated_by_step: Sequence[Sequence[int]],
    confounder_bins: Sequence[int],
    *,
    form: str = "tv",
) -> list[float]:
    """Per-step channel-cap bounds for sequential per-factor credits (THEORY Theorem-seq).

    One row per unit (e.g. per episode); ``treated_by_step[t]`` is step ``t``'s binary action
    factor for every unit, ``confounder_bins`` the shared episode-level hidden variable. Each
    step's omitted-``Z`` credit bias ``|naive_t - adjusted_t|`` is bounded by
    :func:`confounding_bias_bound` at that step's arm split; a cumulative additive credit's
    bias is bounded by the sum. This is the average-budget (information) sibling of the
    uniform-odds :func:`msm_per_step_bounds`: there each step gets an MSM ``gamma`` budget,
    here each step gets the channel budget ``MI(F_t;Z) <= MI(I_t;Z)`` of ITS OWN information
    set — in games, later information sets typically leak more about ``Z``, so later credits
    are systematically less protected (the growing-channel prediction).
    """
    return [
        confounding_bias_bound(outcomes, f_t, confounder_bins, form=form) for f_t in treated_by_step
    ]

causalrl.identification.counterfactual_bounds.counterfactual_interval(scm, *, evidence, interventions, target, n=20000, seed=None)

Bound E[target_{do(interventions)} | evidence] on a fitted SCM.

A fitted discrete mechanism pins P(V | parents) but not the coupling from noise to value, and the two are indistinguishable in L1 and L2 data. So the counterfactual is an interval over admissible couplings, not a point. Its extremes solve a box-and-sum problem: the factual and counterfactual columns of the coupling share known marginals, which bounds each cell, and the conditioning event fixes their sum.

Invertible mechanisms contribute no width — their noise is recoverable — so an all-invertible SCM returns a point, matching :meth:StructuralCausalModel.abduct exactly. So does an intervention that pins target itself, or one that leaves target's parents at their factual values: neither leaves any coupling to choose.

What tight=True claims. The bound is exact given the fitted conditionals: no relaxation is applied to them, and the extremes are solved in closed form rather than approximated. Those conditionals are not exact, though — they are Monte-Carlo estimates, n draws from target's exogenous distribution pushed through its fitted mechanism at each parent configuration, under seed (None gives a fixed default stream, not a random one). So the endpoints carry that estimate's sampling error, and a level whose fitted probability is far below 1/n can be missing from the sampled support altogether, which narrows the reported interval below the identified one. Raise n to shrink both effects; tight describes the estimator, not the number of correct digits.

Source code in src/causalrl/identification/counterfactual_bounds.py

def counterfactual_interval(
    scm: StructuralCausalModel,
    *,
    evidence: Mapping[str, float],
    interventions: Mapping[str, Value],
    target: str,
    n: int = 20_000,
    seed: int | None = None,
) -> CounterfactualBound:
    """Bound ``E[target_{do(interventions)} | evidence]`` on a fitted SCM.

    A fitted discrete mechanism pins ``P(V | parents)`` but not the coupling from noise to value,
    and the two are indistinguishable in L1 and L2 data. So the counterfactual is an interval over
    admissible couplings, not a point. Its extremes solve a box-and-sum problem: the factual and
    counterfactual columns of the coupling share known marginals, which bounds each cell, and the
    conditioning event fixes their sum.

    Invertible mechanisms contribute no width — their noise is recoverable — so an all-invertible
    SCM returns a point, matching :meth:`StructuralCausalModel.abduct` exactly. So does an
    intervention that pins ``target`` itself, or one that leaves ``target``'s parents at their
    factual values: neither leaves any coupling to choose.

    **What ``tight=True`` claims.** The bound is exact *given the fitted conditionals*: no
    relaxation is applied to them, and the extremes are solved in closed form rather than
    approximated. Those conditionals are not exact, though — they are Monte-Carlo estimates,
    ``n`` draws from ``target``'s exogenous distribution pushed through its fitted mechanism at
    each parent configuration, under ``seed`` (``None`` gives a fixed default stream, not a
    random one). So the endpoints carry that estimate's sampling error, and a level whose fitted
    probability is far below ``1/n`` can be missing from the sampled support altogether, which
    narrows the reported interval below the identified one. Raise ``n`` to shrink both effects;
    ``tight`` describes the estimator, not the number of correct digits.
    """
    if scm.provenance != "fitted":
        raise ValueError(
            "counterfactual_interval is for fitted SCMs, whose noise-to-value coupling is not "
            "identified. A specified SCM's mechanisms are asserted by their author, so its "
            "counterfactual is a point — use counterfactual_expectation instead."
        )
    missing = [node for node in scm.graph.nodes if node not in evidence]
    if missing:
        raise KeyError(
            f"evidence must cover every node for a counterfactual on a fully observed unit; "
            f"missing {sorted(missing)}"
        )

    if target in interventions:
        # do(target = v) replaces the target's own mechanism with a constant, so its
        # counterfactual value IS v -- no coupling is involved and nothing upstream can reach it.
        # Checked before the ambiguity analysis below for that reason: an ambiguous node between
        # the intervention and a target the same do() pins cannot make this query unanswerable.
        pinned = _scalar(interventions[target])
        return CounterfactualBound(pinned, pinned, True)

    intervened = set(interventions)
    ambiguous = set(scm.non_invertible_nodes())
    upstream = sorted(_ambiguous_upstream(scm, ambiguous, intervened, target))
    if upstream:
        raise NotIdentifiableError(
            f"node(s) {upstream} are non-invertible and lie upstream of {target!r} on a "
            "directed path from the intervention; composing per-node bounds would be loose rather "
            "than tight. The tight answer needs the neural-causal-model min/max.",
            witness=upstream,
        )

    cf_parents = _counterfactual_parents(scm, target, evidence, interventions)
    factual_parents = {p: float(evidence[p]) for p in scm.graph.parents(target)}

    if target not in ambiguous:
        # Exact: invert the target's noise from the factual unit, replay it on the cf parents.
        mechanism = scm.mechanisms[target]
        # residual is attached dynamically (Tasks 4-5), so it is unknown to the Mechanism
        # protocol -- cast rather than annotate, matching evaluate_holdout's log_prob pattern
        # (fitters.py), so the Unknown does not propagate into the mechanism(...) call below.
        residual = cast("_Residual", mechanism.residual)  # type: ignore[attr-defined]
        noise = residual(_unit(factual_parents), _unit({target: float(evidence[target])})[target])
        value = float(mechanism(_unit(cf_parents), noise).reshape(-1)[0])
        return CounterfactualBound(value, value, True)

    factual = _pmf_at(scm, target, factual_parents, n=n, seed=seed)
    observed = float(evidence[target])
    p_factual = factual.get(observed, 0.0)
    if p_factual <= 0.0:
        raise RealizabilityError(
            f"evidence {target}={observed} has zero fitted probability at the factual parent "
            f"configuration {factual_parents}; the counterfactual is not defined"
        )

    if _same_configuration(cf_parents, factual_parents):
        # The intervention leaves the target's parents where they factually were, so Y_cf and Y_f
        # are the same random variable: every admissible coupling puts all its mass on the
        # diagonal, and the identified set is the point the unit was observed at. Relaxing that to
        # Frechet limits would report a modelling choice as a result. Covers "what if we had done
        # what we did" and placebo / negative-control queries. Placed after the realizability
        # check so evidence the fitted model deems impossible is still refused rather than
        # answered.
        return CounterfactualBound(observed, observed, True)

    counterfactual = _pmf_at(scm, target, cf_parents, n=n, seed=seed)
    values = np.array(sorted(counterfactual), dtype=float)
    marginal = np.array([counterfactual[v] for v in values], dtype=float)
    # Cell pi(v) = P(target_cf = v, target_f = observed). Each cell is capped by its own marginal
    # and by the conditioning mass, floored by what the rest of the table cannot absorb, and the
    # cells sum to p_factual — a box-and-sum problem solved exactly by a fractional fill.
    upper_cell = np.minimum(marginal, p_factual)
    lower_cell = np.maximum(0.0, marginal - (1.0 - p_factual))
    low = _extreme_under_sum(values, lower_cell, upper_cell, p_factual, maximize=False)
    high = _extreme_under_sum(values, lower_cell, upper_cell, p_factual, maximize=True)
    return CounterfactualBound(float(low / p_factual), float(high / p_factual), True)

causalrl.identification.transport_regret.decision_abstain_mask(greedy, units, *, values=(0, 1))

Abstention mask: fire where the policy's immediate greedy decision is not invariant to intervening on the selection-marked witness.

Fires on unit iff greedy(unit, w0) != greedy(unit, w1). This is the per-unit decision-flip rule — by construction it recovers the truly-confused set when "confused" is defined as the same single-decision flip (the mechanistic guarantee the G1 gate's M4 pins). A behaviour-invariant policy has an (almost) empty mask. Distinct from :func:decision_flip_rate's trajectory-divergence mu that scales the regret bound.

Source code in src/causalrl/identification/transport_regret.py

def decision_abstain_mask(
    greedy: Callable[[U, int], object],
    units: Sequence[U],
    *,
    values: tuple[int, int] = (0, 1),
) -> NDArray[np.bool_]:
    """Abstention mask: fire where the policy's *immediate* greedy decision is not invariant to
    intervening on the selection-marked witness.

    Fires on ``unit`` iff ``greedy(unit, w0) != greedy(unit, w1)``. This is the per-unit
    decision-flip rule — by construction it recovers the truly-confused set when "confused" is
    defined as the same single-decision flip (the mechanistic guarantee the G1 gate's M4 pins).
    A behaviour-invariant policy has an (almost) empty mask. Distinct from
    :func:`decision_flip_rate`'s trajectory-divergence ``mu`` that scales the regret bound.
    """
    w0, w1 = values
    mask = np.zeros(len(units), dtype=bool)
    for i, u in enumerate(units):
        mask[i] = greedy(u, w0) != greedy(u, w1)
    return mask

causalrl.identification.transport_regret.decision_flip_rate(trace, units, *, values=(0, 1))

mu — fraction of units whose executed trajectory diverges under the do()-sweep.

trace(unit, w) returns the policy's executed rollout trace (action or visited-state sequence — any hashable summary of what the policy does) from unit under do(W = w), with everything else held fixed. Two traces differing is exactly the event on which V^pi(u; w0) != V^pi(u; w1) can be nonzero, so span * mu is a valid regret bound (module docstring). Use this — not the single-decision flip rate — to scale the bound.

Source code in src/causalrl/identification/transport_regret.py

def decision_flip_rate(
    trace: Callable[[U, int], Hashable],
    units: Sequence[U],
    *,
    values: tuple[int, int] = (0, 1),
) -> float:
    """``mu`` — fraction of units whose *executed trajectory* diverges under the ``do()``-sweep.

    ``trace(unit, w)`` returns the policy's executed rollout trace (action or visited-state
    sequence — any hashable summary of *what the policy does*) from ``unit`` under
    ``do(W = w)``, with everything else held fixed. Two traces differing is exactly the event on
    which ``V^pi(u; w0) != V^pi(u; w1)`` can be nonzero, so ``span * mu`` is a *valid* regret
    bound (module docstring). Use this — not the single-decision flip rate — to scale the bound.
    """
    if not units:
        return 0.0
    w0, w1 = values
    flips = sum(1 for u in units if trace(u, w0) != trace(u, w1))
    return flips / len(units)

causalrl.identification.criteria.is_identifiable(graph, treatment, outcome)

Whether P(outcome | do(treatment)) is identifiable from observational data.

Delegates to the sound and complete ID algorithm (:func:causalrl.identification.id_algorithm.is_identifiable_effect), so it returns a definite boolean for any ADMG — including front-door-style cases the earlier scoped heuristic could only report as unknown. For the estimand itself, see :func:~causalrl.identification.id_algorithm.identify_effect and :func:~causalrl.identification.id_algorithm.estimate_effect.

Source code in src/causalrl/identification/criteria.py

def is_identifiable(graph: CausalGraph, treatment: str, outcome: str) -> bool:
    """Whether ``P(outcome | do(treatment))`` is identifiable from observational data.

    Delegates to the sound and complete ID algorithm
    (:func:`causalrl.identification.id_algorithm.is_identifiable_effect`), so it returns a definite
    boolean for any ADMG — including front-door-style cases the earlier scoped heuristic could only
    report as unknown. For the estimand itself, see
    :func:`~causalrl.identification.id_algorithm.identify_effect` and
    :func:`~causalrl.identification.id_algorithm.estimate_effect`.
    """
    return is_identifiable_effect(graph, {treatment}, {outcome})

causalrl.discovery.orient(cpdag, *, tiers=None)

Orient a CPDAG's remaining undirected edges into a DAG, or refuse.

Resolution order per undirected edge: (1) tiers — an edge between different tiers points from the earlier tier to the later one; (2) acyclicity — if one direction would close a cycle, take the other; (3) refuse. Silently picking an orientation would commit to a choice the data does not identify, so the third case raises and names :func:causalrl.fit_scm_mec, which fits every member of the equivalence class instead. Exception: if tiers are provided and a tier-forced edge would create a cycle (indicating a conflict between tiers and discovered structure), raises immediately with a message naming the specific colliding edges.

Source code in src/causalrl/discovery.py

def orient(cpdag: CPDAG, *, tiers: Sequence[Sequence[str]] | None = None) -> CausalGraph:
    """Orient a CPDAG's remaining undirected edges into a DAG, or refuse.

    Resolution order per undirected edge: (1) ``tiers`` — an edge between different tiers points
    from the earlier tier to the later one; (2) acyclicity — if one direction would close a cycle,
    take the other; (3) refuse. Silently picking an orientation would commit to a choice the data
    does not identify, so the third case raises and names :func:`causalrl.fit_scm_mec`, which fits
    every member of the equivalence class instead. Exception: if ``tiers`` are provided and a
    tier-forced edge would create a cycle (indicating a conflict between tiers and discovered
    structure), raises immediately with a message naming the specific colliding edges.
    """
    rank: dict[str, int] = {}
    if tiers is not None:
        for level, tier in enumerate(tiers):
            for name in tier:
                rank[name] = level
        uncovered = set(cpdag.variables) - set(rank)
        if uncovered:
            raise CausalGraphError(
                f"variable(s) {sorted(uncovered)} not covered by tiers {list(tiers)}"
            )

    directed = set(cpdag.directed_edges)
    pending = sorted(tuple(sorted(edge)) for edge in cpdag.undirected_edges)
    unresolved: list[tuple[str, str]] = []
    # Repeat: each orientation can unlock another edge through the acyclicity rule.
    while pending:
        progressed = False
        deferred: list[tuple[str, str]] = []
        for a, b in pending:
            if a in rank and b in rank and rank[a] != rank[b]:
                # Orient according to tiers, but first check that it doesn't create a cycle.
                tail, head = (a, b) if rank[a] < rank[b] else (b, a)
                cycle_path = _creates_cycle(directed, tail, head)
                if cycle_path is not None:
                    # cycle_path is [head, ...nodes..., tail], forming path head -> ... -> tail.
                    # The new edge tail -> head would close cycle: tail -> head -> ... -> tail.
                    cycle_edges = " -> ".join([tail, *cycle_path])
                    raise CausalGraphError(
                        f"tier-implied edge {tail} -> {head} would create a cycle: {cycle_edges}"
                    )
                directed.add((tail, head))
                progressed = True
                continue
            ab_ok = _creates_cycle(directed, a, b) is None
            ba_ok = _creates_cycle(directed, b, a) is None
            if ab_ok and not ba_ok:
                directed.add((a, b))
                progressed = True
            elif ba_ok and not ab_ok:
                directed.add((b, a))
                progressed = True
            else:
                deferred.append((a, b))
        if not progressed:
            unresolved = deferred
            break
        pending = deferred

    if unresolved:
        raise CausalGraphError(
            f"cannot orient edge(s) {sorted(unresolved)}: neither tiers nor acyclicity decides "
            "them. Pass tiers=..., or use fit_scm_mec to fit every member of the equivalence class."
        )
    return CausalGraph(directed_edges=sorted(directed), nodes=list(cpdag.variables))

causalrl.identification.transport_regret.transport_regret_certificate(diagram, *, action, value, value_range, decision_dependence, max_adjustment_size=3)

Build the transport-regret certificate for P*(value | do(action)) on diagram.

Composes the library transport decision with the decision-regret wrapper:

• transportable / formula from the engine (:func:is_transportable_effect / :func:transport_formula).
• non_transportable_witness = selection-marked variables the transport formula must re-weight (they appear in its adjustment set); the full selection set if the effect is non-identifiable; empty for a direct formula with empty adjustment (certified zero).
• regret_bound = Interval(0, (V_max - V_min) * mu) with mu the supplied decision_dependence (the do()-sweep result, :func:decision_flip_rate — the trajectory-divergence rate, NOT the abstention statistic).

decision_dependence is supplied by the caller because it is a property of the specific policy being certified; the certificate primitive itself stays a pure composition of causal-engine outputs.

Source code in src/causalrl/identification/transport_regret.py

def transport_regret_certificate(
    diagram: SelectionDiagram,
    *,
    action: str,
    value: str,
    value_range: tuple[float, float],
    decision_dependence: float,
    max_adjustment_size: int = 3,
) -> TransportRegretCertificate:
    """Build the transport-regret certificate for ``P*(value | do(action))`` on ``diagram``.

    Composes the library transport decision with the decision-regret wrapper:

    * ``transportable`` / ``formula`` from the engine (:func:`is_transportable_effect` /
      :func:`transport_formula`).
    * ``non_transportable_witness`` = selection-marked variables the transport formula must
      re-weight (they appear in its adjustment set); the full selection set if the effect is
      non-identifiable; empty for a direct formula with empty adjustment (certified zero).
    * ``regret_bound`` = ``Interval(0, (V_max - V_min) * mu)`` with ``mu`` the supplied
      ``decision_dependence`` (the ``do()``-sweep result, :func:`decision_flip_rate` — the
      trajectory-divergence rate, NOT the abstention statistic).

    ``decision_dependence`` is supplied by the caller because it is a property of the *specific
    policy* being certified; the certificate primitive itself stays a pure composition of
    causal-engine outputs.
    """
    if not 0.0 <= decision_dependence <= 1.0:
        raise ValueError(f"decision_dependence must be in [0, 1], got {decision_dependence}")
    v_min, v_max = value_range
    if v_max < v_min:
        raise ValueError(
            f"value_range must be (V_min, V_max) with V_max >= V_min, got {value_range}"
        )

    selection = diagram.selection_variables
    transportable = is_transportable_effect(diagram.graph, [action], [value], selection)
    formula = transport_formula(
        diagram, treatment=action, outcome=value, max_adjustment_size=max_adjustment_size
    )

    if formula is None:
        witness: frozenset[str] = frozenset(selection)
    else:
        witness = frozenset(selection) & formula.adjustment_set
    reweight_required = len(witness) > 0

    span = v_max - v_min
    regret_bound = Interval(0.0, float(span * decision_dependence))

    return TransportRegretCertificate(
        transportable=transportable,
        formula=formula,
        non_transportable_witness=witness,
        reweight_required=reweight_required,
        decision_dependence=float(decision_dependence),
        value_range=(float(v_min), float(v_max)),
        regret_bound=regret_bound,
    )

Structural Models & Data — Complete Reference

causalrl.scm.fitters.ANMFit

Continuous node: additive noise model V = g(parents) + U for an arbitrary regressor.

estimator is a factory returning a fresh object with fit(X, y) / predict(X) — the same duck-typed contract as GFormulaBackdoorAgent(outcome_model=...), so scikit-learn stays optional. The default is a dependency-free ridge on RBF features. U is drawn from the empirical residual distribution, and is recoverable from (parents, value) — invertible=True.

Source code in src/causalrl/scm/fitters.py

class ANMFit:
    """Continuous node: additive noise model ``V = g(parents) + U`` for an arbitrary regressor.

    ``estimator`` is a factory returning a fresh object with ``fit(X, y)`` / ``predict(X)`` — the
    same duck-typed contract as ``GFormulaBackdoorAgent(outcome_model=...)``, so scikit-learn stays
    optional. The default is a dependency-free ridge on RBF features. ``U`` is drawn from the
    empirical residual distribution, and is recoverable from (parents, value) — ``invertible=True``.
    """

    def __init__(
        self, estimator: Callable[[], Any] | None = None, *, n_features: int = 32, seed: int = 0
    ) -> None:
        self._estimator = estimator
        self._n_features = n_features
        self._seed = seed

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        y = np.asarray(child, dtype=float)
        names = sorted(parents)
        if not names:
            return LinearGaussianFit().fit({}, y)
        raw = np.column_stack([np.asarray(parents[name], dtype=float) for name in names])
        model = (
            self._estimator()
            if self._estimator is not None
            else _RBFRidge(self._n_features, self._seed)
        )
        model.fit(raw, y)
        predicted = np.asarray(model.predict(raw), dtype=float)
        residual = y - predicted
        residual_tensor = torch.tensor(residual, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]

        def mechanism(parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
            flat_noise = noise.reshape(-1)
            columns = _parent_columns(parent_values, names, flat_noise.shape[0])
            mean = torch.tensor(  # type: ignore[reportPrivateImportUsage]
                np.asarray(model.predict(columns), dtype=float),
                dtype=torch.float32,  # type: ignore[reportPrivateImportUsage]
            )
            return mean + flat_noise

        fitted_mechanism = FunctionalMechanism(names, mechanism)
        _attach_residual(fitted_mechanism, names, lambda columns: model.predict(columns))
        return FittedMechanism(
            mechanism=fitted_mechanism,
            noise=_Empirical(residual_tensor),
            invertible=True,
            score=_r2(y, predicted),
        )

causalrl.scm.continuous.abduction.AmortizedGaussianAbduction

Bases: Module

Amortized VI encoder q(u | parents, y) = N(mu, sigma) for a fixed mechanism.

Trained by ELBO (Gaussian reparameterisation; standard-normal prior) to reconstruct the observed outcome through the mechanism. General — works for non-invertible mechanisms — so its counterfactuals are kind=EMPIRICAL.

Source code in src/causalrl/scm/continuous/abduction.py

class AmortizedGaussianAbduction(torch.nn.Module):
    """Amortized VI encoder ``q(u | parents, y) = N(mu, sigma)`` for a fixed mechanism.

    Trained by ELBO (Gaussian reparameterisation; standard-normal prior) to reconstruct the observed
    outcome through the mechanism. General — works for non-invertible mechanisms — so its
    counterfactuals are ``kind=EMPIRICAL``.
    """

    def __init__(
        self,
        mechanism: torch.nn.Module,
        parents: list[str],
        hidden: tuple[int, ...] = (32, 32),
        *,
        name: str = "U",
    ) -> None:
        super().__init__()
        self.mechanism = mechanism
        self.parents = parents
        self.name = name
        self.encoder = build_mlp(len(parents) + 1, hidden, 2)

    def _q(self, parent_values: Mapping[str, Tensor], y: Tensor) -> tuple[Tensor, Tensor]:
        cols = [parent_values[p].reshape(-1, 1) for p in self.parents] + [y.reshape(-1, 1)]
        h = self.encoder(torch.cat(cols, dim=1))  # type: ignore[reportPrivateImportUsage,reportUnknownMemberType]
        scale = torch.nn.functional.softplus(h[:, 1]) + 1e-4  # type: ignore[reportPrivateImportUsage]
        return h[:, 0], scale

    def fit(
        self,
        parent_values: Mapping[str, Tensor],
        y: Tensor,
        *,
        steps: int = 400,
        lr: float = 1e-2,
        beta: float = 1.0,
        seed: int = 0,
    ) -> AmortizedGaussianAbduction:
        torch.manual_seed(seed)  # type: ignore[reportPrivateImportUsage]
        # Train the encoder only; the mechanism is fixed (abduction infers noise given it).
        opt = torch.optim.Adam(self.encoder.parameters(), lr=lr)  # type: ignore[reportPrivateImportUsage]
        pv = dict(parent_values)
        target = y.detach()  # a fixed observed target (drop any graph carried in on `y`)
        for _ in range(steps):
            mean, scale = self._q(pv, target)
            u = mean + scale * torch.randn_like(mean)  # type: ignore[reportPrivateImportUsage]
            recon = ((self.mechanism(pv, u) - target) ** 2).mean()
            kl = (0.5 * (mean**2 + scale**2 - 1.0) - torch.log(scale)).mean()  # type: ignore[reportPrivateImportUsage]
            loss = recon + beta * kl
            opt.zero_grad()
            loss.backward()
            opt.step()  # type: ignore[reportUnknownMemberType]
        return self

    def posterior(self, parent_values: Mapping[str, Tensor], y: Tensor) -> AmortizedNoisePosterior:
        with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
            mean, scale = self._q(dict(parent_values), y)
        return AmortizedNoisePosterior(mean, scale, name=self.name)

causalrl.scm.continuous.bayesian_fit.BayesianLinearFit

Continuous node: Bayesian linear structural equation, posterior sampled by NUTS (NumPyro).

Where :class:~causalrl.scm.fitters.LinearGaussianFit solves a single OLS point estimate, this samples the full posterior over (intercept, weights, sigma) and returns the posterior-mean as a :class:~causalrl.scm.mechanisms.LinearGaussianMechanism -- so every invariant the point-estimated continuous families satisfy still holds here: invertible=True, a residual closure is attached, the round-trip identity holds exactly, and evaluate_holdout scores it by R^2 like its siblings. The draws behind that mean are attached as mechanism.posterior, a dict[str, np.ndarray] keyed "intercept", each parent name, and "sigma" -- plain NumPy, not JAX, so a caller without JAX in scope can use them directly. Materialising one SCM per draw is a natural follow-on; this fitter does not build it.

Model: intercept ~ Normal(0, 10), w_j ~ Normal(0, 10) per parent, sigma ~ HalfNormal(1), y ~ Normal(intercept + X @ w, sigma). Both X and y are standardised internally so the weakly-informative priors stay sensible regardless of native scale on either side -- leaving y unstandardised would put sigma's HalfNormal(1) prior on the child's raw scale, silently shrinking both sigma and the weights toward zero whenever that scale is far from 1. The posterior draws are unstandardised back onto the original data scale before being exposed.

Source code in src/causalrl/scm/continuous/bayesian_fit.py

class BayesianLinearFit:
    """Continuous node: Bayesian linear structural equation, posterior sampled by NUTS (NumPyro).

    Where :class:`~causalrl.scm.fitters.LinearGaussianFit` solves a single OLS point estimate,
    this samples the full posterior over ``(intercept, weights, sigma)`` and returns the
    posterior-**mean** as a :class:`~causalrl.scm.mechanisms.LinearGaussianMechanism` -- so every
    invariant the point-estimated continuous families satisfy still holds here: ``invertible=True``,
    a ``residual`` closure is attached, the round-trip identity holds exactly, and
    ``evaluate_holdout`` scores it by R^2 like its siblings. The draws behind that mean are attached
    as ``mechanism.posterior``, a ``dict[str, np.ndarray]`` keyed ``"intercept"``, each parent name,
    and ``"sigma"`` -- plain NumPy, not JAX, so a caller without JAX in scope can use them directly.
    Materialising one SCM per draw is a natural follow-on; this fitter does not build it.

    Model: ``intercept ~ Normal(0, 10)``, ``w_j ~ Normal(0, 10)`` per parent, ``sigma ~
    HalfNormal(1)``, ``y ~ Normal(intercept + X @ w, sigma)``. Both ``X`` and ``y`` are standardised
    internally so the weakly-informative priors stay sensible regardless of native scale on either
    side -- leaving ``y`` unstandardised would put ``sigma``'s ``HalfNormal(1)`` prior on the
    child's raw scale, silently shrinking both sigma and the weights toward zero whenever that
    scale is far from 1. The posterior draws are unstandardised back onto the original data scale
    before being exposed.
    """

    def __init__(self, draws: int = 1000, warmup: int = 1000, seed: int = 0) -> None:
        self.draws = draws
        self.warmup = warmup
        self.seed = seed

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        # Guarded here rather than at module scope: the extra is optional, and the error must name
        # it instead of surfacing a bare ModuleNotFoundError from six lines down.
        try:
            import jax
            import jax.numpy as jnp
            import numpyro
            import numpyro.distributions as dist
            from numpyro.infer import MCMC, NUTS
        except ImportError as exc:  # pragma: no cover - exercised only without the extra
            raise ImportError(
                "BayesianLinearFit requires NumPyro (and JAX); install the 'causalrl[numpyro]' "
                "extra"
            ) from exc

        y = np.asarray(child, dtype=float)
        n = len(y)
        names = sorted(parents)
        reserved = {"intercept", "sigma"} & set(names)
        if reserved:
            raise ValueError(
                f"BayesianLinearFit.fit: parent name(s) {sorted(reserved)} collide with the "
                "reserved posterior keys 'intercept' and 'sigma' (the fitted intercept and "
                "residual noise scale); rename the parent column(s) before fitting."
            )
        n_parents = len(names)
        raw = (
            np.column_stack([np.asarray(parents[name], dtype=float) for name in names])
            if names
            else np.zeros((n, 0))
        )
        # Standardise both X and y so N(0, 10) / HalfNormal(1) stay weakly-informative regardless
        # of native scale on either side; a constant column (std 0) is left unscaled rather than
        # divided by zero.
        mean_x = raw.mean(axis=0)
        scale_x = raw.std(axis=0)
        safe_scale_x = np.where(scale_x > 1e-9, scale_x, 1.0)
        x_std = (raw - mean_x) / safe_scale_x

        mean_y = float(y.mean())
        scale_y = float(y.std())
        safe_scale_y = scale_y if scale_y > 1e-9 else 1.0
        y_std = (y - mean_y) / safe_scale_y

        x_jax = jnp.asarray(x_std)
        y_jax = jnp.asarray(y_std)

        def model(x: Any, y_obs: Any) -> None:
            intercept = numpyro.sample("intercept", dist.Normal(0.0, 10.0))
            if n_parents > 0:
                w = numpyro.sample("w", dist.Normal(0.0, 10.0).expand([n_parents]).to_event(1))
                mean = intercept + x @ w
            else:
                mean = intercept + jnp.zeros(y_obs.shape[0])
            sigma = numpyro.sample("sigma", dist.HalfNormal(1.0))
            with numpyro.plate("data", y_obs.shape[0]):
                numpyro.sample("y", dist.Normal(mean, sigma), obs=y_obs)

        mcmc = MCMC(NUTS(model), num_warmup=self.warmup, num_samples=self.draws, progress_bar=False)
        mcmc.run(jax.random.PRNGKey(self.seed), x_jax, y_jax)
        samples = mcmc.get_samples()

        intercept_std = np.asarray(samples["intercept"], dtype=float)
        sigma_std_draws = np.asarray(samples["sigma"], dtype=float)
        # Unstandardise: w_j = w_std_j * scale_y / scale_x_j; intercept absorbs both mean shifts;
        # sigma is a pure scale parameter so it only picks up scale_y. Re-derived independently
        # from y = mean_y + scale_y * y_std and y_std = intercept_std + x_std @ w_std + eps_std.
        sigma_draws = sigma_std_draws * safe_scale_y
        if n_parents > 0:
            w_std = np.asarray(samples["w"], dtype=float)  # (num_draws, n_parents)
            weight_draws = {
                name: w_std[:, i] * safe_scale_y / safe_scale_x[i] for i, name in enumerate(names)
            }
            intercept_draws = (
                mean_y
                + safe_scale_y * intercept_std
                - sum(weight_draws[name] * mean_x[i] for i, name in enumerate(names))
            )
        else:
            weight_draws = {}
            intercept_draws = mean_y + safe_scale_y * intercept_std

        posterior: dict[str, np.ndarray] = {
            "intercept": intercept_draws,
            **weight_draws,
            "sigma": sigma_draws,
        }

        intercept_mean = float(intercept_draws.mean())
        weight_means = {name: float(weight_draws[name].mean()) for name in names}
        sigma_mean = float(sigma_draws.mean())
        # Same affine mean map LinearGaussianFit inverts, with posterior means in place of OLS
        # coefficients -- so the residual round-trip identity holds here for the same reason.
        mean_fn = _affine_mean_fn(
            np.concatenate(([intercept_mean], [weight_means[name] for name in names]))
        )

        mechanism = LinearGaussianMechanism(names, weight_means, bias=intercept_mean)
        _attach_residual(mechanism, names, mean_fn)
        mechanism.posterior = posterior  # type: ignore[attr-defined]
        # dict[str, float], matching PoissonGLMFit's `coefficients` shape (Task 12) rather than a
        # positional vector -- a generic consumer reads coefficients["intercept"], not index 0.
        mechanism.coefficients = {"intercept": intercept_mean, **weight_means}  # type: ignore[attr-defined]

        predicted = mean_fn(raw)
        return FittedMechanism(
            mechanism=mechanism,
            noise=Normal(0.0, max(sigma_mean, 1e-6)),
            invertible=True,
            score=_r2(y, predicted),
        )

causalrl.protocols.CausalEnvProtocol

Bases: Protocol

The contract a simulator implements so every estimator/bound can drive it (§5.1).

Source code in src/causalrl/protocols.py

@runtime_checkable
class CausalEnvProtocol(Protocol):
    """The contract a simulator implements so every estimator/bound can drive it (§5.1)."""

    def sample(
        self, n: int, *, seed: int | None = None, regime: Regime | None = None
    ) -> TrajectoryLog: ...

    def do(
        self,
        interventions: Mapping[str, Any],
        n: int,
        *,
        seed: int | None = None,
        regime: Regime | None = None,
    ) -> TrajectoryLog: ...

    def noise_ledger(self) -> NoiseLedger | None: ...

causalrl.scm.continuous.mechanisms.ConditionalFlowMechanism

Bases: Module

Conditional normalizing-flow mechanism, invertible in its scalar exogenous noise.

V_i = flow(noise; parents) is a composition of n_blocks conditional affine maps u -> loc_k(parents) + scale_k(parents) * u (with scale_k = softplus(.) + min_scale > 0) interleaved with a fixed invertible LeakyReLU (negative_slope in (0, 1)), ending on an affine map. Each factor is strictly increasing in the scalar u, so the whole map is strictly monotone and the noise is recovered in closed form (:meth:invert) — exact abduction and exact continuous counterfactuals are therefore available, as for :class:LocationScaleMechanism but with a strictly more expressive, non-affine transform.

References: Rezende & Mohamed, Variational Inference with Normalizing Flows (2015); Papamakarios et al., Normalizing Flows for Probabilistic Modeling and Inference (JMLR 2021). Formula-level implementation; no third-party code is ported.

Source code in src/causalrl/scm/continuous/mechanisms.py

class ConditionalFlowMechanism(torch.nn.Module):
    """Conditional normalizing-flow mechanism, invertible in its scalar exogenous noise.

    ``V_i = flow(noise; parents)`` is a composition of ``n_blocks`` conditional affine maps
    ``u -> loc_k(parents) + scale_k(parents) * u`` (with ``scale_k = softplus(.) + min_scale > 0``)
    interleaved with a fixed invertible ``LeakyReLU`` (``negative_slope`` in ``(0, 1)``), ending on
    an affine map. Each factor is strictly increasing in the scalar ``u``, so the whole map is
    strictly monotone and the noise is recovered in closed form (:meth:`invert`) — exact abduction
    and exact continuous counterfactuals are therefore available, as for
    :class:`LocationScaleMechanism` but with a strictly more expressive, non-affine transform.

    References: Rezende & Mohamed, *Variational Inference with Normalizing Flows* (2015);
    Papamakarios et al., *Normalizing Flows for Probabilistic Modeling and Inference* (JMLR 2021).
    Formula-level implementation; no third-party code is ported.
    """

    def __init__(
        self,
        parents: list[str],
        hidden: tuple[int, ...] = (32,),
        *,
        n_blocks: int = 2,
        negative_slope: float = 0.1,
        min_scale: float = 1e-3,
    ) -> None:
        super().__init__()
        if n_blocks < 1:
            raise ValueError("n_blocks must be >= 1")
        if not 0.0 < negative_slope < 1.0:
            raise ValueError("negative_slope must lie in (0, 1) to stay invertible")
        self.parents = parents
        self.negative_slope = float(negative_slope)
        self.min_scale = float(min_scale)
        in_dim = max(len(parents), 1)
        self.locs = torch.nn.ModuleList(  # type: ignore[reportPrivateImportUsage]
            [build_mlp(in_dim, hidden, 1) for _ in range(n_blocks)]
        )
        self.log_scales = torch.nn.ModuleList(  # type: ignore[reportPrivateImportUsage]
            [build_mlp(in_dim, hidden, 1) for _ in range(n_blocks)]
        )

    def _affine(self, k: int, feats: Tensor) -> tuple[Tensor, Tensor]:
        loc = self.locs[k](feats).squeeze(-1)  # type: ignore[reportUnknownMemberType]
        raw = self.log_scales[k](feats).squeeze(-1)  # type: ignore[reportUnknownMemberType]
        scale = torch.nn.functional.softplus(raw) + self.min_scale  # type: ignore[reportPrivateImportUsage]
        return loc, scale

    def forward(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
        feats = _stack_parents(self.parents, parent_values, noise)
        n = len(self.locs)
        u = noise
        for k in range(n):
            loc, scale = self._affine(k, feats)
            u = loc + scale * u
            if k < n - 1:
                u = torch.nn.functional.leaky_relu(u, self.negative_slope)  # type: ignore[reportPrivateImportUsage]
        return u

    def invert(self, parent_values: dict[str, Tensor], observed: Tensor) -> Tensor:
        """Recover the exogenous noise exactly by undoing each factor in reverse order."""
        feats = _stack_parents(self.parents, parent_values, observed)
        n = len(self.locs)
        v = observed
        for k in reversed(range(n)):
            if k < n - 1:
                v = _leaky_relu_inverse(v, self.negative_slope)
            loc, scale = self._affine(k, feats)
            v = (v - loc) / scale
        return v

invert(parent_values, observed)

Recover the exogenous noise exactly by undoing each factor in reverse order.

Source code in src/causalrl/scm/continuous/mechanisms.py

def invert(self, parent_values: dict[str, Tensor], observed: Tensor) -> Tensor:
    """Recover the exogenous noise exactly by undoing each factor in reverse order."""
    feats = _stack_parents(self.parents, parent_values, observed)
    n = len(self.locs)
    v = observed
    for k in reversed(range(n)):
        if k < n - 1:
            v = _leaky_relu_inverse(v, self.negative_slope)
        loc, scale = self._affine(k, feats)
        v = (v - loc) / scale
    return v

causalrl.data.dataset.ConfoundedTrajectoryDataset

Immutable offline log plus empirical behavior statistics per (state, action).

Source code in src/causalrl/data/dataset.py

class ConfoundedTrajectoryDataset:
    """Immutable offline log plus empirical behavior statistics per (state, action)."""

    def __init__(self, transitions: list[Transition], n_states: int, n_actions: int) -> None:
        self._transitions = list(transitions)
        self.n_states = n_states
        self.n_actions = n_actions
        # counts[s][a] = number of times action a was taken in state s
        self._counts = [[0 for _ in range(n_actions)] for _ in range(n_states)]
        self._reward_sums = [[0.0 for _ in range(n_actions)] for _ in range(n_states)]
        for tr in self._transitions:
            self._counts[tr.state][tr.action] += 1
            self._reward_sums[tr.state][tr.action] += tr.reward

    def __len__(self) -> int:
        return len(self._transitions)

    @property
    def transitions(self) -> list[Transition]:
        return list(self._transitions)

    def _state_total(self, state: int) -> int:
        return sum(self._counts[state])

    def behavior_propensity(self, state: int, action: int) -> float:
        """Empirical P(action | state). 0.0 if the state was never visited."""
        total = self._state_total(state)
        if total == 0:
            return 0.0
        return self._counts[state][action] / total

    def mean_reward(self, state: int, action: int) -> float:
        """Empirical E[return | state, action]. 0.0 if (state, action) never logged."""
        n = self._counts[state][action]
        if n == 0:
            return 0.0
        return self._reward_sums[state][action] / n

    # -- LoggedDecisions: the off-policy certification interface (causalrl.data.logged) --

    def outcomes(self) -> Sequence[float]:
        """The logged return of each transition, in log order."""
        return [tr.reward for tr in self._transitions]

    def logging_propensities(self) -> Sequence[float]:
        """Empirical ``pi_behavior(a_i | s_i)`` for the action logged at each transition."""
        return [self.behavior_propensity(tr.state, tr.action) for tr in self._transitions]

    def matches(self, target_actions: Sequence[int]) -> Sequence[bool]:
        """Whether the target policy would have taken the logged action, per transition."""
        self._check_length(target_actions)
        return [
            int(a) == tr.action for a, tr in zip(target_actions, self._transitions, strict=True)
        ]

    def target_propensities(self, target_actions: Sequence[int]) -> Sequence[float] | None:
        """Empirical ``pi_behavior(a | s_i)`` for the action the target policy takes.

        Never ``None``: a tabular log counts every cell, so the propensity of an action nobody
        took in a state is known to be exactly zero.
        """
        self._check_length(target_actions)
        return [
            self.behavior_propensity(tr.state, int(a))
            for a, tr in zip(target_actions, self._transitions, strict=True)
        ]

    def positivity(self, target_actions: Sequence[int]) -> PositivityReport:
        """``(state, action)`` cells the target policy reaches that the logs never played.

        Always ``checkable``: a tabular log counts every cell, so an action nobody took in a state
        is known to have zero behaviour propensity rather than merely unobserved -- the
        distinction :class:`~causalrl.data.logged.PositivityReport` exists to preserve.
        """
        self._check_length(target_actions)
        gaps = sorted(
            {
                (tr.state, int(a))
                for a, tr in zip(target_actions, self._transitions, strict=True)
                if self.behavior_propensity(tr.state, int(a)) <= 0.0
            }
        )
        return PositivityReport(
            checkable=True, gaps=tuple(f"state {s}, action {a}" for s, a in gaps)
        )

    def _check_length(self, target_actions: Sequence[int]) -> None:
        if len(target_actions) != len(self._transitions):
            raise ValueError(
                f"target_actions has {len(target_actions)} entries but the log has "
                f"{len(self._transitions)} transitions: the certificate layer needs the action "
                "the target policy would take at each logged transition, one for one."
            )

behavior_propensity(state, action)

Empirical P(action | state). 0.0 if the state was never visited.

Source code in src/causalrl/data/dataset.py

def behavior_propensity(self, state: int, action: int) -> float:
    """Empirical P(action | state). 0.0 if the state was never visited."""
    total = self._state_total(state)
    if total == 0:
        return 0.0
    return self._counts[state][action] / total

mean_reward(state, action)

Empirical E[return | state, action]. 0.0 if (state, action) never logged.

Source code in src/causalrl/data/dataset.py

def mean_reward(self, state: int, action: int) -> float:
    """Empirical E[return | state, action]. 0.0 if (state, action) never logged."""
    n = self._counts[state][action]
    if n == 0:
        return 0.0
    return self._reward_sums[state][action] / n

outcomes()

The logged return of each transition, in log order.

Source code in src/causalrl/data/dataset.py

def outcomes(self) -> Sequence[float]:
    """The logged return of each transition, in log order."""
    return [tr.reward for tr in self._transitions]

logging_propensities()

Empirical pi_behavior(a_i | s_i) for the action logged at each transition.

Source code in src/causalrl/data/dataset.py

def logging_propensities(self) -> Sequence[float]:
    """Empirical ``pi_behavior(a_i | s_i)`` for the action logged at each transition."""
    return [self.behavior_propensity(tr.state, tr.action) for tr in self._transitions]

matches(target_actions)

Whether the target policy would have taken the logged action, per transition.

Source code in src/causalrl/data/dataset.py

def matches(self, target_actions: Sequence[int]) -> Sequence[bool]:
    """Whether the target policy would have taken the logged action, per transition."""
    self._check_length(target_actions)
    return [
        int(a) == tr.action for a, tr in zip(target_actions, self._transitions, strict=True)
    ]

target_propensities(target_actions)

Empirical pi_behavior(a | s_i) for the action the target policy takes.

Never None: a tabular log counts every cell, so the propensity of an action nobody took in a state is known to be exactly zero.

Source code in src/causalrl/data/dataset.py

def target_propensities(self, target_actions: Sequence[int]) -> Sequence[float] | None:
    """Empirical ``pi_behavior(a | s_i)`` for the action the target policy takes.

    Never ``None``: a tabular log counts every cell, so the propensity of an action nobody
    took in a state is known to be exactly zero.
    """
    self._check_length(target_actions)
    return [
        self.behavior_propensity(tr.state, int(a))
        for a, tr in zip(target_actions, self._transitions, strict=True)
    ]

positivity(target_actions)

(state, action) cells the target policy reaches that the logs never played.

Always checkable: a tabular log counts every cell, so an action nobody took in a state is known to have zero behaviour propensity rather than merely unobserved -- the distinction :class:~causalrl.data.logged.PositivityReport exists to preserve.

Source code in src/causalrl/data/dataset.py

def positivity(self, target_actions: Sequence[int]) -> PositivityReport:
    """``(state, action)`` cells the target policy reaches that the logs never played.

    Always ``checkable``: a tabular log counts every cell, so an action nobody took in a state
    is known to have zero behaviour propensity rather than merely unobserved -- the
    distinction :class:`~causalrl.data.logged.PositivityReport` exists to preserve.
    """
    self._check_length(target_actions)
    gaps = sorted(
        {
            (tr.state, int(a))
            for a, tr in zip(target_actions, self._transitions, strict=True)
            if self.behavior_propensity(tr.state, int(a)) <= 0.0
        }
    )
    return PositivityReport(
        checkable=True, gaps=tuple(f"state {s}, action {a}" for s, a in gaps)
    )

causalrl.data.logged.FeatureDecisionLog

Bases: Generic[ActionT]

A :class:LoggedDecisions over feature-space states and any action type.

Build one from an agent's own logs when its states are vectors or its actions are :data:~causalrl.Intervention assignments. Actions are compared through :func:action_key, so a policy that returns {"dose": 2.0} matches a logged {"dose": 2.0} regardless of mapping identity or key order.

behavior_propensity is optional and exists solely to make positivity checkable: given pi_behavior(a | s) it can ask whether the target policy's action had support at each logged state. Without it the log is still fully usable -- every weight the certificate layer forms uses the logged action's own propensity, which is always present -- but :meth:positivity reports that it could not check rather than that it found nothing.

Source code in src/causalrl/data/logged.py

class FeatureDecisionLog(Generic[ActionT]):
    """A :class:`LoggedDecisions` over feature-space states and any action type.

    Build one from an agent's own logs when its states are vectors or its actions are
    :data:`~causalrl.Intervention` assignments. Actions are compared through :func:`action_key`,
    so a policy that returns ``{"dose": 2.0}`` matches a logged ``{"dose": 2.0}`` regardless of
    mapping identity or key order.

    ``behavior_propensity`` is optional and exists solely to make positivity *checkable*: given
    ``pi_behavior(a | s)`` it can ask whether the target policy's action had support at each logged
    state. Without it the log is still fully usable -- every weight the certificate layer forms
    uses the logged action's own propensity, which is always present -- but
    :meth:`positivity` reports that it could not check rather than that it found nothing.
    """

    def __init__(
        self,
        decisions: Sequence[LoggedDecision[ActionT]],
        *,
        behavior_propensity: Callable[[FloatArray, ActionT], float] | None = None,
    ) -> None:
        if not decisions:
            raise ValueError(
                "FeatureDecisionLog needs at least one logged decision: every quantity the "
                "certificate layer computes is an average over the log, and there is no "
                "off-policy claim to make from an empty one."
            )
        self._decisions = tuple(decisions)
        self._behavior_propensity = behavior_propensity

    def __len__(self) -> int:
        return len(self._decisions)

    @property
    def decisions(self) -> tuple[LoggedDecision[ActionT], ...]:
        """The logged decisions, in order."""
        return self._decisions

    def outcomes(self) -> Sequence[float]:
        return [d.reward for d in self._decisions]

    def logging_propensities(self) -> Sequence[float]:
        return [d.propensity for d in self._decisions]

    def matches(self, target_actions: Sequence[ActionT]) -> Sequence[bool]:
        self._check_length(target_actions)
        return [
            action_key(a) == action_key(d.action)
            for a, d in zip(target_actions, self._decisions, strict=True)
        ]

    def target_propensities(self, target_actions: Sequence[ActionT]) -> Sequence[float] | None:
        self._check_length(target_actions)
        if self._behavior_propensity is None:
            return None
        return [
            float(self._behavior_propensity(d.state, a))
            for a, d in zip(target_actions, self._decisions, strict=True)
        ]

    def positivity(self, target_actions: Sequence[ActionT]) -> PositivityReport:
        propensities = self.target_propensities(target_actions)
        if propensities is None:
            return PositivityReport(checkable=False)
        gaps = tuple(
            sorted(
                {
                    f"action {target_actions[i]!r} unsupported at logged state {i}"
                    for i, e in enumerate(propensities)
                    if e <= 0.0
                }
            )
        )
        return PositivityReport(checkable=True, gaps=gaps)

    def _check_length(self, target_actions: Sequence[ActionT]) -> None:
        if len(target_actions) != len(self._decisions):
            raise ValueError(
                f"target_actions has {len(target_actions)} entries but the log has "
                f"{len(self._decisions)} decisions: the certificate layer needs the action the "
                "target policy would take at each logged decision, one for one."
            )

decisions property

The logged decisions, in order.

causalrl.data.logged.LoggedDecision dataclass

Bases: Generic[ActionT]

One logged decision in feature space, with the propensity that produced it.

The feature-space counterpart of :class:~causalrl.data.dataset.Transition for the certification path rather than the planning path: it carries no successor state, because nothing in off-policy certification of a terminal return looks at one. propensity is supplied rather than counted -- in feature space there are no cells to count.

Source code in src/causalrl/data/logged.py

@dataclass(frozen=True)
class LoggedDecision(Generic[ActionT]):
    """One logged decision in feature space, with the propensity that produced it.

    The feature-space counterpart of :class:`~causalrl.data.dataset.Transition` for the
    *certification* path rather than the planning path: it carries no successor state, because
    nothing in off-policy certification of a terminal return looks at one. ``propensity`` is
    supplied rather than counted -- in feature space there are no cells to count.
    """

    state: FloatArray
    action: ActionT
    reward: float
    propensity: float

    def __post_init__(self) -> None:
        object.__setattr__(self, "state", np.asarray(self.state, dtype=np.float64).reshape(-1))
        if not 0.0 < self.propensity <= 1.0:
            raise ValueError(
                f"propensity={self.propensity} must lie in (0, 1]: it is the logging policy's "
                "probability of the action it actually took, which is positive by construction "
                "for any action that appears in the log. A zero here would divide by zero in "
                "every importance weight downstream."
            )

causalrl.data.logged.LoggedDecisions

Bases: Protocol[ActionT_contra]

The off-policy certification interface: outcomes, logging propensities, target matches.

Deliberately smaller than any concrete log. Nothing here mentions a state representation, because none of the certificate machinery needs one -- the marginal sensitivity model reweights units, and a unit is identified by its position in the log.

Source code in src/causalrl/data/logged.py

class LoggedDecisions(Protocol[ActionT_contra]):
    """The off-policy certification interface: outcomes, logging propensities, target matches.

    Deliberately smaller than any concrete log. Nothing here mentions a state representation,
    because none of the certificate machinery needs one -- the marginal sensitivity model reweights
    *units*, and a unit is identified by its position in the log.
    """

    def __len__(self) -> int:
        """Number of logged decisions."""
        ...

    def outcomes(self) -> Sequence[float]:
        """The logged return of each decision, in log order."""
        ...

    def logging_propensities(self) -> Sequence[float]:
        """``pi_behavior(a_i | s_i)`` for the action actually logged at each decision."""
        ...

    def matches(self, target_actions: Sequence[ActionT_contra]) -> Sequence[bool]:
        """Whether the target policy would have taken the logged action, per decision."""
        ...

    def target_propensities(
        self, target_actions: Sequence[ActionT_contra]
    ) -> Sequence[float] | None:
        """``pi_behavior(pi_target(s_i) | s_i)`` per decision, or ``None`` if unknowable here.

        Note this is the behaviour propensity of the action the *target* policy would take, which
        is a strictly harder thing to know than :meth:`logging_propensities`: it asks about an
        action that may never have been played. A tabular log counts cells and can answer; a
        feature-space log without a behaviour model returns ``None``. The conformal band needs it
        to bound the likelihood ratio at a fresh test point, and ``None`` correctly forces that
        bound to infinity rather than letting an unknown pass as a small number.
        """
        ...

    def positivity(self, target_actions: Sequence[ActionT_contra]) -> PositivityReport:
        """Where the target policy leaves the logs' support, if that is knowable at all.

        Must agree with :meth:`target_propensities`: ``checkable`` exactly when that returns a
        sequence, and non-empty ``gaps`` exactly where its entries are non-positive. The two are
        separate methods because one supplies a bound and the other a human-readable diagnostic,
        and collapsing them would cost the cell labels that make a positivity failure actionable.
        """
        ...

__len__()

Number of logged decisions.

Source code in src/causalrl/data/logged.py

def __len__(self) -> int:
    """Number of logged decisions."""
    ...

outcomes()

The logged return of each decision, in log order.

Source code in src/causalrl/data/logged.py

def outcomes(self) -> Sequence[float]:
    """The logged return of each decision, in log order."""
    ...

logging_propensities()

pi_behavior(a_i | s_i) for the action actually logged at each decision.

Source code in src/causalrl/data/logged.py

def logging_propensities(self) -> Sequence[float]:
    """``pi_behavior(a_i | s_i)`` for the action actually logged at each decision."""
    ...

matches(target_actions)

Whether the target policy would have taken the logged action, per decision.

Source code in src/causalrl/data/logged.py

def matches(self, target_actions: Sequence[ActionT_contra]) -> Sequence[bool]:
    """Whether the target policy would have taken the logged action, per decision."""
    ...

target_propensities(target_actions)

pi_behavior(pi_target(s_i) | s_i) per decision, or None if unknowable here.

Note this is the behaviour propensity of the action the target policy would take, which is a strictly harder thing to know than :meth:logging_propensities: it asks about an action that may never have been played. A tabular log counts cells and can answer; a feature-space log without a behaviour model returns None. The conformal band needs it to bound the likelihood ratio at a fresh test point, and None correctly forces that bound to infinity rather than letting an unknown pass as a small number.

Source code in src/causalrl/data/logged.py

def target_propensities(
    self, target_actions: Sequence[ActionT_contra]
) -> Sequence[float] | None:
    """``pi_behavior(pi_target(s_i) | s_i)`` per decision, or ``None`` if unknowable here.

    Note this is the behaviour propensity of the action the *target* policy would take, which
    is a strictly harder thing to know than :meth:`logging_propensities`: it asks about an
    action that may never have been played. A tabular log counts cells and can answer; a
    feature-space log without a behaviour model returns ``None``. The conformal band needs it
    to bound the likelihood ratio at a fresh test point, and ``None`` correctly forces that
    bound to infinity rather than letting an unknown pass as a small number.
    """
    ...

positivity(target_actions)

Where the target policy leaves the logs' support, if that is knowable at all.

Must agree with :meth:target_propensities: checkable exactly when that returns a sequence, and non-empty gaps exactly where its entries are non-positive. The two are separate methods because one supplies a bound and the other a human-readable diagnostic, and collapsing them would cost the cell labels that make a positivity failure actionable.

Source code in src/causalrl/data/logged.py

def positivity(self, target_actions: Sequence[ActionT_contra]) -> PositivityReport:
    """Where the target policy leaves the logs' support, if that is knowable at all.

    Must agree with :meth:`target_propensities`: ``checkable`` exactly when that returns a
    sequence, and non-empty ``gaps`` exactly where its entries are non-positive. The two are
    separate methods because one supplies a bound and the other a human-readable diagnostic,
    and collapsing them would cost the cell labels that make a positivity failure actionable.
    """
    ...

causalrl.data.logged.PositivityReport dataclass

Whether a target policy strays outside the logs' support -- and whether that is knowable.

checkable is the load-bearing field. A tabular log counts every (state, action) cell, so it can answer the question and checkable is True; a feature-space log without a behaviour-propensity model cannot answer it at all, and reports False with empty gaps rather than an empty gaps that reads as a clean bill of health. Consumers must not treat the two the same: the first licenses a checkable positivity assumption, the second a hedge.

Source code in src/causalrl/data/logged.py

@dataclass(frozen=True)
class PositivityReport:
    """Whether a target policy strays outside the logs' support -- and whether that is knowable.

    ``checkable`` is the load-bearing field. A tabular log counts every ``(state, action)`` cell,
    so it can answer the question and ``checkable`` is ``True``; a feature-space log without a
    behaviour-propensity model cannot answer it at all, and reports ``False`` with empty ``gaps``
    rather than an empty ``gaps`` that reads as a clean bill of health. Consumers must not treat
    the two the same: the first licenses a checkable positivity assumption, the second a hedge.
    """

    checkable: bool
    gaps: tuple[str, ...] = ()

    @property
    def violated(self) -> bool:
        """Whether a gap was actually found. ``False`` when the check could not be run."""
        return bool(self.gaps)

violated property

Whether a gap was actually found. False when the check could not be run.

causalrl.protocols.DictNoiseLedger

A white-box NoiseLedger backed by registered exogenous draws (§5.5).

The black-box posterior path (amortized abduction) arrives in Phase 1; it raises here.

Source code in src/causalrl/protocols.py

class DictNoiseLedger:
    """A white-box `NoiseLedger` backed by registered exogenous draws (§5.5).

    The black-box ``posterior`` path (amortized abduction) arrives in Phase 1; it raises here.
    """

    def __init__(self, draws: Mapping[tuple[int, int], Mapping[str, Any]]) -> None:
        self._draws = dict(draws)

    def draws(self, episode_id: int, entity_id: int) -> Mapping[str, Any] | None:
        return self._draws.get((episode_id, entity_id))

    def posterior(self, evidence: TrajectoryLog) -> NoisePosterior:
        raise NotImplementedError("black-box noise posterior arrives in Phase 1 (abduction)")

causalrl.scm.scm.ExogenousPosterior

Retained exogenous values from abduction, ready for prediction under interventions.

Holds the source SCM and a dict of exogenous tensors (the abducted units). Call :meth:predict to evaluate the model — optionally mutilated by do — on them. Abduct once, predict under many interventions (the efficient counterfactual pattern).

Source code in src/causalrl/scm/scm.py

class ExogenousPosterior:
    """Retained exogenous values from abduction, ready for prediction under interventions.

    Holds the source SCM and a dict of exogenous tensors (the abducted units). Call
    :meth:`predict` to evaluate the model — optionally mutilated by ``do`` — on them. Abduct
    once, predict under many interventions (the efficient counterfactual pattern).
    """

    def __init__(self, scm: StructuralCausalModel, noise: dict[str, Tensor]) -> None:
        self._scm = scm
        self.noise = noise

    def __len__(self) -> int:
        return int(next(iter(self.noise.values())).shape[0]) if self.noise else 0

    def predict(self, *, do: Mapping[str, Value] | None = None) -> dict[str, Tensor]:
        """Evaluate the (optionally do-mutilated) model on the retained exogenous units."""
        model = self._scm.do(do) if do else self._scm
        return model._evaluate(self.noise)  # type: ignore[reportPrivateUsage]

predict(*, do=None)

Evaluate the (optionally do-mutilated) model on the retained exogenous units.

Source code in src/causalrl/scm/scm.py

def predict(self, *, do: Mapping[str, Value] | None = None) -> dict[str, Tensor]:
    """Evaluate the (optionally do-mutilated) model on the retained exogenous units."""
    model = self._scm.do(do) if do else self._scm
    return model._evaluate(self.noise)  # type: ignore[reportPrivateUsage]

causalrl.scm.fit.FitReport

Bases: NamedTuple

Per-node provenance for a fitted SCM -- see :class:NodeFit for what each field reports.

Source code in src/causalrl/scm/fit.py

class FitReport(NamedTuple):
    """Per-node provenance for a fitted SCM -- see :class:`NodeFit` for what each field reports."""

    nodes: tuple[NodeFit, ...]
    n_samples: int

    def summary(self) -> str:
        lines = [f"FitReport(n={self.n_samples})"]
        for fit in self.nodes:
            parents = ", ".join(fit.parents) or "-"
            marker = " PINNED" if fit.pinned else ""
            lines.append(
                f"  {fit.node}: family={fit.family} parents=[{parents}] "
                f"holdout={fit.holdout_score:.3f} invertible={fit.invertible}{marker}"
            )
        return "\n".join(lines)

    @property
    def pinned_nodes(self) -> tuple[str, ...]:
        """Nodes whose equation was supplied rather than learned, in report order."""
        return tuple(fit.node for fit in self.nodes if fit.pinned)

pinned_nodes property

Nodes whose equation was supplied rather than learned, in report order.

causalrl.scm.mechanisms.FunctionalMechanism

Wrap an arbitrary callable f(parent_values, noise) -> Tensor.

Source code in src/causalrl/scm/mechanisms.py

class FunctionalMechanism:
    """Wrap an arbitrary callable f(parent_values, noise) -> Tensor."""

    def __init__(
        self, parents: list[str], fn: Callable[[dict[str, Tensor], Tensor], Tensor]
    ) -> None:
        self.parents = parents
        self._fn = fn

    def __call__(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
        return self._fn(parent_values, noise)

causalrl.scm.fitters.LinearGaussianFit

Continuous node: closed-form OLS mean with Gaussian residual noise.

The noise enters additively and is therefore recoverable from (parents, value), so counterfactuals at this node are identified — invertible=True.

Source code in src/causalrl/scm/fitters.py

class LinearGaussianFit:
    """Continuous node: closed-form OLS mean with Gaussian residual noise.

    The noise enters additively and is therefore recoverable from (parents, value), so
    counterfactuals at this node are identified — ``invertible=True``.
    """

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        y = np.asarray(child, dtype=float)
        names = sorted(parents)
        design = _design(parents, names, len(y))
        coefficients, *_ = np.linalg.lstsq(design, y, rcond=None)
        predicted = design @ coefficients
        residual = y - predicted
        sigma = float(residual.std())
        weights = {name: float(coefficients[i + 1]) for i, name in enumerate(names)}
        mechanism = LinearGaussianMechanism(names, weights, bias=float(coefficients[0]))
        _attach_residual(mechanism, names, _affine_mean_fn(coefficients))
        return FittedMechanism(
            mechanism=mechanism,
            noise=Normal(0.0, max(sigma, 1e-6)),
            invertible=True,
            score=_r2(y, predicted),
        )

causalrl.scm.mechanisms.LinearGaussianMechanism

V_i = sum_j w_j * parent_j + bias + noise.

Source code in src/causalrl/scm/mechanisms.py

class LinearGaussianMechanism:
    """V_i = sum_j w_j * parent_j + bias + noise."""

    def __init__(self, parents: list[str], weights: dict[str, float], bias: float = 0.0) -> None:
        self.parents = parents
        self._weights = weights
        self._bias = bias

    def __call__(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
        # torch.tensor / torch.cat are public API; pyright strict flags them as not
        # re-exported by torch stubs — suppress only these two occurrences.
        out: Tensor = torch.tensor(self._bias, dtype=noise.dtype) + noise  # type: ignore[reportPrivateImportUsage]
        for p in self.parents:
            out = out + self._weights[p] * parent_values[p]
        return out

causalrl.scm.continuous.mechanisms.LocationScaleMechanism

Bases: Module

Invertible location-scale mechanism V_i = loc(parents) + scale(parents) * noise.

scale = softplus(raw_scale(parents)) + min_scale > 0, so the map is invertible in noise: noise = (V_i - loc) / scale. Exact abduction is therefore available via :func:causalrl.scm.continuous.abduction.abduct_location_scale.

Source code in src/causalrl/scm/continuous/mechanisms.py

class LocationScaleMechanism(torch.nn.Module):
    """Invertible location-scale mechanism ``V_i = loc(parents) + scale(parents) * noise``.

    ``scale = softplus(raw_scale(parents)) + min_scale > 0``, so the map is invertible in ``noise``:
    ``noise = (V_i - loc) / scale``. Exact abduction is therefore available via
    :func:`causalrl.scm.continuous.abduction.abduct_location_scale`.
    """

    def __init__(
        self, parents: list[str], hidden: tuple[int, ...] = (32,), min_scale: float = 1e-3
    ) -> None:
        super().__init__()
        self.parents = parents
        self.min_scale = float(min_scale)
        self.loc_net = build_mlp(max(len(parents), 1), hidden, 1)
        self.scale_net = build_mlp(max(len(parents), 1), hidden, 1)

    def _loc_scale(self, feats: Tensor) -> tuple[Tensor, Tensor]:
        loc = self.loc_net(feats).squeeze(-1)  # type: ignore[reportUnknownMemberType]
        raw = self.scale_net(feats).squeeze(-1)  # type: ignore[reportUnknownMemberType]
        scale = torch.nn.functional.softplus(raw) + self.min_scale  # type: ignore[reportPrivateImportUsage]
        return loc, scale

    def forward(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
        loc, scale = self._loc_scale(_stack_parents(self.parents, parent_values, noise))
        return loc + scale * noise

    def invert(self, parent_values: dict[str, Tensor], observed: Tensor) -> Tensor:
        """Recover the exogenous noise exactly: ``(observed - loc) / scale``."""
        loc, scale = self._loc_scale(_stack_parents(self.parents, parent_values, observed))
        return (observed - loc) / scale

invert(parent_values, observed)

Recover the exogenous noise exactly: (observed - loc) / scale.

Source code in src/causalrl/scm/continuous/mechanisms.py

def invert(self, parent_values: dict[str, Tensor], observed: Tensor) -> Tensor:
    """Recover the exogenous noise exactly: ``(observed - loc) / scale``."""
    loc, scale = self._loc_scale(_stack_parents(self.parents, parent_values, observed))
    return (observed - loc) / scale

causalrl.scm.continuous.mechanisms.MLPMechanism

Bases: Module

Additive-noise neural mechanism V_i = MLP(parents) + noise.

Source code in src/causalrl/scm/continuous/mechanisms.py

class MLPMechanism(torch.nn.Module):
    """Additive-noise neural mechanism ``V_i = MLP(parents) + noise``."""

    def __init__(self, parents: list[str], hidden: tuple[int, ...] = (32, 32)) -> None:
        super().__init__()
        self.parents = parents
        self.net = build_mlp(max(len(parents), 1), hidden, 1)

    def forward(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
        feats = _stack_parents(self.parents, parent_values, noise)
        return self.net(feats).squeeze(-1) + noise  # type: ignore[reportUnknownMemberType]

causalrl.scm.mechanisms.Mechanism

Bases: Protocol

A structural equation V_i = f_i(parents, noise).

Source code in src/causalrl/scm/mechanisms.py

@runtime_checkable
class Mechanism(Protocol):
    """A structural equation V_i = f_i(parents, noise)."""

    parents: list[str]

    def __call__(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor: ...

causalrl.scm.continuous.nuts.NUTSNoisePosterior

A NoisePosterior backed by NUTS draws of a mechanism's per-unit exogenous noise.

mean / std summarise the posterior per unit; sample(n) returns n posterior draws (each a full per-unit noise vector), tiling if more are requested than were drawn. Pure NumPy — it holds no JAX state, so it composes with the torch/NumPy stack unchanged.

Source code in src/causalrl/scm/continuous/nuts.py

class NUTSNoisePosterior:
    """A ``NoisePosterior`` backed by NUTS draws of a mechanism's per-unit exogenous noise.

    ``mean`` / ``std`` summarise the posterior per unit; ``sample(n)`` returns ``n`` posterior draws
    (each a full per-unit noise vector), tiling if more are requested than were drawn. Pure NumPy —
    it holds no JAX state, so it composes with the torch/NumPy stack unchanged.
    """

    def __init__(self, samples: FloatArray, *, name: str = "U") -> None:
        self._samples = np.asarray(samples, dtype=np.float64)  # (num_samples, n_units)
        self.name = name
        self.mean: FloatArray = self._samples.mean(axis=0)
        self.std: FloatArray = self._samples.std(axis=0)

    def sample(self, n: int, *, seed: int | None = None) -> Mapping[str, Any]:
        s = self._samples
        if n <= s.shape[0]:
            return {self.name: s[:n]}
        reps = (n + s.shape[0] - 1) // s.shape[0]
        return {self.name: np.tile(s, (reps, 1))[:n]}

causalrl.scm.fitters.NeuralFit

Continuous node: V = net(parents) + U, an MLP mean with Gaussian residual noise.

Fires :class:causalrl.scm.mechanisms.NeuralMechanism — the neural causal model primitive. The noise head is additive, so the mechanism stays invertible and counterfactuals at this node remain identified; a general (non-additive) net would not.

Source code in src/causalrl/scm/fitters.py

class NeuralFit:
    """Continuous node: ``V = net(parents) + U``, an MLP mean with Gaussian residual noise.

    Fires :class:`causalrl.scm.mechanisms.NeuralMechanism` — the neural causal model primitive.
    The noise head is additive, so the mechanism stays invertible and counterfactuals at this node
    remain identified; a general (non-additive) net would not.
    """

    def __init__(
        self, *, hidden: int = 32, epochs: int = 200, lr: float = 0.01, seed: int = 0
    ) -> None:
        self.hidden = hidden
        self.epochs = epochs
        self.lr = lr
        self.seed = seed

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        y = np.asarray(child, dtype=float)
        names = sorted(parents)
        if not names:
            return LinearGaussianFit().fit({}, y)
        x = torch.tensor(  # type: ignore[reportPrivateImportUsage]
            np.column_stack([np.asarray(parents[n], dtype=float) for n in names]),
            dtype=torch.float32,  # type: ignore[reportPrivateImportUsage]
        )
        target = torch.tensor(y, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
        # Seed INSIDE a forked RNG: the fit stays deterministic for a given self.seed without
        # reseeding the caller's process-global stream, which would silently make every
        # subsequent torch draw in their program a function of when they happened to fit a
        # mechanism. Same pattern, and same reason, as StructuralCausalModel._sample_exogenous.
        with torch.random.fork_rng():  # type: ignore[reportUnknownMemberType]
            torch.manual_seed(self.seed)  # type: ignore[reportUnknownMemberType]
            net = torch.nn.Sequential(
                torch.nn.Linear(len(names), self.hidden),
                torch.nn.Tanh(),
                torch.nn.Linear(self.hidden, 1),
            )
            optimizer = torch.optim.Adam(net.parameters(), lr=self.lr)
            for _ in range(self.epochs):
                optimizer.zero_grad()
                loss = torch.nn.functional.mse_loss(net(x).squeeze(-1), target)
                loss.backward()  # type: ignore[reportUnknownMemberType]
                optimizer.step()  # type: ignore[reportUnknownMemberType]
        # Fitting is done: freeze the net so every later call (mechanism(...), residual(...),
        # forward sampling through StructuralCausalModel._evaluate) returns a plain tensor
        # instead of silently re-building a live autograd graph through the trained weights.
        for parameter in net.parameters():
            parameter.requires_grad_(False)
        predicted = net(x).squeeze(-1)
        residual = target - predicted
        sigma = float(residual.std())

        # NeuralMechanism concatenates [parents..., noise]; wrap so the noise stays additive.
        mechanism = NeuralMechanism(names, _AdditiveHead(net))

        def mean_fn(columns: np.ndarray) -> np.ndarray:
            with torch.no_grad():
                batch = torch.tensor(columns, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
                return net(batch).squeeze(-1).numpy()

        _attach_residual(mechanism, names, mean_fn)
        return FittedMechanism(
            mechanism=mechanism,
            noise=Normal(0.0, max(sigma, 1e-6)),
            invertible=True,
            score=_r2(y, predicted.detach().numpy()),
        )

causalrl.scm.mechanisms.NeuralMechanism

Bases: Module

V_i = net([parents, noise]). Makes the SCM a neural causal model (NCM).

Source code in src/causalrl/scm/mechanisms.py

class NeuralMechanism(torch.nn.Module):
    """V_i = net([parents, noise]). Makes the SCM a neural causal model (NCM)."""

    def __init__(self, parents: list[str], net: torch.nn.Module) -> None:
        super().__init__()
        self.parents = parents
        self.net = net

    def forward(self, parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
        cols = [parent_values[p].reshape(-1, 1) for p in self.parents] + [noise.reshape(-1, 1)]
        return self.net(torch.cat(cols, dim=1)).squeeze(-1)  # type: ignore[reportPrivateImportUsage]

causalrl.scm.fit.NodeFit

Bases: NamedTuple

What was fitted at one node, and what that licenses.

holdout_score evaluates the DEPLOYED (train-fitted) mechanism against data the fit never saw -- mean log-likelihood for a discrete (non-invertible) mechanism, R^2 for a continuous (invertible, additive-noise) one -- via :func:causalrl.scm.fitters.evaluate_holdout. It sits on the same scale as each family's in-sample score, but is a genuine out-of-sample number: an overfit deployed mechanism scores worse here even though its in-sample score looks fine.

pinned marks a node whose equation was SUPPLIED rather than learned (see :class:causalrl.scm.fitters.PinnedMechanism). It changes how holdout_score should be read: for a learned node the score measures how well a family fitted, and for a pinned one it measures whether the equation the caller asserted survives contact with held-out data.

Source code in src/causalrl/scm/fit.py

class NodeFit(NamedTuple):
    """What was fitted at one node, and what that licenses.

    ``holdout_score`` evaluates the DEPLOYED (train-fitted) mechanism against data the fit never
    saw -- mean log-likelihood for a discrete (non-invertible) mechanism, R^2 for a continuous
    (invertible, additive-noise) one -- via :func:`causalrl.scm.fitters.evaluate_holdout`. It sits
    on the same scale as each family's in-sample ``score``, but is a genuine out-of-sample number:
    an overfit deployed mechanism scores worse here even though its in-sample ``score`` looks fine.

    ``pinned`` marks a node whose equation was SUPPLIED rather than learned (see
    :class:`causalrl.scm.fitters.PinnedMechanism`). It changes how ``holdout_score`` should be
    read: for a learned node the score measures how well a family fitted, and for a pinned one it
    measures whether the equation the caller asserted survives contact with held-out data.
    """

    node: str
    family: str
    parents: tuple[str, ...]
    holdout_score: float
    invertible: bool
    pinned: bool = False

causalrl.protocols.NoiseLedger

Bases: Protocol

Explicit noise handling for Level-3 queries (§5.5, I6).

Either registered white-box draws (draws) or a posterior obtained by abduction (posterior, Phase 1). No counterfactual API may silently assume access to noise.

Source code in src/causalrl/protocols.py

@runtime_checkable
class NoiseLedger(Protocol):
    """Explicit noise handling for Level-3 queries (§5.5, I6).

    Either registered white-box draws (``draws``) or a posterior obtained by abduction
    (``posterior``, Phase 1). No counterfactual API may silently assume access to noise.
    """

    def draws(self, episode_id: int, entity_id: int) -> Mapping[str, Any] | None: ...

    def posterior(self, evidence: TrajectoryLog) -> NoisePosterior: ...

causalrl.protocols.NoisePosterior

Bases: Protocol

A posterior over exogenous noise (the Phase-1 abduction output).

Source code in src/causalrl/protocols.py

@runtime_checkable
class NoisePosterior(Protocol):
    """A posterior over exogenous noise (the Phase-1 abduction output)."""

    def sample(self, n: int, *, seed: int | None = None) -> Mapping[str, Any]: ...

causalrl.scm.fitters.PoissonGLMFit

Count node: a Poisson GLM with a log link, fitted by IRLS.

log E[V | parents] = intercept + w · parents. Linear in the parents where a conditional probability table is exponential in them, which is what makes lag-unrolled count data (6-12 lagged parents) tractable at all — and it keeps the point-process structure a table discards.

The mechanism is V = F⁻¹(U | parents) with U ~ Uniform(0, 1), the same canonical construction :class:TabularCPT uses. A Poisson draw is not recoverable from (parents, count), so invertible=False: counterfactuals at this node are genuinely unidentified and the L3 guard correctly refuses them. log_prob is attached so :func:evaluate_holdout scores this family by mean held-out log-likelihood rather than R².

Implements the standard IRLS/Fisher-scoring fit for a Poisson GLM (McCullagh & Nelder, Generalized Linear Models, 2nd ed., §2.5); no external code is ported.

Source code in src/causalrl/scm/fitters.py

class PoissonGLMFit:
    """Count node: a Poisson GLM with a log link, fitted by IRLS.

    ``log E[V | parents] = intercept + w · parents``. Linear in the parents where a conditional
    probability table is exponential in them, which is what makes lag-unrolled count data (6-12
    lagged parents) tractable at all — and it keeps the point-process structure a table discards.

    The mechanism is ``V = F⁻¹(U | parents)`` with ``U ~ Uniform(0, 1)``, the same canonical
    construction :class:`TabularCPT` uses. A Poisson draw is not recoverable from
    ``(parents, count)``, so ``invertible=False``: counterfactuals at this node are genuinely
    unidentified and the L3 guard correctly refuses them. ``log_prob`` is attached so
    :func:`evaluate_holdout` scores this family by mean held-out log-likelihood rather than R².

    Implements the standard IRLS/Fisher-scoring fit for a Poisson GLM (McCullagh & Nelder,
    *Generalized Linear Models*, 2nd ed., §2.5); no external code is ported.
    """

    def __init__(self, max_iter: int = 50, tol: float = 1e-8) -> None:
        self.max_iter = max_iter
        self.tol = tol

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        y = np.asarray(child, dtype=float)
        names = sorted(parents)
        design = _design(parents, names, len(y))
        beta = _poisson_irls(design, y, self.max_iter, self.tol)
        intercept = float(beta[0])
        weights = {name: float(beta[i + 1]) for i, name in enumerate(names)}
        coefficients = {"intercept": intercept, **weights}

        def rate(parent_values: dict[str, Tensor], n: int) -> Tensor:
            eta = torch.full((n,), intercept, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
            for name in names:
                eta = eta + weights[name] * parent_values[name].reshape(-1).float()
            return eta.clamp(-30.0, 30.0).exp()

        def mechanism(parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
            flat_noise = noise.reshape(-1).float()
            n = flat_noise.shape[0]
            if n == 0:
                return torch.zeros(0)  # type: ignore[reportPrivateImportUsage]
            lam = rate(parent_values, n)
            # Cap the support so the vectorised CDF is a finite (n, cap+1) table. A draw landing
            # at the cap is astronomically unlikely at these rates (the Poisson tail decays
            # super-exponentially), so clamping to it below is silently correct, not an
            # approximation worth raising over.
            lam_max = float(lam.max())
            cap = max(20, int(np.ceil(lam_max + 10.0 * np.sqrt(lam_max))))
            cells = n * (cap + 1)
            if cells > _MAX_POISSON_CELLS:
                raise ValueError(
                    f"PoissonGLMFit's sampler would need a support table of {cells} cells "
                    f"(n={n} rows x {cap + 1} columns, lambda up to {lam_max:.3g} in this batch) "
                    f"-- above _MAX_POISSON_CELLS={_MAX_POISSON_CELLS}. This table is rebuilt on "
                    f"every mechanism() call and its cost is the PRODUCT of the batch size and the "
                    f"per-row support width, so either factor alone -- a large batch or a large "
                    f"lambda -- can trip it; the numbers above show which one did here. The cap is "
                    f"a fixed module constant with no parameter to raise it, so the fix is on the "
                    f"calling side: a lambda this large usually means an unstable upstream "
                    f"mechanism (e.g. a lag chain whose coefficients amplify rather than decay) "
                    f"rather than real point-process data, so check the fitted coefficients "
                    f"feeding this node; if lambda is ordinary, ask for fewer rows per call."
                )
            counts = torch.arange(cap + 1, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
            log_pmf = (
                counts.unsqueeze(0) * lam.log().unsqueeze(1)
                - lam.unsqueeze(1)
                - (counts + 1.0).lgamma().unsqueeze(0)
            )
            cum = log_pmf.exp().cumsum(dim=1)
            picked = (flat_noise.unsqueeze(1) > cum).sum(dim=1)
            return picked.clamp(max=cap).float()

        def log_prob(parent_values: dict[str, Tensor], value: Tensor) -> Tensor:
            """Per-sample log P(value | parents) under the fitted Poisson rate -- what
            evaluate_holdout uses to score this non-invertible family, and what mechanism()'s
            inverse-CDF table is itself built from (one Poisson log-pmf, not two)."""
            flat_value = value.reshape(-1).float()
            lam = rate(parent_values, flat_value.shape[0])
            return flat_value * lam.log() - lam - (flat_value + 1.0).lgamma()

        fitted_mechanism = FunctionalMechanism(names, mechanism)
        fitted_mechanism.coefficients = coefficients  # type: ignore[attr-defined]
        fitted_mechanism.log_prob = log_prob  # type: ignore[attr-defined]

        # In-sample mean log-likelihood of the training child values under the fitted rate --
        # same convention as TabularCPT.score, and routed through the same log_prob used for
        # holdout scoring rather than a second formula.
        child_tensor = torch.tensor(y, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
        with torch.no_grad():
            score = float(log_prob(_parent_tensors(parents, names), child_tensor).mean())

        return FittedMechanism(
            mechanism=fitted_mechanism,
            noise=Uniform(0.0, 1.0),
            invertible=False,
            score=score,
        )

causalrl.regime.Regime dataclass

A labeled data-generating configuration (§5.3).

Source code in src/causalrl/regime.py

@dataclass(frozen=True)
class Regime:
    """A labeled data-generating configuration (§5.3)."""

    name: str
    selection: frozenset[str] = frozenset()
    parameters: tuple[tuple[str, Any], ...] = ()

    def __post_init__(self) -> None:
        # normalize parameter order so equality/hashing ignore insertion order
        ordered = tuple(sorted(self.parameters, key=lambda kv: kv[0]))
        object.__setattr__(self, "parameters", ordered)

    @classmethod
    def create(
        cls,
        name: str,
        *,
        selection: Iterable[str] | None = None,
        parameters: Mapping[str, Any] | None = None,
    ) -> Regime:
        params = parameters or {}
        pairs = tuple((k, params[k]) for k in sorted(params))
        return cls(name, frozenset(selection or ()), pairs)

    @property
    def params(self) -> dict[str, Any]:
        return dict(self.parameters)

    def __or__(self, other: Regime) -> Regime:
        """Merge two regimes; raise on a shared parameter set to conflicting values."""
        a, b = self.params, other.params
        for key in a.keys() & b.keys():
            if a[key] != b[key]:
                raise ValueError(f"regime conflict on parameter {key!r}: {a[key]!r} vs {b[key]!r}")
        return Regime.create(
            f"{self.name}|{other.name}",
            selection=self.selection | other.selection,
            parameters={**a, **b},
        )

    def to_domain(self, *, experiments: frozenset[frozenset[str]] = frozenset()) -> Domain:
        """Project to a transport :class:`Domain` (selection-marked variables carry over)."""
        return Domain(self.name, self.selection, experiments)

    def to_json(self) -> str:
        return json.dumps(
            {
                "name": self.name,
                "selection": sorted(self.selection),
                "parameters": [list(kv) for kv in self.parameters],
            }
        )

    @staticmethod
    def from_json(s: str) -> Regime:
        d: dict[str, Any] = json.loads(s)
        raw_params: Any = d.get("parameters", [])
        params: dict[str, Any] = {str(k): v for k, v in raw_params}
        raw_selection: Any = d.get("selection", [])
        return Regime.create(
            d["name"], selection={str(x) for x in raw_selection}, parameters=params
        )

__or__(other)

Merge two regimes; raise on a shared parameter set to conflicting values.

Source code in src/causalrl/regime.py

def __or__(self, other: Regime) -> Regime:
    """Merge two regimes; raise on a shared parameter set to conflicting values."""
    a, b = self.params, other.params
    for key in a.keys() & b.keys():
        if a[key] != b[key]:
            raise ValueError(f"regime conflict on parameter {key!r}: {a[key]!r} vs {b[key]!r}")
    return Regime.create(
        f"{self.name}|{other.name}",
        selection=self.selection | other.selection,
        parameters={**a, **b},
    )

to_domain(*, experiments=frozenset())

Project to a transport :class:Domain (selection-marked variables carry over).

Source code in src/causalrl/regime.py

def to_domain(self, *, experiments: frozenset[frozenset[str]] = frozenset()) -> Domain:
    """Project to a transport :class:`Domain` (selection-marked variables carry over)."""
    return Domain(self.name, self.selection, experiments)

causalrl.protocols.SCMCausalEnv

Conform a StructuralCausalModel (or any see/do model) to CausalEnvProtocol.

Duck-typed and torch-free: sample calls scm.see(n, seed=...); do calls scm.do(interventions).see(n, ...), packing each into a TrajectoryLog (one entity per unit, single step). The white-box noise_ledger (wrapping ExogenousPosterior) arrives in Phase 1; it returns None here.

Source code in src/causalrl/protocols.py

class SCMCausalEnv:
    """Conform a ``StructuralCausalModel`` (or any ``see``/``do`` model) to `CausalEnvProtocol`.

    Duck-typed and torch-free: ``sample`` calls ``scm.see(n, seed=...)``; ``do`` calls
    ``scm.do(interventions).see(n, ...)``, packing each into a TrajectoryLog (one entity
    per unit, single step). The white-box ``noise_ledger`` (wrapping ``ExogenousPosterior``)
    arrives in Phase 1; it returns ``None`` here.
    """

    def __init__(self, scm: Any, *, regime_label: str = "observed") -> None:
        self._scm = scm
        self._regime_label = regime_label

    def sample(
        self, n: int, *, seed: int | None = None, regime: Regime | None = None
    ) -> TrajectoryLog:
        return _samples_to_log(self._scm.see(n, seed=seed), n, self._regime_label)

    def do(
        self,
        interventions: Mapping[str, Any],
        n: int,
        *,
        seed: int | None = None,
        regime: Regime | None = None,
    ) -> TrajectoryLog:
        mutilated = self._scm.do(interventions)
        return _samples_to_log(mutilated.see(n, seed=seed), n, self._regime_label)

    def noise_ledger(self) -> NoiseLedger | None:
        return None

causalrl.scm.fitters.TabularCPT

Discrete node: a Laplace-smoothed conditional probability table, sampled by inverse CDF.

The mechanism is V = F^-1(U | parents) with U ~ Uniform(0, 1). That construction is one of many couplings reproducing the same P(V | parents), and the data cannot distinguish them — hence invertible=False, which makes counterfactuals at this node an interval rather than a point.

Every parent is discretised by its distinct observed values, and the table holds one row per parent configuration — the product of those cardinalities. So this family fits a discrete node with discrete parents. A continuous parent contributes one level per observed value and the product explodes: the commonest real-data shape, a binary treatment with two continuous confounders at n=800, already asks for 409,600 rows. fit therefore raises ValueError above :data:_MAX_CPT_ROWS rather than exhausting memory, or (where it fits) silently becoming a nearest-neighbour memoriser of the training rows. Pass a continuous family for that node instead — fit_scm(..., families={"A": ANMFit()}).

alpha is the Laplace pseudo-count and must be positive: at alpha=0 an unobserved parent configuration would give an all-zero count row and every draw there would silently be the smallest level.

Source code in src/causalrl/scm/fitters.py

class TabularCPT:
    """Discrete node: a Laplace-smoothed conditional probability table, sampled by inverse CDF.

    The mechanism is ``V = F^-1(U | parents)`` with ``U ~ Uniform(0, 1)``. That construction is
    one of many couplings reproducing the same ``P(V | parents)``, and the data cannot
    distinguish them — hence ``invertible=False``, which makes counterfactuals at this node an
    interval rather than a point.

    **Every parent is discretised by its distinct observed values**, and the table holds one row
    per parent configuration — the product of those cardinalities. So this family fits a discrete
    node with *discrete* parents. A continuous parent contributes one level per observed value and
    the product explodes: the commonest real-data shape, a binary treatment with two continuous
    confounders at n=800, already asks for 409,600 rows. ``fit`` therefore raises ``ValueError``
    above :data:`_MAX_CPT_ROWS` rather than exhausting memory, or (where it fits) silently
    becoming a nearest-neighbour memoriser of the training rows. Pass a continuous family for that
    node instead — ``fit_scm(..., families={"A": ANMFit()})``.

    ``alpha`` is the Laplace pseudo-count and must be positive: at ``alpha=0`` an unobserved parent
    configuration would give an all-zero count row and every draw there would silently be the
    smallest level.
    """

    def __init__(self, alpha: float = 1.0) -> None:
        if alpha <= 0.0:
            raise ValueError(
                f"TabularCPT(alpha={alpha}) must be positive: a non-positive pseudo-count leaves "
                "an unobserved parent configuration with an all-zero count row, whose 0/0 "
                "normalisation makes every draw there the smallest level rather than a draw from "
                "any distribution. Use a small alpha (e.g. 1e-6) for near-maximum-likelihood."
            )
        self.alpha = alpha

    def fit(self, parents: dict[str, np.ndarray], child: np.ndarray) -> FittedMechanism:
        values = np.unique(child)
        parent_names = sorted(parents)
        levels = {name: np.unique(parents[name]) for name in parent_names}
        strides: list[int] = []
        size = 1
        for name in parent_names:
            strides.append(size)
            size *= len(levels[name])
        if size > _MAX_CPT_ROWS:
            cardinalities = ", ".join(f"{name}={len(levels[name])}" for name in parent_names)
            raise ValueError(
                f"TabularCPT would need {size} rows -- one per parent configuration, the product "
                f"of the parents' distinct observed values ({cardinalities}) -- above the "
                f"_MAX_CPT_ROWS={_MAX_CPT_ROWS} limit. A parent with that many levels is "
                f"continuous, not tabular, so the table would be a nearest-neighbour memoriser of "
                f"the {len(child)} training rows rather than a conditional distribution. Fit this "
                f"node with a continuous family instead: fit_scm(..., families={{...: ANMFit()}})."
            )

        rows = (
            np.zeros(len(child), dtype=int)
            if not parent_names
            else self._config_index(parents, parent_names, levels, strides)
        )
        counts = np.full((size, len(values)), self.alpha, dtype=float)
        col_of = {v: j for j, v in enumerate(values)}
        for row, value in zip(rows, child, strict=True):
            counts[row, col_of[value]] += 1.0
        table = counts / counts.sum(axis=1, keepdims=True)

        # Mean conditional log-likelihood of the training child values under the fitted table.
        columns: list[int] = [col_of[v] for v in child]
        score = float(np.log(table[rows, columns]).mean())

        cum = torch.tensor(np.cumsum(table, axis=1), dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
        value_tensor = torch.tensor(values, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
        table_tensor = torch.tensor(table, dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
        level_tensors = {
            name: torch.tensor(levels[name], dtype=torch.float32)  # type: ignore[reportPrivateImportUsage]
            for name in parent_names
        }

        def row_index(parent_values: dict[str, Tensor], n: int) -> Tensor:
            row = torch.zeros(n, dtype=torch.long)  # type: ignore[reportPrivateImportUsage]
            for name, stride in zip(parent_names, strides, strict=True):
                column = parent_values[name].reshape(-1).float()
                # Nearest level, so an unseen/off-grid parent value maps to its closest bucket.
                distance = (column.unsqueeze(1) - level_tensors[name].unsqueeze(0)).abs()
                row = row + distance.argmin(dim=1) * stride
            return row

        def mechanism(parent_values: dict[str, Tensor], noise: Tensor) -> Tensor:
            n = noise.reshape(-1).shape[0]
            row = row_index(parent_values, n)
            picked = (noise.reshape(-1).unsqueeze(1) > cum[row]).sum(dim=1)
            return value_tensor[picked.clamp(max=len(values) - 1)]

        def log_prob(parent_values: dict[str, Tensor], value: Tensor) -> Tensor:
            """Per-sample log P(value | parents) under the fitted table -- ``evaluate_holdout``'s
            score for this non-invertible family, since there is no residual to exploit here."""
            flat = value.reshape(-1).float()
            row = row_index(parent_values, flat.shape[0])
            # Nearest fitted value, so an unseen/off-grid child value scores against its closest
            # bucket instead of a lookup error -- degrades gracefully, matching how mechanism()
            # already handles an off-grid parent value.
            distance = (flat.unsqueeze(1) - value_tensor.unsqueeze(0)).abs()
            col = distance.argmin(dim=1)
            return table_tensor[row, col].log()

        fitted_mechanism = FunctionalMechanism(parent_names, mechanism)
        fitted_mechanism.log_prob = log_prob  # type: ignore[attr-defined]
        return FittedMechanism(
            mechanism=fitted_mechanism,
            noise=Uniform(0.0, 1.0),
            invertible=False,
            score=score,
        )

    @staticmethod
    def _config_index(
        parents: dict[str, np.ndarray],
        parent_names: list[str],
        levels: dict[str, np.ndarray],
        strides: list[int],
    ) -> np.ndarray:
        rows = np.zeros(len(next(iter(parents.values()))), dtype=int)
        for name, stride in zip(parent_names, strides, strict=True):
            rows += np.searchsorted(levels[name], parents[name]) * stride
        return rows

causalrl.data.trajectory.TrajectoryLog

A columnar, long-format trajectory log (§5.4).

Source code in src/causalrl/data/trajectory.py

class TrajectoryLog:
    """A columnar, long-format trajectory log (§5.4)."""

    def __init__(
        self, columns: Mapping[str, Any], metadata: Mapping[str, Any] | None = None
    ) -> None:
        missing = [c for c in _COLUMNS if c not in columns]
        if missing:
            raise ValueError(f"TrajectoryLog missing columns: {missing}")
        n = len(columns["entity_id"])
        for c in _COLUMNS:
            if len(columns[c]) != n:
                raise ValueError(f"column {c!r} has length {len(columns[c])}, expected {n}")
        self._entity_id: NDArray[Any] = np.asarray(columns["entity_id"], dtype=np.int64)
        self._episode_id: NDArray[Any] = np.asarray(columns["episode_id"], dtype=np.int64)
        self._t: NDArray[Any] = np.asarray(columns["t"], dtype=np.int64)
        self._kind: NDArray[Any] = np.asarray([str(x) for x in columns["kind"]], dtype=object)
        self._name: NDArray[Any] = np.asarray([str(x) for x in columns["name"]], dtype=object)
        self._value: NDArray[Any] = np.empty(n, dtype=object)
        for i, v in enumerate(columns["value"]):
            self._value[i] = _coerce_value(v)
        self._regime: NDArray[Any] = np.asarray([str(x) for x in columns["regime"]], dtype=object)
        self._observed: NDArray[Any] = np.asarray(columns["observed"], dtype=bool)
        self._metadata: dict[str, Any] = dict(metadata or {})

    def __len__(self) -> int:
        return int(self._entity_id.shape[0])

    @property
    def metadata(self) -> dict[str, Any]:
        return dict(self._metadata)

    def column(self, name: str) -> NDArray[Any]:
        arrays: dict[str, NDArray[Any]] = {
            "entity_id": self._entity_id,
            "episode_id": self._episode_id,
            "t": self._t,
            "kind": self._kind,
            "name": self._name,
            "value": self._value,
            "regime": self._regime,
            "observed": self._observed,
        }
        if name not in arrays:
            raise KeyError(name)
        return arrays[name]

    @classmethod
    def from_rows(
        cls, rows: Sequence[Mapping[str, Any]], metadata: Mapping[str, Any] | None = None
    ) -> TrajectoryLog:
        cols: dict[str, list[Any]] = {c: [] for c in _COLUMNS}
        for r in rows:
            cols["entity_id"].append(r["entity_id"])
            cols["episode_id"].append(r["episode_id"])
            cols["t"].append(r["t"])
            cols["kind"].append(r["kind"])
            cols["name"].append(r["name"])
            cols["value"].append(r["value"])
            cols["regime"].append(r.get("regime", "observed"))
            cols["observed"].append(r.get("observed", True))
        return cls(cols, metadata)

    def values_by_name(self, name: str) -> NDArray[Any]:
        """All ``value`` cells whose ``name`` equals ``name``, in row order."""
        return self._value[self._name == name]

    def scan(self, batch_size: int) -> Iterator[TrajectoryLog]:
        """Iterate the log in row batches (streaming access; metadata carried on each batch)."""
        if batch_size <= 0:
            raise ValueError("batch_size must be positive")
        n = len(self)
        for start in range(0, n, batch_size):
            yield self._slice(start, min(start + batch_size, n))

    def _slice(self, start: int, stop: int) -> TrajectoryLog:
        cols = {c: self.column(c)[start:stop] for c in _COLUMNS}
        return TrajectoryLog(cols, self._metadata)

    def pivot(self) -> tuple[list[tuple[int, int, int]], dict[str, NDArray[Any]]]:
        """Long -> wide: sorted ``(entity_id, episode_id, t)`` keys and one array per variable name.

        Cells missing a value for a key are ``None`` (dense-case helper; last write wins).
        """
        keys = sorted(
            {
                (int(self._entity_id[i]), int(self._episode_id[i]), int(self._t[i]))
                for i in range(len(self))
            }
        )
        key_index = {k: j for j, k in enumerate(keys)}
        names = sorted({str(x) for x in self._name})
        table: dict[str, NDArray[Any]] = {
            nm: np.full(len(keys), None, dtype=object) for nm in names
        }
        for i in range(len(self)):
            k = (int(self._entity_id[i]), int(self._episode_id[i]), int(self._t[i]))
            table[str(self._name[i])][key_index[k]] = self._value[i]
        return keys, table

    def fingerprint(self) -> str:
        """A stable content hash over the columns and metadata (I8 data-fingerprint)."""
        payload = {
            "entity_id": self._entity_id.tolist(),
            "episode_id": self._episode_id.tolist(),
            "t": self._t.tolist(),
            "kind": [str(x) for x in self._kind],
            "name": [str(x) for x in self._name],
            "value": [self._value[i] for i in range(len(self))],
            "regime": [str(x) for x in self._regime],
            "observed": self._observed.tolist(),
            "metadata": self._metadata,
        }
        blob = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=_json_default)
        return hashlib.sha256(blob.encode("utf-8")).hexdigest()

    # --- ConfoundedTrajectoryDataset bridge (§5.4; acceptance #4) ---

    @classmethod
    def from_confounded_dataset(cls, ds: ConfoundedTrajectoryDataset) -> TrajectoryLog:
        """Lossless bridge in: one ``(entity_id=0, episode_id, t)`` row-group per transition."""
        rows: list[dict[str, Any]] = []
        episode = 0
        t = 0
        for tr in ds.transitions:
            base: dict[str, Any] = {
                "entity_id": 0,
                "episode_id": episode,
                "t": t,
                "regime": "observed",
                "observed": True,
            }
            rows.append({**base, "kind": "obs", "name": "state", "value": int(tr.state)})
            rows.append({**base, "kind": "action", "name": "action", "value": int(tr.action)})
            rows.append({**base, "kind": "reward", "name": "reward", "value": float(tr.reward)})
            rows.append({**base, "kind": "obs", "name": "next_state", "value": int(tr.next_state)})
            rows.append({**base, "kind": "done", "name": "done", "value": bool(tr.done)})
            if tr.done:
                episode += 1
                t = 0
            else:
                t += 1
        meta: dict[str, Any] = {
            "n_states": ds.n_states,
            "n_actions": ds.n_actions,
            "source": "ConfoundedTrajectoryDataset",
        }
        return cls.from_rows(rows, meta)

    def to_confounded_dataset(self) -> ConfoundedTrajectoryDataset:
        """Lossless bridge out. Requires ``n_states``/``n_actions`` in metadata."""
        if "n_states" not in self._metadata or "n_actions" not in self._metadata:
            raise ValueError("log lacks n_states/n_actions metadata; not a dataset bridge")
        groups: dict[tuple[int, int, int], dict[str, Any]] = {}
        for i in range(len(self)):
            k = (int(self._entity_id[i]), int(self._episode_id[i]), int(self._t[i]))
            groups.setdefault(k, {})[str(self._name[i])] = self._value[i]
        transitions = [
            Transition(
                state=int(groups[k]["state"]),
                action=int(groups[k]["action"]),
                reward=float(groups[k]["reward"]),
                next_state=int(groups[k]["next_state"]),
                done=bool(groups[k]["done"]),
            )
            for k in sorted(groups)
        ]
        return ConfoundedTrajectoryDataset(
            transitions,
            n_states=int(self._metadata["n_states"]),
            n_actions=int(self._metadata["n_actions"]),
        )

    # --- Arrow/Parquet IO (lazy; optional [data] extra). value union -> typed sub-columns,
    #     because Parquet has no union type. ---

    def to_arrow(self) -> Any:
        pa = _pyarrow()
        vtype: list[str] = []
        vint: list[int | None] = []
        vfloat: list[float | None] = []
        vbool: list[bool | None] = []
        vlist: list[list[float] | None] = []
        for i in range(len(self)):
            v = self._value[i]
            tag, iv, fv, bv, lv = _encode_value(v)
            vtype.append(tag)
            vint.append(iv)
            vfloat.append(fv)
            vbool.append(bv)
            vlist.append(lv)
        table = pa.table(
            {
                "entity_id": self._entity_id.tolist(),
                "episode_id": self._episode_id.tolist(),
                "t": self._t.tolist(),
                "kind": [str(x) for x in self._kind],
                "name": [str(x) for x in self._name],
                "regime": [str(x) for x in self._regime],
                "observed": self._observed.tolist(),
                "value_type": vtype,
                "value_int": pa.array(vint, type=pa.int64()),
                "value_float": pa.array(vfloat, type=pa.float64()),
                "value_bool": pa.array(vbool, type=pa.bool_()),
                "value_list": pa.array(vlist, type=pa.list_(pa.float64())),
            }
        )
        meta = {b"causalrl_metadata": json.dumps(self._metadata).encode("utf-8")}
        return table.replace_schema_metadata(meta)

    @classmethod
    def from_arrow(cls, table: Any) -> TrajectoryLog:
        d: dict[Any, Any] = {name: table.column(name).to_pylist() for name in table.column_names}
        values: list[Any] = [
            _decode_value(
                d["value_type"][i],
                d["value_int"][i],
                d["value_float"][i],
                d["value_bool"][i],
                d["value_list"][i],
            )
            for i in range(table.num_rows)
        ]
        cols: dict[str, Sequence[Any]] = {
            "entity_id": d["entity_id"],
            "episode_id": d["episode_id"],
            "t": d["t"],
            "kind": d["kind"],
            "name": d["name"],
            "value": values,
            "regime": d["regime"],
            "observed": d["observed"],
        }
        schema_meta: Any = table.schema.metadata
        raw = schema_meta.get(b"causalrl_metadata") if schema_meta else None
        metadata: dict[str, Any] = json.loads(raw) if raw else {}
        return cls(cols, metadata)

    def to_parquet(self, path: str | os.PathLike[str]) -> None:
        _, pq = _parquet()
        pq.write_table(self.to_arrow(), str(path))

    @classmethod
    def from_parquet(cls, path: str | os.PathLike[str]) -> TrajectoryLog:
        _, pq = _parquet()
        return cls.from_arrow(pq.read_table(str(path)))

    def sorted_by_key(self) -> TrajectoryLog:
        """Return a copy with rows sorted by ``(entity_id, episode_id, t)`` (stable within a key).

        Key-contiguous ordering lets the streaming estimators join a decision's cells with an
        O(1) carry-over buffer instead of holding the whole log (plan §9).
        """
        n = len(self)
        order = np.lexsort((np.arange(n), self._t, self._episode_id, self._entity_id))
        cols = {c: self.column(c)[order] for c in _COLUMNS}
        return TrajectoryLog(cols, self._metadata)

    @classmethod
    def iter_parquet_batches(
        cls, path: str | os.PathLike[str], batch_size: int
    ) -> Iterator[TrajectoryLog]:
        """Stream a Parquet log in row batches without materialising it (plan §9; ``[data]`` extra).

        Yields one :class:`TrajectoryLog` per Arrow record batch; log-level metadata is not carried
        on the per-batch logs (the streaming estimators consume named value columns, not metadata).
        """
        if batch_size <= 0:
            raise ValueError("batch_size must be positive")
        pa, pq = _parquet()
        parquet_file = pq.ParquetFile(str(path))
        for record_batch in parquet_file.iter_batches(batch_size=batch_size):
            yield cls.from_arrow(pa.Table.from_batches([record_batch]))

values_by_name(name)

All value cells whose name equals name, in row order.

Source code in src/causalrl/data/trajectory.py

def values_by_name(self, name: str) -> NDArray[Any]:
    """All ``value`` cells whose ``name`` equals ``name``, in row order."""
    return self._value[self._name == name]

scan(batch_size)

Iterate the log in row batches (streaming access; metadata carried on each batch).

Source code in src/causalrl/data/trajectory.py

def scan(self, batch_size: int) -> Iterator[TrajectoryLog]:
    """Iterate the log in row batches (streaming access; metadata carried on each batch)."""
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    n = len(self)
    for start in range(0, n, batch_size):
        yield self._slice(start, min(start + batch_size, n))

pivot()

Long -> wide: sorted (entity_id, episode_id, t) keys and one array per variable name.

Cells missing a value for a key are None (dense-case helper; last write wins).

Source code in src/causalrl/data/trajectory.py

def pivot(self) -> tuple[list[tuple[int, int, int]], dict[str, NDArray[Any]]]:
    """Long -> wide: sorted ``(entity_id, episode_id, t)`` keys and one array per variable name.

    Cells missing a value for a key are ``None`` (dense-case helper; last write wins).
    """
    keys = sorted(
        {
            (int(self._entity_id[i]), int(self._episode_id[i]), int(self._t[i]))
            for i in range(len(self))
        }
    )
    key_index = {k: j for j, k in enumerate(keys)}
    names = sorted({str(x) for x in self._name})
    table: dict[str, NDArray[Any]] = {
        nm: np.full(len(keys), None, dtype=object) for nm in names
    }
    for i in range(len(self)):
        k = (int(self._entity_id[i]), int(self._episode_id[i]), int(self._t[i]))
        table[str(self._name[i])][key_index[k]] = self._value[i]
    return keys, table

fingerprint()

A stable content hash over the columns and metadata (I8 data-fingerprint).

Source code in src/causalrl/data/trajectory.py

def fingerprint(self) -> str:
    """A stable content hash over the columns and metadata (I8 data-fingerprint)."""
    payload = {
        "entity_id": self._entity_id.tolist(),
        "episode_id": self._episode_id.tolist(),
        "t": self._t.tolist(),
        "kind": [str(x) for x in self._kind],
        "name": [str(x) for x in self._name],
        "value": [self._value[i] for i in range(len(self))],
        "regime": [str(x) for x in self._regime],
        "observed": self._observed.tolist(),
        "metadata": self._metadata,
    }
    blob = json.dumps(payload, sort_keys=True, separators=(",", ":"), default=_json_default)
    return hashlib.sha256(blob.encode("utf-8")).hexdigest()

from_confounded_dataset(ds) classmethod

Lossless bridge in: one (entity_id=0, episode_id, t) row-group per transition.

Source code in src/causalrl/data/trajectory.py

@classmethod
def from_confounded_dataset(cls, ds: ConfoundedTrajectoryDataset) -> TrajectoryLog:
    """Lossless bridge in: one ``(entity_id=0, episode_id, t)`` row-group per transition."""
    rows: list[dict[str, Any]] = []
    episode = 0
    t = 0
    for tr in ds.transitions:
        base: dict[str, Any] = {
            "entity_id": 0,
            "episode_id": episode,
            "t": t,
            "regime": "observed",
            "observed": True,
        }
        rows.append({**base, "kind": "obs", "name": "state", "value": int(tr.state)})
        rows.append({**base, "kind": "action", "name": "action", "value": int(tr.action)})
        rows.append({**base, "kind": "reward", "name": "reward", "value": float(tr.reward)})
        rows.append({**base, "kind": "obs", "name": "next_state", "value": int(tr.next_state)})
        rows.append({**base, "kind": "done", "name": "done", "value": bool(tr.done)})
        if tr.done:
            episode += 1
            t = 0
        else:
            t += 1
    meta: dict[str, Any] = {
        "n_states": ds.n_states,
        "n_actions": ds.n_actions,
        "source": "ConfoundedTrajectoryDataset",
    }
    return cls.from_rows(rows, meta)

to_confounded_dataset()

Lossless bridge out. Requires n_states/n_actions in metadata.

Source code in src/causalrl/data/trajectory.py

def to_confounded_dataset(self) -> ConfoundedTrajectoryDataset:
    """Lossless bridge out. Requires ``n_states``/``n_actions`` in metadata."""
    if "n_states" not in self._metadata or "n_actions" not in self._metadata:
        raise ValueError("log lacks n_states/n_actions metadata; not a dataset bridge")
    groups: dict[tuple[int, int, int], dict[str, Any]] = {}
    for i in range(len(self)):
        k = (int(self._entity_id[i]), int(self._episode_id[i]), int(self._t[i]))
        groups.setdefault(k, {})[str(self._name[i])] = self._value[i]
    transitions = [
        Transition(
            state=int(groups[k]["state"]),
            action=int(groups[k]["action"]),
            reward=float(groups[k]["reward"]),
            next_state=int(groups[k]["next_state"]),
            done=bool(groups[k]["done"]),
        )
        for k in sorted(groups)
    ]
    return ConfoundedTrajectoryDataset(
        transitions,
        n_states=int(self._metadata["n_states"]),
        n_actions=int(self._metadata["n_actions"]),
    )

sorted_by_key()

Return a copy with rows sorted by (entity_id, episode_id, t) (stable within a key).

Key-contiguous ordering lets the streaming estimators join a decision's cells with an O(1) carry-over buffer instead of holding the whole log (plan §9).

Source code in src/causalrl/data/trajectory.py

def sorted_by_key(self) -> TrajectoryLog:
    """Return a copy with rows sorted by ``(entity_id, episode_id, t)`` (stable within a key).

    Key-contiguous ordering lets the streaming estimators join a decision's cells with an
    O(1) carry-over buffer instead of holding the whole log (plan §9).
    """
    n = len(self)
    order = np.lexsort((np.arange(n), self._t, self._episode_id, self._entity_id))
    cols = {c: self.column(c)[order] for c in _COLUMNS}
    return TrajectoryLog(cols, self._metadata)

iter_parquet_batches(path, batch_size) classmethod

Stream a Parquet log in row batches without materialising it (plan §9; [data] extra).

Yields one :class:TrajectoryLog per Arrow record batch; log-level metadata is not carried on the per-batch logs (the streaming estimators consume named value columns, not metadata).

Source code in src/causalrl/data/trajectory.py

@classmethod
def iter_parquet_batches(
    cls, path: str | os.PathLike[str], batch_size: int
) -> Iterator[TrajectoryLog]:
    """Stream a Parquet log in row batches without materialising it (plan §9; ``[data]`` extra).

    Yields one :class:`TrajectoryLog` per Arrow record batch; log-level metadata is not carried
    on the per-batch logs (the streaming estimators consume named value columns, not metadata).
    """
    if batch_size <= 0:
        raise ValueError("batch_size must be positive")
    pa, pq = _parquet()
    parquet_file = pq.ParquetFile(str(path))
    for record_batch in parquet_file.iter_batches(batch_size=batch_size):
        yield cls.from_arrow(pa.Table.from_batches([record_batch]))

causalrl.data.dataset.Transition dataclass

A single logged transition. reward is the terminal return (0 except on done).

Source code in src/causalrl/data/dataset.py

@dataclass(frozen=True)
class Transition:
    """A single logged transition. `reward` is the terminal return (0 except on done)."""

    state: int
    action: int
    reward: float
    next_state: int
    done: bool

causalrl.scm.continuous.abduction.abduct_invertible(mechanism, parent_values, observed, *, name='U')

Exactly recover any invertible mechanism's exogenous noise from an observed outcome.

Works for the location-scale and conditional-flow families alike (anything implementing :class:InvertibleMechanism); the recovered noise is a point mass, so counterfactuals built on it may claim kind=IDENTIFIED.

Source code in src/causalrl/scm/continuous/abduction.py

def abduct_invertible(
    mechanism: InvertibleMechanism,
    parent_values: Mapping[str, Tensor],
    observed: Tensor,
    *,
    name: str = "U",
) -> PointNoisePosterior:
    """Exactly recover any invertible mechanism's exogenous noise from an observed outcome.

    Works for the location-scale and conditional-flow families alike (anything implementing
    :class:`InvertibleMechanism`); the recovered noise is a point mass, so counterfactuals built on
    it may claim ``kind=IDENTIFIED``.
    """
    return PointNoisePosterior({name: mechanism.invert(dict(parent_values), observed)})

causalrl.scm.continuous.abduction.abduct_location_scale(mechanism, parent_values, observed, *, name='U')

Exactly recover a location-scale mechanism's noise (see :func:abduct_invertible).

Source code in src/causalrl/scm/continuous/abduction.py

def abduct_location_scale(
    mechanism: LocationScaleMechanism,
    parent_values: Mapping[str, Tensor],
    observed: Tensor,
    *,
    name: str = "U",
) -> PointNoisePosterior:
    """Exactly recover a location-scale mechanism's noise (see :func:`abduct_invertible`)."""
    return abduct_invertible(mechanism, parent_values, observed, name=name)

causalrl.scm.continuous.nuts.abduct_nuts(forward, parent_values, observed, *, noise_scale=0.1, num_warmup=500, num_samples=1000, seed=0, name='U')

Draw the posterior over a mechanism's exogenous noise by NUTS (NumPyro).

Model: a standard-normal prior u_i ~ N(0, 1) per unit and a Gaussian likelihood observed_i ~ N(forward(parents, u)_i, noise_scale). forward must be written in jax.numpy (elementwise over units). Returns a pure-NumPy :class:NUTSNoisePosterior. Requires the [numpyro] extra.

Source code in src/causalrl/scm/continuous/nuts.py

def abduct_nuts(
    forward: JaxForward,
    parent_values: Mapping[str, Any],
    observed: Any,
    *,
    noise_scale: float = 0.1,
    num_warmup: int = 500,
    num_samples: int = 1000,
    seed: int = 0,
    name: str = "U",
) -> NUTSNoisePosterior:
    """Draw the posterior over a mechanism's exogenous noise by NUTS (NumPyro).

    Model: a standard-normal prior ``u_i ~ N(0, 1)`` per unit and a Gaussian likelihood
    ``observed_i ~ N(forward(parents, u)_i, noise_scale)``. ``forward`` must be written in
    ``jax.numpy`` (elementwise over units). Returns a pure-NumPy :class:`NUTSNoisePosterior`.
    Requires the ``[numpyro]`` extra.
    """
    _require_numpyro()
    import jax
    import jax.numpy as jnp
    import numpyro
    import numpyro.distributions as dist
    from numpyro.infer import MCMC, NUTS

    y = jnp.asarray(np.asarray(observed, dtype=np.float64))
    n_units = int(y.shape[0])
    pv = {k: jnp.asarray(np.asarray(v, dtype=np.float64)) for k, v in parent_values.items()}

    def model() -> None:
        with numpyro.plate("units", n_units):
            u = numpyro.sample("u", dist.Normal(0.0, 1.0))
            mean = forward(pv, u)
            numpyro.sample("obs", dist.Normal(mean, noise_scale), obs=y)

    mcmc = MCMC(NUTS(model), num_warmup=num_warmup, num_samples=num_samples, progress_bar=False)
    mcmc.run(jax.random.PRNGKey(seed))
    draws = np.asarray(mcmc.get_samples()["u"], dtype=np.float64)  # (num_samples, n_units)
    return NUTSNoisePosterior(draws, name=name)

causalrl.scm.unrolled.build_unrolled_scm(transition, horizon, *, state0_dist, latents, actions=None, process_noise_dist=None, state_prefix='state')

Build a time-unrolled :class:StructuralCausalModel over horizon steps.

Parameters

transition: Callable (state_t, action_t, latents, noise_t) -> state_{t+1}. Receives the current state tensor (n, state_dim), the step's action constant, a {name: tensor} mapping of the shared latent values, and the step's per-step exogenous tensor. Deterministic transitions simply ignore noise_t. horizon: Number of transition steps T. The model has states {prefix}_0 .. {prefix}_T. state0_dist: Exogenous distribution for the initial state {prefix}_0. (For the deterministic abduction pattern the actual start is pinned via abduct(known=...), but the node still needs a declared exogenous distribution.) latents: Shared latent node(s), e.g. {"F": Bernoulli(0.5)}. Each is a parent of every transition step and is therefore interveneable via do / pinnable via abduct. actions: Optional per-step action constants; actions[t] is passed to the transition that produces {prefix}_{t+1}. Must have length horizon when given. None passes None as every action. process_noise_dist: Optional per-step exogenous distribution attached to each produced state node. Defaults to a point mass at zero (fully deterministic dynamics). state_prefix: Node-name prefix for the state chain (default "state").

Returns

StructuralCausalModel With nodes {prefix}_0 .. {prefix}_T and one node per latent.

Source code in src/causalrl/scm/unrolled.py

def build_unrolled_scm(
    transition: Transition,
    horizon: int,
    *,
    state0_dist: Distribution,
    latents: Mapping[str, Distribution],
    actions: Sequence[object] | None = None,
    process_noise_dist: Distribution | None = None,
    state_prefix: str = "state",
) -> StructuralCausalModel:
    """Build a time-unrolled :class:`StructuralCausalModel` over ``horizon`` steps.

    Parameters
    ----------
    transition:
        Callable ``(state_t, action_t, latents, noise_t) -> state_{t+1}``. Receives the current
        state tensor ``(n, state_dim)``, the step's action constant, a ``{name: tensor}`` mapping
        of the shared latent values, and the step's per-step exogenous tensor. Deterministic
        transitions simply ignore ``noise_t``.
    horizon:
        Number of transition steps ``T``. The model has states ``{prefix}_0 .. {prefix}_T``.
    state0_dist:
        Exogenous distribution for the initial state ``{prefix}_0``. (For the deterministic
        abduction pattern the actual start is pinned via ``abduct(known=...)``, but the node
        still needs a declared exogenous distribution.)
    latents:
        Shared latent node(s), e.g. ``{"F": Bernoulli(0.5)}``. Each is a parent of *every*
        transition step and is therefore interveneable via ``do`` / pinnable via ``abduct``.
    actions:
        Optional per-step action constants; ``actions[t]`` is passed to the transition that
        produces ``{prefix}_{t+1}``. Must have length ``horizon`` when given. ``None`` passes
        ``None`` as every action.
    process_noise_dist:
        Optional per-step exogenous distribution attached to each produced state node. Defaults
        to a point mass at zero (fully deterministic dynamics).
    state_prefix:
        Node-name prefix for the state chain (default ``"state"``).

    Returns
    -------
    StructuralCausalModel
        With nodes ``{prefix}_0 .. {prefix}_T`` and one node per latent.
    """
    if horizon < 1:
        raise ValueError(f"horizon must be >= 1, got {horizon}")
    if actions is not None and len(actions) != horizon:
        raise ValueError(f"actions must have length horizon={horizon}, got {len(actions)}")
    latent_names = list(latents)
    if not latent_names:
        raise ValueError("at least one shared latent node is required")
    reserved = {f"{state_prefix}_{t}" for t in range(horizon + 1)}
    clash = reserved & set(latent_names)
    if clash:
        raise ValueError(f"latent name(s) collide with state nodes: {sorted(clash)}")

    noise_dist: Distribution = process_noise_dist or _Degenerate()

    def state_name(t: int) -> str:
        return f"{state_prefix}_{t}"

    nodes: list[str] = latent_names + [state_name(t) for t in range(horizon + 1)]
    directed_edges: list[tuple[str, str]] = []
    mechanisms: dict[str, object] = {}
    exogenous: dict[str, Distribution] = {}

    # Shared latents: root nodes whose value IS their exogenous draw (identity mechanism).
    for name in latent_names:
        exogenous[name] = latents[name]
        mechanisms[name] = FunctionalMechanism([], lambda pa, u: u)

    # Initial state: root node, value is its exogenous draw.
    exogenous[state_name(0)] = state0_dist
    mechanisms[state_name(0)] = FunctionalMechanism([], lambda pa, u: u)

    # Transition steps: state_{t+1} = transition(state_t, action_t, latents, noise_t).
    for t in range(horizon):
        cur, nxt = state_name(t), state_name(t + 1)
        parents = [cur, *latent_names]
        directed_edges.append((cur, nxt))
        for name in latent_names:
            directed_edges.append((name, nxt))
        action_t = None if actions is None else actions[t]

        def make_step_fn(
            _cur: str = cur,
            _action: object = action_t,
        ) -> Callable[[dict[str, Tensor], Tensor], Tensor]:
            def step_fn(pa: dict[str, Tensor], u: Tensor) -> Tensor:
                state = pa[_cur]
                latent_vals = {name: pa[name] for name in latent_names}
                return transition(state, _action, latent_vals, u)

            return step_fn

        mechanisms[nxt] = FunctionalMechanism(parents, make_step_fn())
        exogenous[nxt] = noise_dist

    # Every node (latents, state_0..state_T) appears in ``directed_edges`` for horizon >= 1,
    # so edges alone define the node set. Passing ``nodes=`` *and* edges can double-list nodes
    # in some CausalGraph builds, which then fails the SCM's exact-coverage check — so rely on
    # the edges. ``nodes`` is retained only for the assertion below.
    graph = CausalGraph(directed_edges=directed_edges)
    assert set(graph.nodes) == set(nodes), (set(graph.nodes), set(nodes))
    return StructuralCausalModel(graph, mechanisms, exogenous)  # type: ignore[arg-type]

causalrl.conformal.core.certify_conformal_interval(prediction, cal_y, cal_pred, *, alpha=0.1, weights=None, test_weight=None)

A kind=EMPIRICAL certificate for a split-conformal prediction interval.

value is left None (the claim is coverage of an individual outcome, not a point estimate); the interval is in ci and alpha is the miscoverage. The (weighted-)exchangeability requirement is recorded as a non-checkable :class:Assumption.

The query is "see" and is not configurable: residuals of a fitted prediction carry no intervention, and weights only move the observational law to a shifted one, so neither an interventional nor a counterfactual label would be earned. For the off-policy claim, use :func:conformal_action_value, which computes the propensity ratio and records the causal assumptions that license it.

Source code in src/causalrl/conformal/core.py

def certify_conformal_interval(
    prediction: float,
    cal_y: Sequence[float],
    cal_pred: Sequence[float],
    *,
    alpha: float = 0.1,
    weights: Sequence[float] | None = None,
    test_weight: float | None = None,
) -> Certificate:
    """A ``kind=EMPIRICAL`` certificate for a split-conformal prediction interval.

    ``value`` is left ``None`` (the claim is coverage of an individual outcome, not a point
    estimate); the interval is in ``ci`` and ``alpha`` is the miscoverage. The
    (weighted-)exchangeability requirement is recorded as a non-checkable :class:`Assumption`.

    The query is ``"see"`` and is not configurable: residuals of a fitted prediction carry no
    intervention, and ``weights`` only move the *observational* law to a shifted one, so neither
    an interventional nor a counterfactual label would be earned. For the off-policy claim, use
    :func:`conformal_action_value`, which computes the propensity ratio and records the causal
    assumptions that license it.
    """
    interval = split_conformal_interval(
        prediction, cal_y, cal_pred, alpha, weights=weights, test_weight=test_weight
    )
    shifted = weights is not None
    return Certificate(
        claim=(
            f"P(Y in [{interval.lower:.4g}, {interval.upper:.4g}]) >= {1.0 - alpha:.2f} (conformal)"
        ),
        estimand=EstimandSpec(query="see", target="interval"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=alpha,
        assumptions=(
            Assumption(
                name="weighted-exchangeability" if shifted else "exchangeability",
                params={"n_calibration": int(np.size(cal_y))},
                checkable=False,
            ),
        ),
        method="weighted-split-conformal" if shifted else "split-conformal",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
        ci=interval,
    )

causalrl.scm.continuous.abduction.certify_counterfactual(claim, ppc, *, exact, alpha=None)

Wrap a counterfactual query in a certificate.

kind is IDENTIFIED for the exact-inversion path and EMPIRICAL for amortized VI (I2); the PPC is recorded as a checkable assumption.

Source code in src/causalrl/scm/continuous/abduction.py

def certify_counterfactual(
    claim: str, ppc: dict[str, float], *, exact: bool, alpha: float | None = None
) -> Certificate:
    """Wrap a counterfactual query in a certificate.

    ``kind`` is IDENTIFIED for the exact-inversion path and EMPIRICAL for amortized VI (I2); the PPC
    is recorded as a checkable assumption.
    """
    method = "exact-inversion" if exact else "amortized-vi"
    return Certificate(
        claim=claim,
        estimand=EstimandSpec(query="counterfactual", target="mean"),
        kind=Kind.IDENTIFIED if exact else Kind.EMPIRICAL,
        value=None,
        alpha=alpha,
        assumptions=(
            Assumption(name="abduction", params={"method": method}, checkable=True, diagnostic=ppc),
        ),
        method=method,
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
    )

causalrl.scm.continuous.nuts.certify_nuts_counterfactual(claim, ppc, *, alpha=None)

Wrap a NUTS-abducted counterfactual in an EMPIRICAL certificate (sampling evidence; I2).

The posterior-predictive check is recorded as a checkable assumption.

Source code in src/causalrl/scm/continuous/nuts.py

def certify_nuts_counterfactual(
    claim: str, ppc: dict[str, float], *, alpha: float | None = None
) -> Certificate:
    """Wrap a NUTS-abducted counterfactual in an ``EMPIRICAL`` certificate (sampling evidence; I2).

    The posterior-predictive check is recorded as a checkable assumption.
    """
    return Certificate(
        claim=claim,
        estimand=EstimandSpec(query="counterfactual", target="mean"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=alpha,
        assumptions=(
            Assumption(name="abduction", params={"method": "nuts"}, checkable=True, diagnostic=ppc),
        ),
        method="nuts",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
    )

causalrl.conformal.core.conformal_quantile(scores, alpha=0.1, *, weights=None, test_weight=None)

Conformal quantile of nonconformity scores at miscoverage alpha.

Returns the smallest score whose (weight-)normalised rank reaches 1 - alpha once the test point's +inf mass is included, so prediction ± result (or the CQR adjustment) is marginally valid at level 1 - alpha under (weighted) exchangeability. Returns inf (a valid, infinite interval) when the calibration set is too small for the requested level. weights are calibration likelihood ratios dP_test/dP_cal; test_weight defaults to their mean.

Source code in src/causalrl/conformal/core.py

def conformal_quantile(
    scores: Sequence[float],
    alpha: float = 0.1,
    *,
    weights: Sequence[float] | None = None,
    test_weight: float | None = None,
) -> float:
    """Conformal quantile of nonconformity ``scores`` at miscoverage ``alpha``.

    Returns the smallest score whose (weight-)normalised rank reaches ``1 - alpha`` once the test
    point's ``+inf`` mass is included, so ``prediction ± result`` (or the CQR adjustment) is
    marginally valid at level ``1 - alpha`` under (weighted) exchangeability. Returns ``inf`` (a
    valid, infinite interval) when the calibration set is too small for the requested level.
    ``weights`` are calibration likelihood ratios ``dP_test/dP_cal``; ``test_weight`` defaults to
    their mean.
    """
    s = np.asarray(scores, dtype=np.float64)
    n = s.size
    if n == 0:
        raise ValueError("need at least one calibration score")
    if not 0.0 < alpha < 1.0:
        raise ValueError("alpha must be in (0, 1)")
    w = np.ones(n) if weights is None else np.asarray(weights, dtype=np.float64)
    if w.shape != s.shape:
        raise ValueError("weights must match scores in length")
    wt = float(w.mean()) if test_weight is None else float(test_weight)
    order = np.argsort(s)
    s_sorted = s[order]
    w_sorted = w[order]
    total = float(w_sorted.sum()) + wt
    cum = np.cumsum(w_sorted) / total
    idx = int(np.searchsorted(cum, 1.0 - alpha, side="left"))
    return float("inf") if idx >= n else float(s_sorted[idx])

causalrl.conformal.core.cqr_interval(pred_lo, pred_hi, cal_y, cal_lo, cal_hi, alpha=0.1, *, weights=None, test_weight=None)

Conformalized quantile regression interval (Romano, Patterson & Candes 2019).

cal_lo/cal_hi are the calibration lower/upper quantile predictions; the conformity score is max(cal_lo - y, y - cal_hi) and the returned interval widens the test quantile band [pred_lo, pred_hi] by the conformal quantile, giving adaptive finite-sample coverage.

Source code in src/causalrl/conformal/core.py

def cqr_interval(
    pred_lo: float,
    pred_hi: float,
    cal_y: Sequence[float],
    cal_lo: Sequence[float],
    cal_hi: Sequence[float],
    alpha: float = 0.1,
    *,
    weights: Sequence[float] | None = None,
    test_weight: float | None = None,
) -> Interval:
    """Conformalized quantile regression interval (Romano, Patterson & Candes 2019).

    ``cal_lo``/``cal_hi`` are the calibration lower/upper quantile predictions; the conformity score
    is ``max(cal_lo - y, y - cal_hi)`` and the returned interval widens the test quantile band
    ``[pred_lo, pred_hi]`` by the conformal quantile, giving adaptive finite-sample coverage.
    """
    y = np.asarray(cal_y, dtype=np.float64)
    lo = np.asarray(cal_lo, dtype=np.float64)
    hi = np.asarray(cal_hi, dtype=np.float64)
    scores = np.maximum(lo - y, y - hi)
    q = conformal_quantile(scores.tolist(), alpha, weights=weights, test_weight=test_weight)
    return Interval(float(pred_lo) - q, float(pred_hi) + q)

causalrl.scm.fit.fit_scm_mec(data, *, cpdag, max_members=32, **kwargs)

Fit one SCM per DAG in cpdag's Markov equivalence class.

The returned list is the SCM belief: observational data cannot choose among these, so every member fits it equally well. Sub-project 3 (active interventional discovery) shrinks this set by intervening. Raises ValueError above max_members naming the true class size — equivalence classes are exponential and silently truncating one would misreport the belief.

A second ValueError fires before enumerating when the CPDAG has too many undirected edges: finding the class means trying all 2 ** k orientations of them, and that work is exponential in k no matter how small the class turns out to be, so max_members cannot bound it (see :data:_MAX_MEC_ENUMERATION). The pre-check can only name 2 ** k, an upper bound on the class size rather than the size itself, and so over-refuses: the honest reading is "this search is too big to run", not "this class is too big to fit".

Source code in src/causalrl/scm/fit.py

def fit_scm_mec(
    data: Mapping[str, np.ndarray],
    *,
    cpdag: CPDAG,
    max_members: int = 32,
    **kwargs: Any,
) -> list[StructuralCausalModel]:
    """Fit one SCM per DAG in ``cpdag``'s Markov equivalence class.

    The returned list is the SCM *belief*: observational data cannot choose among these, so every
    member fits it equally well. Sub-project 3 (active interventional discovery) shrinks this set
    by intervening. Raises ``ValueError`` above ``max_members`` naming the true class size —
    equivalence classes are exponential and silently truncating one would misreport the belief.

    A second ``ValueError`` fires *before* enumerating when the CPDAG has too many undirected
    edges: finding the class means trying all ``2 ** k`` orientations of them, and that work is
    exponential in ``k`` no matter how small the class turns out to be, so ``max_members`` cannot
    bound it (see :data:`_MAX_MEC_ENUMERATION`). The pre-check can only name ``2 ** k``, an upper
    bound on the class size rather than the size itself, and so over-refuses: the honest reading
    is "this search is too big to run", not "this class is too big to fit".
    """
    undirected_count = len(cpdag.undirected_edges)
    candidates = 1 << undirected_count
    if candidates > _MAX_MEC_ENUMERATION:
        raise ValueError(
            f"cpdag has {undirected_count} undirected edges, so enumerating its equivalence class "
            f"means orienting them 2**{undirected_count} = {candidates} ways -- above the "
            f"_MAX_MEC_ENUMERATION={_MAX_MEC_ENUMERATION} search budget. That count is an upper "
            f"bound on the class size, not the size itself, which is only known after the search; "
            f"max_members caps the class and so cannot cap this. Narrow the CPDAG with tiers via "
            f"orient() (or supply interventional data) before fitting."
        )
    members = _enumerate_mec(cpdag)
    if not members:
        # No orientation of the undirected edges is acyclic, so the CPDAG's *directed* edges
        # already contain a cycle and no DAG is consistent with it. Returning [] here would hand
        # back an empty belief that reads exactly like "not fitted yet" -- a caller's next query
        # then fails with a misleading message about calling this function, when the real fault is
        # an inconsistent graph. This function already refuses to truncate a class it cannot
        # enumerate; refusing an empty one is the same principle.
        raise NotIdentifiableError(
            f"no DAG is consistent with this CPDAG, so its equivalence class is empty: the "
            f"directed edges {sorted(cpdag.directed_edges)} admit no acyclic orientation of the "
            f"{undirected_count} undirected one(s). A cyclic directed set usually means the "
            f"orientations came from separate sources that were not reconciled -- "
            f"discover_interventional orients the edges incident to each intervention target "
            f"independently, so two or more targets can disagree. Reconcile them (or drop a "
            f"target) before fitting.",
            witness=sorted(cpdag.directed_edges),
        )
    if len(members) > max_members:
        raise ValueError(
            f"equivalence class has {len(members)} members, above max_members={max_members}; "
            "raise the cap deliberately or narrow the CPDAG with tiers via orient()"
        )
    return [fit_scm(data, graph=graph, **kwargs) for graph in members]

causalrl.data.dataset.generate_logs(env, n_episodes, seed)

Roll out an env's confounded behavior_policy to build an offline dataset.

The env must expose n_states, n_actions, reset(seed=...), step(action), and behavior_policy(observation). Rewards are recorded as the per-step reward; for the finite-horizon envs here the terminal step carries the return.

Source code in src/causalrl/data/dataset.py

def generate_logs(env: RolloutEnv, n_episodes: int, seed: int) -> ConfoundedTrajectoryDataset:
    """Roll out an env's confounded behavior_policy to build an offline dataset.

    The env must expose ``n_states``, ``n_actions``, ``reset(seed=...)``, ``step(action)``,
    and ``behavior_policy(observation)``. Rewards are recorded as the per-step reward; for
    the finite-horizon envs here the terminal step carries the return.
    """
    transitions: list[Transition] = []
    obs, _ = env.reset(seed=seed)
    for ep in range(n_episodes):
        if ep > 0:
            obs, _ = env.reset()
        done = False
        while not done:
            state = int(obs["state"])
            action = int(env.behavior_policy(obs))
            next_obs, reward, terminated, truncated, _info = env.step(action)
            done = bool(terminated or truncated)
            transitions.append(
                Transition(state, action, float(reward), int(next_obs["state"]), done)
            )
            obs = next_obs
    return ConfoundedTrajectoryDataset(transitions, n_states=env.n_states, n_actions=env.n_actions)

causalrl.conformal.core.split_conformal_interval(prediction, cal_y, cal_pred, alpha=0.1, *, weights=None, test_weight=None)

Symmetric split-conformal interval prediction ± q from absolute calibration residuals.

Source code in src/causalrl/conformal/core.py

def split_conformal_interval(
    prediction: float,
    cal_y: Sequence[float],
    cal_pred: Sequence[float],
    alpha: float = 0.1,
    *,
    weights: Sequence[float] | None = None,
    test_weight: float | None = None,
) -> Interval:
    """Symmetric split-conformal interval ``prediction ± q`` from absolute calibration residuals."""
    resid = np.abs(np.asarray(cal_y, dtype=np.float64) - np.asarray(cal_pred, dtype=np.float64))
    q = conformal_quantile(resid.tolist(), alpha, weights=weights, test_weight=test_weight)
    return Interval(float(prediction) - q, float(prediction) + q)

Agents, Environments & Interventions — Complete Reference

causalrl.agents.base.Agent

Bases: ABC

Minimal agent interface: choose an action, then learn from the reward.

Implement this directly for an online learner, one whose policy moves in response to each reward it sees. If learning happens in a batch step instead, subclass :class:BatchAgent.

Source code in src/causalrl/agents/base.py

class Agent(ABC):
    """Minimal agent interface: choose an action, then learn from the reward.

    Implement this directly for an **online** learner, one whose policy moves in response to each
    reward it sees. If learning happens in a batch step instead, subclass :class:`BatchAgent`.
    """

    @abstractmethod
    def act(self, observation: dict[str, Any]) -> int: ...

    @abstractmethod
    def update(self, observation: dict[str, Any], action: int, reward: float) -> None: ...

    def observe_transition(  # noqa: B027
        self, state: int, action: int, next_state: int, done: bool
    ) -> None:
        """Optional model-learning hook: observe a `(s, a, s', done)` transition.

        Default no-op. Model-based agents (e.g. DOVI's value iteration) override this to
        build an empirical transition model; reward-only agents ignore it.
        """
        pass

    def observe_step(
        self,
        observation: dict[str, Any],
        action: int,
        reward: float,
        next_observation: dict[str, Any],
        done: bool,
    ) -> None:
        """Observe one transition in the environment's own observation space.

        This is the hook a *driver* should call. :meth:`observe_transition` takes state indices,
        which forces whoever drives the loop to discretise before the agent sees anything -- and an
        agent whose states are feature vectors then has no way to receive the transition at all.
        Here the agent is handed the raw observations and decides for itself how to represent them.

        The default is the tabular reading, so every existing agent and driver keeps its exact
        behaviour: pull ``observation["state"]`` from each endpoint and forward to
        :meth:`observe_transition`. A feature-space agent overrides this to encode instead.
        """
        self.observe_transition(
            int(observation["state"]), action, int(next_observation["state"]), done
        )

observe_transition(state, action, next_state, done)

Optional model-learning hook: observe a (s, a, s', done) transition.

Default no-op. Model-based agents (e.g. DOVI's value iteration) override this to build an empirical transition model; reward-only agents ignore it.

Source code in src/causalrl/agents/base.py

def observe_transition(  # noqa: B027
    self, state: int, action: int, next_state: int, done: bool
) -> None:
    """Optional model-learning hook: observe a `(s, a, s', done)` transition.

    Default no-op. Model-based agents (e.g. DOVI's value iteration) override this to
    build an empirical transition model; reward-only agents ignore it.
    """
    pass

observe_step(observation, action, reward, next_observation, done)

Observe one transition in the environment's own observation space.

This is the hook a driver should call. :meth:observe_transition takes state indices, which forces whoever drives the loop to discretise before the agent sees anything -- and an agent whose states are feature vectors then has no way to receive the transition at all. Here the agent is handed the raw observations and decides for itself how to represent them.

The default is the tabular reading, so every existing agent and driver keeps its exact behaviour: pull observation["state"] from each endpoint and forward to :meth:observe_transition. A feature-space agent overrides this to encode instead.

Source code in src/causalrl/agents/base.py

def observe_step(
    self,
    observation: dict[str, Any],
    action: int,
    reward: float,
    next_observation: dict[str, Any],
    done: bool,
) -> None:
    """Observe one transition in the environment's own observation space.

    This is the hook a *driver* should call. :meth:`observe_transition` takes state indices,
    which forces whoever drives the loop to discretise before the agent sees anything -- and an
    agent whose states are feature vectors then has no way to receive the transition at all.
    Here the agent is handed the raw observations and decides for itself how to represent them.

    The default is the tabular reading, so every existing agent and driver keeps its exact
    behaviour: pull ``observation["state"]`` from each endpoint and forward to
    :meth:`observe_transition`. A feature-space agent overrides this to encode instead.
    """
    self.observe_transition(
        int(observation["state"]), action, int(next_observation["state"]), done
    )

causalrl.magames.AgentType dataclass

A template shared by a class of agents: an action set and a payoff over the joint profile.

payoff(own_action, others_actions, params) receives the other agents' actions in population order (self excluded) and the instance's params; sharing one AgentType across instances is the parameter sharing.

Source code in src/causalrl/magames/population.py

@dataclass(frozen=True)
class AgentType:
    """A template shared by a class of agents: an action set and a payoff over the joint profile.

    ``payoff(own_action, others_actions, params)`` receives the other agents' actions in population
    order (self excluded) and the instance's ``params``; sharing one ``AgentType`` across instances
    *is* the parameter sharing.
    """

    name: str
    actions: tuple[int, ...]
    payoff: PayoffTemplate

causalrl.agents.mbrl.BackdoorAdjustedAgent

Bases: BatchAgent

Active deconfounded optimizer: pick the action with the highest back-door-adjusted value E[Y | do(A=a)] = Σ_z P(z) · E[Y | A=a, Z=z], with the adjustment set read from the graph via :func:~causalrl.backdoor_adjustment_set.

Unlike the certify-gated agent (whose ceiling is the behavior policy), this recovers the interventional optimum from confounded logs, given an observed admissible adjustment set. It is fitted on columnar data {treatment, outcome, *adjustment} (equal-length arrays).

:meth:fit keeps the per-stratum conditional means, so :meth:act is a policy over the adjustment set — argmax_a Ehat[Y | do(a), Z=z] — not the one action that wins on average.

Source code in src/causalrl/agents/mbrl.py

class BackdoorAdjustedAgent(BatchAgent):
    """Active deconfounded optimizer: pick the action with the highest back-door-adjusted value
    ``E[Y | do(A=a)] = Σ_z P(z) · E[Y | A=a, Z=z]``, with the adjustment set read from the graph via
    :func:`~causalrl.backdoor_adjustment_set`.

    Unlike the certify-gated agent (whose ceiling is the behavior policy), this *recovers* the
    interventional optimum from confounded logs, given an observed admissible adjustment set. It is
    fitted on columnar data ``{treatment, outcome, *adjustment}`` (equal-length arrays).

    :meth:`fit` keeps the per-stratum conditional means, so :meth:`act` is a *policy* over the
    adjustment set — ``argmax_a Ehat[Y | do(a), Z=z]`` — not the one action that wins on average.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        graph: CausalGraph,
        treatment: str = "A",
        outcome: str = "Y",
    ) -> None:
        self.n_actions = n_actions
        self.treatment = treatment
        self.outcome = outcome
        self.adjustment: tuple[str, ...] = tuple(
            sorted(backdoor_adjustment_set(graph, treatment, outcome))
        )
        self._best_action = 0
        self.values: list[float] = [0.0] * n_actions
        self._stratum_values: dict[tuple[Any, ...], tuple[float, ...]] = {}

    def fit(self, data: Mapping[str, np.ndarray]) -> None:
        """Estimate each action's back-door-adjusted value, cache the per-stratum action values
        :meth:`act` conditions on, and record the marginal argmax."""
        self.values = [self._adjusted_value(a, data) for a in range(self.n_actions)]
        self._best_action = int(np.argmax(np.asarray(self.values)))
        self._stratum_values = _stratum_action_values(
            data, self.treatment, self.outcome, self.adjustment, self.n_actions
        )

    def _adjusted_value(self, action: int, data: Mapping[str, np.ndarray]) -> float:
        return _backdoor_value(action, data, self.treatment, self.outcome, self.adjustment)

    def act(self, observation: dict[str, Any]) -> int:
        """Contextual policy ``argmax_a Ehat[Y | do(A=a), Z=z]`` at the adjustment values in
        ``observation`` (supplied by name, e.g. ``{"Z": 1}``).

        Conditioning on the back-door set is what makes this a policy rather than a single action:
        the arm that wins *after* averaging over Z need not win *inside* a stratum. With none of
        the adjustment variables supplied — or an empty adjustment set, or a stratum never logged —
        the marginal decision ``argmax_a E[Y | do(a)]`` is returned, which is the right call when
        there is no context to condition on. A partial adjustment set raises.
        """
        return _stratum_action(
            self._stratum_values, observation, self.adjustment, self._best_action
        )

fit(data)

Estimate each action's back-door-adjusted value, cache the per-stratum action values :meth:act conditions on, and record the marginal argmax.

Source code in src/causalrl/agents/mbrl.py

def fit(self, data: Mapping[str, np.ndarray]) -> None:
    """Estimate each action's back-door-adjusted value, cache the per-stratum action values
    :meth:`act` conditions on, and record the marginal argmax."""
    self.values = [self._adjusted_value(a, data) for a in range(self.n_actions)]
    self._best_action = int(np.argmax(np.asarray(self.values)))
    self._stratum_values = _stratum_action_values(
        data, self.treatment, self.outcome, self.adjustment, self.n_actions
    )

act(observation)

Contextual policy argmax_a Ehat[Y | do(A=a), Z=z] at the adjustment values in observation (supplied by name, e.g. {"Z": 1}).

Conditioning on the back-door set is what makes this a policy rather than a single action: the arm that wins after averaging over Z need not win inside a stratum. With none of the adjustment variables supplied — or an empty adjustment set, or a stratum never logged — the marginal decision argmax_a E[Y | do(a)] is returned, which is the right call when there is no context to condition on. A partial adjustment set raises.

Source code in src/causalrl/agents/mbrl.py

def act(self, observation: dict[str, Any]) -> int:
    """Contextual policy ``argmax_a Ehat[Y | do(A=a), Z=z]`` at the adjustment values in
    ``observation`` (supplied by name, e.g. ``{"Z": 1}``).

    Conditioning on the back-door set is what makes this a policy rather than a single action:
    the arm that wins *after* averaging over Z need not win *inside* a stratum. With none of
    the adjustment variables supplied — or an empty adjustment set, or a stratum never logged —
    the marginal decision ``argmax_a E[Y | do(a)]`` is returned, which is the right call when
    there is no context to condition on. A partial adjustment set raises.
    """
    return _stratum_action(
        self._stratum_values, observation, self.adjustment, self._best_action
    )

causalrl.agents.scbandit.BruteForceInterventionTS

Bases: _ArmSubsetThompsonSampling

Thompson sampling over every arm (no pruning) — the unguided baseline.

Source code in src/causalrl/agents/scbandit.py

class BruteForceInterventionTS(_ArmSubsetThompsonSampling):
    """Thompson sampling over every arm (no pruning) — the unguided baseline."""

    def __init__(self, arms: list[dict[str, int]], seed: int | None = None) -> None:
        super().__init__(list(range(len(arms))), seed)

causalrl.magames.CCEPolytope dataclass

The CCE polytope of a (possibly intervened) finite game, in distribution-over-profiles form.

profiles are the joint actions (in agents order) consistent with the intervention; deviation_gains[k, j] is the payoff gain of constraint_labels[k] = (agent, action) switching unilaterally to action when the realized profile is profiles[j]. A distribution mu over profiles is an epsilon-CCE iff deviation_gains @ mu is componentwise at most epsilon.

Source code in src/causalrl/magames/cce.py

@dataclass(frozen=True)
class CCEPolytope:
    """The CCE polytope of a (possibly intervened) finite game, in distribution-over-profiles form.

    ``profiles`` are the joint actions (in ``agents`` order) consistent with the intervention;
    ``deviation_gains[k, j]`` is the payoff gain of ``constraint_labels[k] = (agent, action)``
    switching unilaterally to ``action`` when the realized profile is ``profiles[j]``. A
    distribution ``mu`` over ``profiles`` is an ``epsilon``-CCE iff ``deviation_gains @ mu`` is
    componentwise at most ``epsilon``.
    """

    agents: tuple[str, ...]
    profiles: tuple[tuple[int, ...], ...]
    deviation_gains: FloatArray
    constraint_labels: tuple[tuple[str, int], ...]

causalrl.envs.base.CausalEnv

Bases: Env

Base class for environments backed by a StructuralCausalModel.

Subclasses define how SCM samples map to (observation, reward) and which node is the action. The SCM is available as self.scm for see/do/counterfactual queries.

Source code in src/causalrl/envs/base.py

class CausalEnv(gym.Env):  # type: ignore[type-arg]
    """Base class for environments backed by a StructuralCausalModel.

    Subclasses define how SCM samples map to (observation, reward) and which node is the
    action. The SCM is available as `self.scm` for see/do/counterfactual queries.
    """

    def __init__(self, scm: StructuralCausalModel) -> None:
        super().__init__()
        self.scm = scm

causalrl.agents.mbrl.CertifiedPolicyAgent

Bases: BatchAgent

Ship the best deterministic policy whose improvement over behavior certifies robust to hidden confounding; abstain to the empirical behavior policy otherwise.

Setting alpha adds the finite-sample layer: a candidate must also clear :func:causalrl.certify_policy's conformal lower-confidence-bound gate, i.e. its calibrated worst-case return (:func:causalrl.conformal_action_value, weights = the propensity ratio pi/pi_behavior) must be at least the behavior policy's. This is the safe-policy-improvement reading of the agent: a candidate that improves the mean but degrades the downside, or that has too little effective support to calibrate a bound at all, is refused. None (the default) runs the confounding layer alone.

Source code in src/causalrl/agents/mbrl.py

class CertifiedPolicyAgent(BatchAgent):
    """Ship the best deterministic policy whose improvement over behavior certifies robust to hidden
    confounding; abstain to the empirical behavior policy otherwise.

    Setting ``alpha`` adds the finite-sample layer: a candidate must also clear
    :func:`causalrl.certify_policy`'s conformal lower-confidence-bound gate, i.e. its calibrated
    worst-case return (:func:`causalrl.conformal_action_value`, weights = the propensity ratio
    ``pi/pi_behavior``) must be at least the behavior policy's. This is the safe-policy-improvement
    reading of the agent: a candidate that improves the *mean* but degrades the downside, or that
    has too little effective support to calibrate a bound at all, is refused. ``None`` (the
    default) runs the confounding layer alone.
    """

    def __init__(
        self,
        n_states: int,
        n_actions: int,
        *,
        gamma_max: float = 5.0,
        alpha: float | None = None,
    ) -> None:
        self.n_states = n_states
        self.n_actions = n_actions
        self.gamma_max = gamma_max
        self.alpha = alpha
        self.policy: list[int] = [0] * n_states

    def _behavior_policy(self, dataset: ConfoundedTrajectoryDataset) -> list[int]:
        """Empirical behavior policy: the most-logged action per state (abstention target)."""
        return [
            max(range(self.n_actions), key=lambda a: dataset.behavior_propensity(s, a))
            for s in range(self.n_states)
        ]

    def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
        transitions = dataset.transitions
        best_policy = self._behavior_policy(dataset)  # abstention default
        best_contrast = 0.0
        for candidate in itertools.product(range(self.n_actions), repeat=self.n_states):
            # Skip a policy that assigns a never-logged action in some state: an unseen action's
            # value is not identified from the logs, and certify_policy has no support to bound it.
            if any(
                dataset.behavior_propensity(s, candidate[s]) == 0.0 for s in range(self.n_states)
            ):
                continue
            target_actions = [candidate[tr.state] for tr in transitions]
            cert = certify_policy(
                dataset, target_actions, gamma_max=self.gamma_max, alpha=self.alpha
            )
            if cert.certified and cert.naive_contrast > best_contrast:
                best_contrast = cert.naive_contrast
                best_policy = list(candidate)
        self.policy = best_policy

    def act(self, observation: dict[str, Any]) -> int:
        return int(self.policy[int(observation["state"])])

causalrl.envs.suite.confounded_context.ConfoundedContextualBandit

Bases: CausalEnv

A 2-context, 2-action bandit with a hidden confounder U on a backdoor A <- U -> Y.

Source code in src/causalrl/envs/suite/confounded_context.py

class ConfoundedContextualBandit(CausalEnv):
    """A 2-context, 2-action bandit with a hidden confounder ``U`` on a backdoor ``A <- U -> Y``."""

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]
    n_states: int = 2
    n_actions: int = 2

    def __init__(self, gamma: float = 0.9, shift: bool = False, *, seed: int | None = None) -> None:
        super().__init__(_build_scm(shift))
        self.gamma = float(gamma)
        self.shift = bool(shift)
        self.action_space = gym.spaces.Discrete(2)
        self.observation_space = gym.spaces.Dict({"state": gym.spaces.Discrete(2)})
        self._rng = np.random.default_rng(seed)
        self._c = 0
        self._u = 0

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        self._c = int(self._rng.integers(0, 2))
        self._u = int(self._rng.integers(0, 2))
        return {"state": self._c}, {}

    def behavior_policy(self, observation: dict[str, Any]) -> int:
        """Confounded logging policy: play ``a = U`` with prob ``gamma``, else uniform."""
        if self._rng.random() < self.gamma:
            return self._u
        return int(self._rng.integers(0, 2))

    def _reward_prob(self, c: int, a: int, u: int) -> float:
        c_opt = (1 - c) if self.shift else c
        return min(1.0, max(0.0, 0.05 + 0.20 * float(a == c_opt) + 0.60 * float(u)))

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
        p = self._reward_prob(self._c, int(action), self._u)
        reward = float(self._rng.random() < p)
        return {"state": self._c}, reward, True, False, {"u": self._u}

    def true_policy_value(self, policy: Sequence[int]) -> float:
        """Exact interventional value of a deterministic ``policy`` (marginalizing U)."""

        def _q(c: int) -> float:
            a = int(policy[c])
            return 0.5 * (self._reward_prob(c, a, 0) + self._reward_prob(c, a, 1))

        return float(np.mean([_q(c) for c in range(self.n_states)]))

behavior_policy(observation)

Confounded logging policy: play a = U with prob gamma, else uniform.

Source code in src/causalrl/envs/suite/confounded_context.py

def behavior_policy(self, observation: dict[str, Any]) -> int:
    """Confounded logging policy: play ``a = U`` with prob ``gamma``, else uniform."""
    if self._rng.random() < self.gamma:
        return self._u
    return int(self._rng.integers(0, 2))

true_policy_value(policy)

Exact interventional value of a deterministic policy (marginalizing U).

Source code in src/causalrl/envs/suite/confounded_context.py

def true_policy_value(self, policy: Sequence[int]) -> float:
    """Exact interventional value of a deterministic ``policy`` (marginalizing U)."""

    def _q(c: int) -> float:
        a = int(policy[c])
        return 0.5 * (self._reward_prob(c, a, 0) + self._reward_prob(c, a, 1))

    return float(np.mean([_q(c) for c in range(self.n_states)]))

causalrl.envs.suite.gridworld.ConfoundedGridworld

Bases: ConfoundedMDP

A size x size gridworld with a hidden confounder that biases the logging policy.

State = row * size + col (0-indexed), plus a terminal state = sizesize. Start at top-left (state 0), goal at bottom-right. Actions: 0=up, 1=down, 2=left, 3=right. Reaching the goal yields terminal return 1.0; hitting the per-episode 'hazard' cell (chosen by the hidden confounder U) yields 0.0 and ends the episode. Horizon = 2size. The behavior policy, knowing U, avoids the hazard — so the logs are confounded.

Source code in src/causalrl/envs/suite/gridworld.py

class ConfoundedGridworld(ConfoundedMDP):
    """A size x size gridworld with a hidden confounder that biases the logging policy.

    State = row * size + col (0-indexed), plus a terminal state = size*size. Start at
    top-left (state 0), goal at bottom-right. Actions: 0=up, 1=down, 2=left, 3=right.
    Reaching the goal yields terminal return 1.0; hitting the per-episode 'hazard' cell
    (chosen by the hidden confounder U) yields 0.0 and ends the episode. Horizon = 2*size.
    The behavior policy, knowing U, avoids the hazard — so the logs are confounded.
    """

    n_actions = 4

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]  # gymnasium Env.metadata is an instance var in the base

    def __init__(self, size: int = 3, seed: int | None = None) -> None:
        super().__init__()
        self.size = size
        self.n_states = size * size + 1
        self.horizon = 2 * size
        self.action_space = gym.spaces.Discrete(4)
        self.observation_space = gym.spaces.Dict(
            {
                "state": gym.spaces.Discrete(self.n_states),
                "t": gym.spaces.Discrete(self.horizon + 1),
            }
        )
        self._rng = np.random.default_rng(seed)
        self._terminal = size * size
        self._row = 0
        self._col = 0
        self._t = 0
        self._u = 0

    def _state_index(self) -> int:
        return self._row * self.size + self._col

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, int], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        self._row = 0
        self._col = 0
        self._t = 0
        self._u = int(self._rng.integers(1, self.size * self.size - 1))
        return {"state": 0, "t": 0}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, int], float, bool, bool, dict[str, Any]]:
        if action == 0:
            self._row = max(0, self._row - 1)
        elif action == 1:
            self._row = min(self.size - 1, self._row + 1)
        elif action == 2:
            self._col = max(0, self._col - 1)
        elif action == 3:
            self._col = min(self.size - 1, self._col + 1)
        self._t += 1
        idx = self._state_index()
        if idx == self._u:
            return {"state": self._terminal, "t": self._t}, 0.0, True, False, {"hazard": True}
        if self._row == self.size - 1 and self._col == self.size - 1:
            return {"state": self._terminal, "t": self._t}, 1.0, True, False, {}
        if self._t >= self.horizon:
            return {"state": self._terminal, "t": self._t}, 0.0, True, False, {}
        return {"state": idx, "t": self._t}, 0.0, False, False, {}

    def behavior_policy(self, observation: dict[str, int]) -> int:
        """Knows the hazard U: biases toward down/right but avoids stepping onto the hazard."""
        candidates = [1, 3]  # down, right (toward goal)
        safe: list[int] = []
        for a in candidates:
            r, c = self._row, self._col
            if a == 1:
                r = min(self.size - 1, r + 1)
            else:
                c = min(self.size - 1, c + 1)
            if r * self.size + c != self._u:
                safe.append(a)
        choices = safe or candidates
        return int(self._rng.choice(choices))

behavior_policy(observation)

Knows the hazard U: biases toward down/right but avoids stepping onto the hazard.

Source code in src/causalrl/envs/suite/gridworld.py

def behavior_policy(self, observation: dict[str, int]) -> int:
    """Knows the hazard U: biases toward down/right but avoids stepping onto the hazard."""
    candidates = [1, 3]  # down, right (toward goal)
    safe: list[int] = []
    for a in candidates:
        r, c = self._row, self._col
        if a == 1:
            r = min(self.size - 1, r + 1)
        else:
            c = min(self.size - 1, c + 1)
        if r * self.size + c != self._u:
            safe.append(a)
    choices = safe or candidates
    return int(self._rng.choice(choices))

causalrl.envs.base.ConfoundedMDP

Bases: CausalEnv

Finite-horizon tabular MDP whose logging policy depends on an unobserved confounder.

Subclasses set n_states, n_actions, horizon, draw a fresh UC per episode in reset, and implement step dynamics. behavior_policy is the UC-dependent logging policy used to generate offline data; it may read the per-episode UC via self.

scm is set to None here because confounded MDPs are defined by their dynamics and behavior policy rather than an explicit SCM; the attribute is kept for interface compatibility with CausalEnv.

Source code in src/causalrl/envs/base.py

class ConfoundedMDP(CausalEnv):
    """Finite-horizon tabular MDP whose logging policy depends on an unobserved confounder.

    Subclasses set `n_states`, `n_actions`, `horizon`, draw a fresh UC per episode in
    `reset`, and implement step dynamics. `behavior_policy` is the UC-dependent logging
    policy used to generate offline data; it may read the per-episode UC via `self`.

    `scm` is set to None here because confounded MDPs are defined by their dynamics and
    behavior policy rather than an explicit SCM; the attribute is kept for interface
    compatibility with CausalEnv.
    """

    n_states: int
    n_actions: int
    horizon: int

    # Satisfies CausalEnv.scm without requiring a real SCM; subclasses may override.
    scm: StructuralCausalModel | None = None  # type: ignore[assignment]

    def __init__(self) -> None:
        # Do not call super().__init__(scm) — no SCM is required for confounded MDPs.
        gym.Env.__init__(self)  # type: ignore[reportUnknownMemberType]

    @abstractmethod
    def behavior_policy(self, observation: dict[str, int]) -> int:
        """The (confounded) logging policy. May depend on the hidden per-episode UC."""

behavior_policy(observation) abstractmethod

The (confounded) logging policy. May depend on the hidden per-episode UC.

Source code in src/causalrl/envs/base.py

@abstractmethod
def behavior_policy(self, observation: dict[str, int]) -> int:
    """The (confounded) logging policy. May depend on the hidden per-episode UC."""

causalrl.envs.suite.continuous_confounded.ContinuousConfoundedBandit

One-decision bandit with a continuous observed confounder Z and a nonlinear arm-1 reward.

Z is the back-door confounder; the reward is continuous (Gaussian observation noise), so the outcome model must be a function approximator rather than a per-stratum mean.

Source code in src/causalrl/envs/suite/continuous_confounded.py

class ContinuousConfoundedBandit:
    """One-decision bandit with a continuous observed confounder ``Z`` and a nonlinear arm-1 reward.

    ``Z`` is the back-door confounder; the reward is continuous (Gaussian observation noise), so the
    outcome model must be a function approximator rather than a per-stratum mean.
    """

    n_states: int = 1
    n_actions: int = 2

    def __init__(self, *, gamma: float = 1.0, seed: int | None = None) -> None:
        if not 0.0 <= gamma <= 1.0:
            raise ValueError("gamma must be in [0, 1]")
        self.gamma = float(gamma)
        self.graph = CausalGraph(directed_edges=[("Z", "A"), ("Z", "Y"), ("A", "Y")])
        self._rng = np.random.default_rng(seed)

    def _reward_mean(self, a: np.ndarray, z: np.ndarray) -> np.ndarray:
        """Vectorized ``E[Y | A=a, Z=z]``: flat 0.5 for arm 0, the bump ``q(1, z)`` for arm 1."""
        arm1 = _BASE + _AMP * np.exp(-((z - _CENTER) ** 2) / (2 * _WIDTH**2))
        return np.where(a == 1, arm1, 0.5)

    def sample(self, n: int, *, seed: int | None = None) -> dict[str, np.ndarray]:
        """Draw ``n`` confounded observations ``{"Z", "A", "Y"}`` (``Z``/``Y`` continuous)."""
        rng = np.random.default_rng(seed) if seed is not None else self._rng
        z = rng.random(n)
        a = (rng.random(n) < np.clip(0.5 + _SLOPE * self.gamma * (2 * z - 1), 0.03, 0.97)).astype(
            int
        )
        y = self._reward_mean(a, z) + rng.normal(0.0, _NOISE, n)
        return {"Z": z, "A": a, "Y": y}

    def true_action_value(self, action: int, *, n_mc: int = 200_000) -> float:
        """Exact ``E[Y | do(A=action)] = E_z[E[Y | action, z]]`` over ``Z ~ Uniform(0, 1)``."""
        z = np.linspace(0.0, 1.0, n_mc)
        return float(self._reward_mean(np.full(n_mc, action), z).mean())

    def optimal_action(self) -> int:
        return max(range(self.n_actions), key=self.true_action_value)

    def optimal_value(self) -> float:
        return max(self.true_action_value(a) for a in range(self.n_actions))

sample(n, *, seed=None)

Draw n confounded observations {"Z", "A", "Y"} (Z/Y continuous).

Source code in src/causalrl/envs/suite/continuous_confounded.py

def sample(self, n: int, *, seed: int | None = None) -> dict[str, np.ndarray]:
    """Draw ``n`` confounded observations ``{"Z", "A", "Y"}`` (``Z``/``Y`` continuous)."""
    rng = np.random.default_rng(seed) if seed is not None else self._rng
    z = rng.random(n)
    a = (rng.random(n) < np.clip(0.5 + _SLOPE * self.gamma * (2 * z - 1), 0.03, 0.97)).astype(
        int
    )
    y = self._reward_mean(a, z) + rng.normal(0.0, _NOISE, n)
    return {"Z": z, "A": a, "Y": y}

true_action_value(action, *, n_mc=200000)

Exact E[Y | do(A=action)] = E_z[E[Y | action, z]] over Z ~ Uniform(0, 1).

Source code in src/causalrl/envs/suite/continuous_confounded.py

def true_action_value(self, action: int, *, n_mc: int = 200_000) -> float:
    """Exact ``E[Y | do(A=action)] = E_z[E[Y | action, z]]`` over ``Z ~ Uniform(0, 1)``."""
    z = np.linspace(0.0, 1.0, n_mc)
    return float(self._reward_mean(np.full(n_mc, action), z).mean())

causalrl.envs.suite.counterfactual_bandit.CounterfactualBanditEnv

Bases: CausalEnv

One-step three-arm bandit. Observe intent intuition = I in {0,1,2}, play arm X in {0,1,2}, receive reward Y (0.8 if X == U else 0.15). The counterfactual-optimal policy plays X = I.

Source code in src/causalrl/envs/suite/counterfactual_bandit.py

class CounterfactualBanditEnv(CausalEnv):
    """One-step three-arm bandit. Observe intent ``intuition = I in {0,1,2}``, play arm
    ``X in {0,1,2}``, receive reward ``Y`` (``0.8`` if ``X == U`` else ``0.15``). The
    counterfactual-optimal policy plays ``X = I``."""

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]  # gymnasium Env.metadata is an instance var in the base

    def __init__(self, seed: int | None = None) -> None:
        super().__init__(build_counterfactual_scm())
        self.action_space = gym.spaces.Discrete(_N_ARMS)
        self.observation_space = gym.spaces.Dict({"intuition": gym.spaces.Discrete(_N_ARMS)})
        self._rng: torch.Generator = torch.Generator()  # type: ignore[reportPrivateImportUsage]
        if seed is not None:
            self._rng.manual_seed(seed)
        self._u = 0

    def _draw_confounder(self) -> None:
        self._u = int(torch.randint(0, _N_ARMS, (1,), generator=self._rng).item())  # type: ignore[reportPrivateImportUsage]

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, int], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng.manual_seed(seed)
        self._draw_confounder()
        return {"intuition": self._u}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, int], float, bool, bool, dict[str, Any]]:
        p = _P_WIN if action == self._u else _P_LOSE
        reward = float(torch.rand(1, generator=self._rng).item() < p)  # type: ignore[reportPrivateImportUsage]
        return {"intuition": self._u}, reward, True, False, {"u": self._u}

causalrl.envs.suite.dtr.DTREnv

Bases: ConfoundedMDP

Confounded (single-stage) dynamic treatment regime.

A context subtype Z ~ Bernoulli(0.5) is OBSERVED by the agent (the state). An unobserved confounder U ~ Bernoulli(0.5) is hidden from the agent but drives the logging policy's treatment choice. The agent picks a treatment a in {0, 1} and receives a Bernoulli reward whose success probability depends on (Z, a, U).

state = 0 / 1 : observed subtype Z state = 2 : terminal Horizon 1 (one treatment decision per episode).

The reward table is chosen so that: - the do-optimal policy is a = Z (matched treatment), with average value 0.75; - the confounder U inflates the apparent reward of treatment 1, so a logging policy that follows U over-selects treatment 1. In the resulting confounded logs, treatment 1 looks best for BOTH subtypes, so a naive offline learner picks a = 1 everywhere and is WRONG for subtype Z = 0 (true value 0.55 vs the optimal 0.70) — average value 0.675.

This is the canonical causal-RL lesson: naively trusting confounded logs yields a confidently biased, suboptimal policy.

Source code in src/causalrl/envs/suite/dtr.py

class DTREnv(ConfoundedMDP):
    """Confounded (single-stage) dynamic treatment regime.

    A context subtype ``Z ~ Bernoulli(0.5)`` is OBSERVED by the agent (the state). An
    unobserved confounder ``U ~ Bernoulli(0.5)`` is hidden from the agent but drives the
    logging policy's treatment choice. The agent picks a treatment ``a in {0, 1}`` and
    receives a Bernoulli reward whose success probability depends on ``(Z, a, U)``.

    state = 0 / 1 : observed subtype Z
    state = 2     : terminal
    Horizon 1 (one treatment decision per episode).

    The reward table is chosen so that:
    - the do-optimal policy is a = Z (matched treatment), with average value 0.75;
    - the confounder U inflates the apparent reward of treatment 1, so a logging policy that
      follows U over-selects treatment 1. In the resulting confounded logs, treatment 1
      looks best for BOTH subtypes, so a naive offline learner picks a = 1 everywhere and is
      WRONG for subtype Z = 0 (true value 0.55 vs the optimal 0.70) — average value 0.675.

    This is the canonical causal-RL lesson: naively trusting confounded logs yields a
    confidently biased, suboptimal policy.
    """

    n_states = 3
    n_actions = 2
    horizon = 1

    # P(R = 1 | Z, a, U). do-value(Z, a) = mean over U ~ Bernoulli(0.5).
    # do-values: (0,0)=0.70 (0,1)=0.55 (1,0)=0.20 (1,1)=0.80 -> optimal a = Z, value 0.75.
    _REWARD_PROB: ClassVar[dict[tuple[int, int, int], float]] = {
        (0, 0, 0): 0.70,
        (0, 0, 1): 0.70,
        (0, 1, 0): 0.20,
        (0, 1, 1): 0.90,
        (1, 0, 0): 0.20,
        (1, 0, 1): 0.20,
        (1, 1, 0): 0.70,
        (1, 1, 1): 0.90,
    }
    _BEHAVIOR_FOLLOW_U = 0.85

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]  # gymnasium Env.metadata is an instance var in the base

    def __init__(self, seed: int | None = None) -> None:
        super().__init__()
        self.action_space = gym.spaces.Discrete(2)
        self.observation_space = gym.spaces.Dict(
            {"state": gym.spaces.Discrete(3), "t": gym.spaces.Discrete(2)}
        )
        self._rng = np.random.default_rng(seed)
        self._z = 0
        self._u = 0

    def do_value(self, z: int, a: int) -> float:
        """The interventional value E[R | do(a), Z=z], averaged over U ~ Bernoulli(0.5)."""
        return 0.5 * self._REWARD_PROB[(z, a, 0)] + 0.5 * self._REWARD_PROB[(z, a, 1)]

    @property
    def optimal_value(self) -> float:
        """Average value of the do-optimal policy a*(Z) = argmax_a do_value(Z, a)."""
        return 0.5 * sum(max(self.do_value(z, 0), self.do_value(z, 1)) for z in (0, 1))

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, int], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        self._z = int(self._rng.integers(0, 2))
        self._u = int(self._rng.integers(0, 2))
        return {"state": self._z, "t": 0}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, int], float, bool, bool, dict[str, Any]]:
        p = self._REWARD_PROB[(self._z, action, self._u)]
        reward = float(self._rng.random() < p)
        return {"state": _TERMINAL, "t": 1}, reward, True, False, {"u": self._u}

    def behavior_policy(self, observation: dict[str, int]) -> int:
        """Confounded logging: the clinician observes the hidden comorbidity U and prescribes
        it with probability 0.85 (else the opposite), independent of the subtype Z. This ties
        the logged action to U, which also drives the reward — the source of confounding."""
        if self._rng.random() < self._BEHAVIOR_FOLLOW_U:
            return self._u
        return 1 - self._u

optimal_value property

Average value of the do-optimal policy a*(Z) = argmax_a do_value(Z, a).

do_value(z, a)

The interventional value E[R | do(a), Z=z], averaged over U ~ Bernoulli(0.5).

Source code in src/causalrl/envs/suite/dtr.py

def do_value(self, z: int, a: int) -> float:
    """The interventional value E[R | do(a), Z=z], averaged over U ~ Bernoulli(0.5)."""
    return 0.5 * self._REWARD_PROB[(z, a, 0)] + 0.5 * self._REWARD_PROB[(z, a, 1)]

behavior_policy(observation)

Confounded logging: the clinician observes the hidden comorbidity U and prescribes it with probability 0.85 (else the opposite), independent of the subtype Z. This ties the logged action to U, which also drives the reward — the source of confounding.

Source code in src/causalrl/envs/suite/dtr.py

def behavior_policy(self, observation: dict[str, int]) -> int:
    """Confounded logging: the clinician observes the hidden comorbidity U and prescribes
    it with probability 0.85 (else the opposite), independent of the subtype Z. This ties
    the logged action to U, which also drives the reward — the source of confounding."""
    if self._rng.random() < self._BEHAVIOR_FOLLOW_U:
        return self._u
    return 1 - self._u

causalrl.agents.deep_deconfounded.DeepDeconfoundedQ

Bases: Agent

DQN-style agent whose bootstrap targets are clamped into the Manski causal bounds.

Function approximation can extrapolate Q outside the causally-valid envelope; clamping each (state, action) target into [lower, upper] keeps the learned values consistent with what the confounded offline data can support, while online experience tightens them. A small one-step bootstrap is used (sufficient for the toy demo); a full offline-RL backbone (d3rlpy) is the designated reuse path at real scale.

Source code in src/causalrl/agents/deep_deconfounded.py

class DeepDeconfoundedQ(Agent):
    """DQN-style agent whose bootstrap targets are clamped into the Manski causal bounds.

    Function approximation can extrapolate Q outside the causally-valid envelope; clamping
    each (state, action) target into [lower, upper] keeps the learned values consistent with
    what the confounded offline data can support, while online experience tightens them.
    A small one-step bootstrap is used (sufficient for the toy demo); a full offline-RL
    backbone (d3rlpy) is the designated reuse path at real scale.
    """

    def __init__(self, n_states: int, n_actions: int, seed: int | None = None) -> None:
        self.n_states = n_states
        self.n_actions = n_actions
        torch.manual_seed(0 if seed is None else seed)  # type: ignore[reportPrivateImportUsage]
        self._rng = np.random.default_rng(seed)
        self._q = _QNet(n_states, n_actions)
        self._opt = torch.optim.Adam(self._q.parameters(), lr=1e-2)  # type: ignore[reportPrivateImportUsage]
        self._eps = 0.2
        self._lower = np.zeros((n_states, n_actions))
        self._upper = np.ones((n_states, n_actions))

    def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
        for (s, a), (lo, hi) in bounds_table(dataset).items():
            self._lower[s, a] = lo
            self._upper[s, a] = hi

    def bound(self, state: int, action: int) -> tuple[float, float]:
        return float(self._lower[state, action]), float(self._upper[state, action])

    def clamp_target(self, state: int, action: int, target: float) -> float:
        lo, hi = self.bound(state, action)
        return float(min(max(target, lo), hi))

    def _onehot(self, state: int) -> torch.Tensor:
        x = torch.zeros(self.n_states)  # type: ignore[reportPrivateImportUsage]
        x[state] = 1.0
        return x

    def act(self, observation: dict[str, Any]) -> int:
        s = int(observation["state"])
        if self._rng.random() < self._eps:
            return int(self._rng.integers(0, self.n_actions))
        with torch.no_grad():  # type: ignore[reportPrivateImportUsage]
            q = self._q(self._onehot(s))
        return int(torch.argmax(q).item())  # type: ignore[reportPrivateImportUsage]

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        s = int(observation["state"])
        target = self.clamp_target(s, action, float(reward))
        q = self._q(self._onehot(s))
        pred = q[action]
        loss = (pred - torch.tensor(target)) ** 2  # type: ignore[reportPrivateImportUsage]
        self._opt.zero_grad()
        loss.backward()
        self._opt.step()  # type: ignore[reportUnknownMemberType]

causalrl.agents.mbrl.DiscoveryBackdoorAgent

Bases: BatchAgent

Learns the structure: discovers the causal skeleton from data, orients it with the known temporal tier order (covariates precede treatment precede outcome — standard in DTR/medicine), takes the treatment's earlier-tier neighbours as the back-door set, then adjusts. The M1 upgrade of the handed-the-graph :class:BackdoorAdjustedAgent. Fitted with :meth:discover_and_fit.

The discovered adjustment set is also the policy's context: :meth:act returns argmax_a Ehat[Y | do(a), Z=z] over the variables discovery selected.

Source code in src/causalrl/agents/mbrl.py

class DiscoveryBackdoorAgent(BatchAgent):
    """Learns the structure: discovers the causal skeleton from data, orients it with the known
    temporal tier order (covariates precede treatment precede outcome — standard in DTR/medicine),
    takes the treatment's earlier-tier neighbours as the back-door set, then adjusts. The M1 upgrade
    of the handed-the-graph :class:`BackdoorAdjustedAgent`. Fitted with :meth:`discover_and_fit`.

    The *discovered* adjustment set is also the policy's context: :meth:`act` returns
    ``argmax_a Ehat[Y | do(a), Z=z]`` over the variables discovery selected."""

    def __init__(
        self,
        n_actions: int,
        *,
        variables: Sequence[str],
        treatment: str = "A",
        outcome: str = "Y",
    ) -> None:
        self.n_actions = n_actions
        self.variables = tuple(variables)
        self.treatment = treatment
        self.outcome = outcome
        self.adjustment: tuple[str, ...] = ()
        self._best_action = 0
        self.values: list[float] = [0.0] * n_actions
        self._stratum_values: dict[tuple[Any, ...], tuple[float, ...]] = {}

    def discover_and_fit(
        self,
        data: Mapping[str, np.ndarray],
        *,
        tiers: Sequence[Sequence[str]],
        threshold: float = 0.01,
    ) -> None:
        """Discover the skeleton, orient it by the temporal ``tiers`` (ordered groups, earliest
        first), take the treatment's earlier-tier neighbours as the back-door set, and adjust.

        Temporal tiering is standard domain knowledge in DTR / medicine and makes adjustment-set
        recovery reliable (pure interventional edge-orientation was not, on this graph).
        """
        cpdag = discover(data, self.variables, threshold=threshold)
        neighbours: set[str] = set()
        for a, b in cpdag.directed_edges:
            if self.treatment in (a, b):
                neighbours.update({a, b} - {self.treatment})
        for edge in cpdag.undirected_edges:
            if self.treatment in edge:
                neighbours.update(v for v in edge if v != self.treatment)
        tier_of = {v: i for i, group in enumerate(tiers) for v in group}
        adj = [v for v in neighbours if tier_of[v] < tier_of[self.treatment]]
        self.adjustment = tuple(sorted(adj))
        self.values = [
            _backdoor_value(a, data, self.treatment, self.outcome, self.adjustment)
            for a in range(self.n_actions)
        ]
        self._best_action = int(np.argmax(np.asarray(self.values)))
        self._stratum_values = _stratum_action_values(
            data, self.treatment, self.outcome, self.adjustment, self.n_actions
        )

    def act(self, observation: dict[str, Any]) -> int:
        """Contextual policy ``argmax_a Ehat[Y | do(A=a), Z=z]`` over the *discovered* adjustment
        set, at the values in ``observation`` (supplied by name, e.g. ``{"Z": 1}``).

        With none of the discovered variables supplied — or an empty discovered set, or a stratum
        never logged — the marginal decision ``argmax_a E[Y | do(a)]`` is returned; a partial
        context raises.
        """
        return _stratum_action(
            self._stratum_values, observation, self.adjustment, self._best_action
        )

discover_and_fit(data, *, tiers, threshold=0.01)

Discover the skeleton, orient it by the temporal tiers (ordered groups, earliest first), take the treatment's earlier-tier neighbours as the back-door set, and adjust.

Temporal tiering is standard domain knowledge in DTR / medicine and makes adjustment-set recovery reliable (pure interventional edge-orientation was not, on this graph).

Source code in src/causalrl/agents/mbrl.py

def discover_and_fit(
    self,
    data: Mapping[str, np.ndarray],
    *,
    tiers: Sequence[Sequence[str]],
    threshold: float = 0.01,
) -> None:
    """Discover the skeleton, orient it by the temporal ``tiers`` (ordered groups, earliest
    first), take the treatment's earlier-tier neighbours as the back-door set, and adjust.

    Temporal tiering is standard domain knowledge in DTR / medicine and makes adjustment-set
    recovery reliable (pure interventional edge-orientation was not, on this graph).
    """
    cpdag = discover(data, self.variables, threshold=threshold)
    neighbours: set[str] = set()
    for a, b in cpdag.directed_edges:
        if self.treatment in (a, b):
            neighbours.update({a, b} - {self.treatment})
    for edge in cpdag.undirected_edges:
        if self.treatment in edge:
            neighbours.update(v for v in edge if v != self.treatment)
    tier_of = {v: i for i, group in enumerate(tiers) for v in group}
    adj = [v for v in neighbours if tier_of[v] < tier_of[self.treatment]]
    self.adjustment = tuple(sorted(adj))
    self.values = [
        _backdoor_value(a, data, self.treatment, self.outcome, self.adjustment)
        for a in range(self.n_actions)
    ]
    self._best_action = int(np.argmax(np.asarray(self.values)))
    self._stratum_values = _stratum_action_values(
        data, self.treatment, self.outcome, self.adjustment, self.n_actions
    )

act(observation)

Contextual policy argmax_a Ehat[Y | do(A=a), Z=z] over the discovered adjustment set, at the values in observation (supplied by name, e.g. {"Z": 1}).

With none of the discovered variables supplied — or an empty discovered set, or a stratum never logged — the marginal decision argmax_a E[Y | do(a)] is returned; a partial context raises.

Source code in src/causalrl/agents/mbrl.py

def act(self, observation: dict[str, Any]) -> int:
    """Contextual policy ``argmax_a Ehat[Y | do(A=a), Z=z]`` over the *discovered* adjustment
    set, at the values in ``observation`` (supplied by name, e.g. ``{"Z": 1}``).

    With none of the discovered variables supplied — or an empty discovered set, or a stratum
    never logged — the marginal decision ``argmax_a E[Y | do(a)]`` is returned; a partial
    context raises.
    """
    return _stratum_action(
        self._stratum_values, observation, self.adjustment, self._best_action
    )

causalrl.agents.scbandit.FixedSetThompsonSampling

Bases: _ArmSubsetThompsonSampling

Thompson sampling restricted to a single fixed intervention set (the naive baseline).

Source code in src/causalrl/agents/scbandit.py

class FixedSetThompsonSampling(_ArmSubsetThompsonSampling):
    """Thompson sampling restricted to a single fixed intervention set (the naive baseline)."""

    def __init__(
        self, arms: list[dict[str, int]], intervention_set: Iterable[str], seed: int | None = None
    ) -> None:
        target = frozenset(intervention_set)
        allowed = [i for i, arm in enumerate(arms) if frozenset(arm.keys()) == target]
        super().__init__(allowed, seed)

causalrl.agents.mbrl.FunctionApproxBackdoorAgent

Bases: BatchAgent

Function-approximation back-door: fit qhat(a, z) by ridge regression on RBF features of a CONTINUOUS confounder, then back-door-adjust by Monte-Carlo integrating qhat(a, .) over the observed confounder sample, and ship the argmax. The M3 upgrade that carries the deconfounding recipe past the tabular regime to a learned continuous estimator.

The single continuous confounder is read from graph via :func:~causalrl.backdoor_adjustment_set. Fit on columnar {confounder, treatment, outcome} via :meth:fit.

The fitted qhat(a, .) is kept, so :meth:act evaluates it at the observed z and returns argmax_a qhat(a, z) — the same model the back-door integral averages, read per unit instead of summed away. A thin reward bump can make an arm interventionally worse on average and still optimal inside the bump; only a contextual policy can play it there.

Source code in src/causalrl/agents/mbrl.py

class FunctionApproxBackdoorAgent(BatchAgent):
    """Function-approximation back-door: fit ``qhat(a, z)`` by ridge regression on RBF features of a
    CONTINUOUS confounder, then back-door-adjust by Monte-Carlo integrating ``qhat(a, .)`` over the
    observed confounder sample, and ship the argmax. The M3 upgrade that carries the deconfounding
    recipe past the tabular regime to a learned continuous estimator.

    The single continuous confounder is read from ``graph`` via
    :func:`~causalrl.backdoor_adjustment_set`. Fit on columnar
    ``{confounder, treatment, outcome}`` via :meth:`fit`.

    The fitted ``qhat(a, .)`` is kept, so :meth:`act` evaluates it at the observed ``z`` and returns
    ``argmax_a qhat(a, z)`` — the same model the back-door integral averages, read per unit instead
    of summed away. A thin reward bump can make an arm interventionally *worse* on average and still
    optimal inside the bump; only a contextual policy can play it there.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        graph: CausalGraph,
        treatment: str = "A",
        outcome: str = "Y",
        n_centers: int = 12,
        bandwidth: float = 0.08,
        ridge: float = 1e-2,
    ) -> None:
        adjustment = sorted(backdoor_adjustment_set(graph, treatment, outcome))
        if len(adjustment) != 1:
            raise ValueError(
                "FunctionApproxBackdoorAgent supports a single continuous confounder; "
                f"graph gives adjustment set {adjustment}"
            )
        self.n_actions = n_actions
        self.treatment = treatment
        self.outcome = outcome
        self.confounder = adjustment[0]
        self.n_centers = n_centers
        self.bandwidth = bandwidth
        self.ridge = ridge
        self._best_action = 0
        self.values: list[float] = [0.0] * n_actions
        self._centers: np.ndarray = np.zeros(0)
        self._weights: list[np.ndarray | None] = [None] * n_actions

    def fit(self, data: Mapping[str, np.ndarray]) -> None:
        """Fit the per-action RBF-ridge outcome model, cache it for :meth:`act`, and select the
        back-door-adjusted (marginal) argmax."""
        z = np.asarray(data[self.confounder], dtype=float)
        a = np.asarray(data[self.treatment])
        y = np.asarray(data[self.outcome], dtype=float)
        self._centers = np.linspace(float(z.min()), float(z.max()), self.n_centers)
        self._weights = [
            self._fit_action_weights(action, z, a, y) for action in range(self.n_actions)
        ]
        # The back-door value is qhat(a, .) Monte-Carlo integrated over the observed confounder.
        self.values = list(self._action_values(_rbf_features(z, self._centers, self.bandwidth)))
        self._best_action = int(np.argmax(np.asarray(self.values)))

    def _fit_action_weights(
        self, action: int, z: np.ndarray, a: np.ndarray, y: np.ndarray
    ) -> np.ndarray | None:
        """Ridge weights of ``qhat(action, .)`` over the RBF design, fit on that action's rows
        (``None`` when the action is never played)."""
        mask = a == action
        if not mask.any():
            return None
        phi = _rbf_features(z[mask], self._centers, self.bandwidth)
        gram = phi.T @ phi + self.ridge * np.eye(phi.shape[1])
        return np.linalg.solve(gram, phi.T @ y[mask])

    def _action_values(self, phi: np.ndarray) -> list[float]:
        """Each action's mean ``qhat(a, .)`` over the design rows ``phi``.

        Over the whole confounder sample this is the back-door-adjusted value ``E[Y | do(a)]``;
        over a single row it is the per-unit ``qhat(a, z)``. A never-played action scores 0.0, as
        in the back-door formula.
        """
        return [0.0 if w is None else float((phi @ w).mean()) for w in self._weights]

    def act(self, observation: dict[str, Any]) -> int:
        """Contextual policy ``argmax_a qhat(a, z)`` at the confounder value in ``observation``
        (supplied by name, e.g. ``{"Z": 0.85}``).

        This is the fitted outcome model read at one point rather than integrated over the sample,
        so the arm an interventional *average* rejects can still be played where it wins. Without
        the confounder in ``observation`` (or before :meth:`fit`) the marginal decision
        ``argmax_a E[Y | do(a)]`` is returned.
        """
        if self._centers.size == 0 or self.confounder not in observation:
            return int(self._best_action)
        z = np.asarray([float(observation[self.confounder])], dtype=float)
        at_z = self._action_values(_rbf_features(z, self._centers, self.bandwidth))
        return int(np.argmax(np.asarray(at_z)))

fit(data)

Fit the per-action RBF-ridge outcome model, cache it for :meth:act, and select the back-door-adjusted (marginal) argmax.

Source code in src/causalrl/agents/mbrl.py

def fit(self, data: Mapping[str, np.ndarray]) -> None:
    """Fit the per-action RBF-ridge outcome model, cache it for :meth:`act`, and select the
    back-door-adjusted (marginal) argmax."""
    z = np.asarray(data[self.confounder], dtype=float)
    a = np.asarray(data[self.treatment])
    y = np.asarray(data[self.outcome], dtype=float)
    self._centers = np.linspace(float(z.min()), float(z.max()), self.n_centers)
    self._weights = [
        self._fit_action_weights(action, z, a, y) for action in range(self.n_actions)
    ]
    # The back-door value is qhat(a, .) Monte-Carlo integrated over the observed confounder.
    self.values = list(self._action_values(_rbf_features(z, self._centers, self.bandwidth)))
    self._best_action = int(np.argmax(np.asarray(self.values)))

act(observation)

Contextual policy argmax_a qhat(a, z) at the confounder value in observation (supplied by name, e.g. {"Z": 0.85}).

This is the fitted outcome model read at one point rather than integrated over the sample, so the arm an interventional average rejects can still be played where it wins. Without the confounder in observation (or before :meth:fit) the marginal decision argmax_a E[Y | do(a)] is returned.

Source code in src/causalrl/agents/mbrl.py

def act(self, observation: dict[str, Any]) -> int:
    """Contextual policy ``argmax_a qhat(a, z)`` at the confounder value in ``observation``
    (supplied by name, e.g. ``{"Z": 0.85}``).

    This is the fitted outcome model read at one point rather than integrated over the sample,
    so the arm an interventional *average* rejects can still be played where it wins. Without
    the confounder in ``observation`` (or before :meth:`fit`) the marginal decision
    ``argmax_a E[Y | do(a)]`` is returned.
    """
    if self._centers.size == 0 or self.confounder not in observation:
        return int(self._best_action)
    z = np.asarray([float(observation[self.confounder])], dtype=float)
    at_z = self._action_values(_rbf_features(z, self._centers, self.bandwidth))
    return int(np.argmax(np.asarray(at_z)))

causalrl.agents.mbrl.GFormulaBackdoorAgent

Bases: BatchAgent

Multivariate g-formula (standardization) back-door for many, mixed-type covariates.

Fits a per-action outcome model Ehat_a[Y | X] over ALL covariates X (a T-learner), then standardizes E[Y | do(a)] = mean_i Ehat_a(X_i) and ships the argmax. This is the strategy for the realistic confounding of real datasets, where per-stratum adjustment degenerates (each stratum has ~1 row) and a single-confounder RBF is too narrow. The default outcome model is a numpy ridge over [1, standardized X] (dependency-free); pass outcome_model — a factory returning a fresh sklearn-style estimator (fit/predict) — for a flexible offline model.

:meth:fit keeps the per-action models, so :meth:act is the policy the T-learner already implies: argmax_a Ehat_a(x), whose sign for a binary treatment is exactly the sign of :meth:cate. That is the same CATE-to-policy conversion :func:~causalrl.interop.econml.policy_from_econml_cate performs for a third-party estimator.

Source code in src/causalrl/agents/mbrl.py

class GFormulaBackdoorAgent(BatchAgent):
    """Multivariate g-formula (standardization) back-door for many, mixed-type covariates.

    Fits a per-action outcome model ``Ehat_a[Y | X]`` over ALL covariates X (a T-learner), then
    standardizes ``E[Y | do(a)] = mean_i Ehat_a(X_i)`` and ships the argmax. This is the strategy
    for the realistic confounding of real datasets, where per-stratum adjustment degenerates (each
    stratum has ~1 row) and a single-confounder RBF is too narrow. The default outcome model is a
    numpy ridge over ``[1, standardized X]`` (dependency-free); pass ``outcome_model`` — a factory
    returning a fresh sklearn-style estimator (``fit``/``predict``) — for a flexible offline model.

    :meth:`fit` keeps the per-action models, so :meth:`act` is the *policy the T-learner already
    implies*: ``argmax_a Ehat_a(x)``, whose sign for a binary treatment is exactly the sign of
    :meth:`cate`. That is the same CATE-to-policy conversion
    :func:`~causalrl.interop.econml.policy_from_econml_cate` performs for a third-party estimator.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        covariates: Sequence[str],
        treatment: str = "A",
        outcome: str = "Y",
        outcome_model: Callable[[], Any] | None = None,
        ridge: float = 1.0,
    ) -> None:
        self.n_actions = n_actions
        self.covariates = tuple(covariates)
        self.treatment = treatment
        self.outcome = outcome
        self._outcome_model = outcome_model
        self.ridge = ridge
        self._best_action = 0
        self.values: list[float] = [0.0] * n_actions
        self._models: list[_OutcomeModel | None] = [None] * n_actions

    def fit(self, data: Mapping[str, np.ndarray]) -> GFormulaBackdoorAgent:
        """Fit the per-action outcome models, keep them for the contextual :meth:`act`, standardize
        to ``E[Y|do(a)]``, and select the marginal argmax."""
        covariate_matrix = self._covariate_matrix(data)
        treatment = np.asarray(data[self.treatment])
        outcome = np.asarray(data[self.outcome], dtype=float)
        self._models = self._fit_action_models(covariate_matrix, treatment, outcome)
        self.values = [_mean_prediction(model, covariate_matrix) for model in self._models]
        self._best_action = int(np.argmax(np.asarray(self.values)))
        return self

    def _covariate_matrix(self, data: Mapping[str, np.ndarray]) -> np.ndarray:
        return np.column_stack([np.asarray(data[c], dtype=float) for c in self.covariates])

    def _fit_action_models(
        self, x: np.ndarray, treatment: np.ndarray, outcome: np.ndarray
    ) -> list[_OutcomeModel | None]:
        """Fit one outcome model per action on that action's rows (``None`` if never played)."""
        return [self._fit_action_model(a, x, treatment, outcome) for a in range(self.n_actions)]

    def _fit_action_model(
        self, action: int, x: np.ndarray, treatment: np.ndarray, outcome: np.ndarray
    ) -> _OutcomeModel | None:
        mask = treatment == action
        if not mask.any():
            return None
        x_a, y_a = x[mask], outcome[mask]
        if self._outcome_model is not None:
            model = self._outcome_model()
            model.fit(x_a, y_a)
            return model
        mean, std = _standardization(x_a)
        phi_a = _ridge_design(x_a, mean, std)
        gram = phi_a.T @ phi_a + self.ridge * np.eye(phi_a.shape[1])
        weights = np.linalg.solve(gram, phi_a.T @ y_a)
        return _RidgeOutcomeModel(weights=weights, mean=mean, std=std)

    def cate(self, data: Mapping[str, np.ndarray]) -> np.ndarray:
        """Per-unit CATE ``Ehat_1(X_i) - Ehat_0(X_i)`` for each row of ``data`` (the T-learner
        individual-effect estimate; binary treatment). ``data`` carries the covariates plus the
        observed ``treatment``/``outcome`` the per-action models are fit on."""
        if self.n_actions != 2:
            raise ValueError("cate is defined only for a binary treatment")
        x = self._covariate_matrix(data)
        treatment = np.asarray(data[self.treatment])
        outcome = np.asarray(data[self.outcome], dtype=float)
        models = self._fit_action_models(x, treatment, outcome)
        mu0, mu1 = _predict(models[0], x), _predict(models[1], x)
        if mu0 is None or mu1 is None:
            raise ValueError("both treatment arms must appear in the data to estimate CATE")
        return mu1 - mu0

    @property
    def contrast(self) -> float:
        """``E[Y|do(1)] - E[Y|do(0)]`` (the ATE) for a binary treatment."""
        if self.n_actions != 2:
            raise ValueError("contrast is defined only for a binary treatment")
        return self.values[1] - self.values[0]

    def act(self, observation: dict[str, Any]) -> int:
        """Contextual policy ``argmax_a Ehat_a(x)`` at the covariates in ``observation`` (supplied
        by name, e.g. ``{"age": 42, "smoke": 1}``).

        For a binary treatment this is exactly the sign of :meth:`cate` at that row — the per-unit
        decision the T-learner already estimates, shipped instead of discarded. With none of
        :attr:`covariates` in ``observation`` (or before :meth:`fit`) the marginal decision
        ``argmax_a E[Y | do(a)]`` is returned; a partial covariate vector raises rather than
        silently answering a different query.
        """
        key = _context_key(observation, self.covariates)
        if key is None or all(model is None for model in self._models):
            return int(self._best_action)
        x = np.asarray([[float(value) for value in key]], dtype=float)
        return int(np.argmax(np.asarray([_mean_prediction(m, x) for m in self._models])))

contrast property

E[Y|do(1)] - E[Y|do(0)] (the ATE) for a binary treatment.

fit(data)

Fit the per-action outcome models, keep them for the contextual :meth:act, standardize to E[Y|do(a)], and select the marginal argmax.

Source code in src/causalrl/agents/mbrl.py

def fit(self, data: Mapping[str, np.ndarray]) -> GFormulaBackdoorAgent:
    """Fit the per-action outcome models, keep them for the contextual :meth:`act`, standardize
    to ``E[Y|do(a)]``, and select the marginal argmax."""
    covariate_matrix = self._covariate_matrix(data)
    treatment = np.asarray(data[self.treatment])
    outcome = np.asarray(data[self.outcome], dtype=float)
    self._models = self._fit_action_models(covariate_matrix, treatment, outcome)
    self.values = [_mean_prediction(model, covariate_matrix) for model in self._models]
    self._best_action = int(np.argmax(np.asarray(self.values)))
    return self

cate(data)

Per-unit CATE Ehat_1(X_i) - Ehat_0(X_i) for each row of data (the T-learner individual-effect estimate; binary treatment). data carries the covariates plus the observed treatment/outcome the per-action models are fit on.

Source code in src/causalrl/agents/mbrl.py

def cate(self, data: Mapping[str, np.ndarray]) -> np.ndarray:
    """Per-unit CATE ``Ehat_1(X_i) - Ehat_0(X_i)`` for each row of ``data`` (the T-learner
    individual-effect estimate; binary treatment). ``data`` carries the covariates plus the
    observed ``treatment``/``outcome`` the per-action models are fit on."""
    if self.n_actions != 2:
        raise ValueError("cate is defined only for a binary treatment")
    x = self._covariate_matrix(data)
    treatment = np.asarray(data[self.treatment])
    outcome = np.asarray(data[self.outcome], dtype=float)
    models = self._fit_action_models(x, treatment, outcome)
    mu0, mu1 = _predict(models[0], x), _predict(models[1], x)
    if mu0 is None or mu1 is None:
        raise ValueError("both treatment arms must appear in the data to estimate CATE")
    return mu1 - mu0

act(observation)

Contextual policy argmax_a Ehat_a(x) at the covariates in observation (supplied by name, e.g. {"age": 42, "smoke": 1}).

For a binary treatment this is exactly the sign of :meth:cate at that row — the per-unit decision the T-learner already estimates, shipped instead of discarded. With none of :attr:covariates in observation (or before :meth:fit) the marginal decision argmax_a E[Y | do(a)] is returned; a partial covariate vector raises rather than silently answering a different query.

Source code in src/causalrl/agents/mbrl.py

def act(self, observation: dict[str, Any]) -> int:
    """Contextual policy ``argmax_a Ehat_a(x)`` at the covariates in ``observation`` (supplied
    by name, e.g. ``{"age": 42, "smoke": 1}``).

    For a binary treatment this is exactly the sign of :meth:`cate` at that row — the per-unit
    decision the T-learner already estimates, shipped instead of discarded. With none of
    :attr:`covariates` in ``observation`` (or before :meth:`fit`) the marginal decision
    ``argmax_a E[Y | do(a)]`` is returned; a partial covariate vector raises rather than
    silently answering a different query.
    """
    key = _context_key(observation, self.covariates)
    if key is None or all(model is None for model in self._models):
        return int(self._best_action)
    x = np.asarray([[float(value) for value in key]], dtype=float)
    return int(np.argmax(np.asarray([_mean_prediction(m, x) for m in self._models])))

causalrl.intervention.Intervention = Mapping[str, Any] module-attribute

An assignment {variable: value} — the argument StructuralCausalModel.do accepts.

The empty mapping is the observational regime (do() on nothing), matching the frozenset() that :func:causalrl.pomis emits when not intervening is possibly optimal.

causalrl.magames.KindNotLicensedError

Bases: TypeError

Raised when a stronger Kind is requested than the learner topology licenses (I2).

Source code in src/causalrl/magames/equilibrium.py

class KindNotLicensedError(TypeError):
    """Raised when a stronger ``Kind`` is requested than the learner topology licenses (I2)."""

causalrl.magames.LearnerTopology

Bases: Enum

How a population learns; caps the strongest Kind a certificate about it may claim.

Source code in src/causalrl/magames/population.py

class LearnerTopology(enum.Enum):
    """How a population learns; caps the strongest ``Kind`` a certificate about it may claim."""

    SINGLE_LEARNER = "single_learner"  # one learner, fixed population -> IDENTIFIED-capable
    INDEPENDENT_LEARNERS = "independent_learners"  # simultaneous learners -> EMPIRICAL only
    CENTRALIZED = "centralized"  # centralised training -> EMPIRICAL only

causalrl.envs.suite.mabuc.MABUCEnv

Bases: CausalEnv

One-step bandit. Each episode: observe intuition I, choose arm X in {0,1}, receive reward Y. The optimal policy plays X = I (the lucky arm).

Source code in src/causalrl/envs/suite/mabuc.py

class MABUCEnv(CausalEnv):
    """One-step bandit. Each episode: observe intuition I, choose arm X in {0,1},
    receive reward Y. The optimal policy plays X = I (the lucky arm)."""

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]  # gymnasium Env.metadata is an instance var in the base

    def __init__(self, seed: int | None = None) -> None:
        super().__init__(build_mabuc_scm())
        self.action_space = gym.spaces.Discrete(2)
        self.observation_space = gym.spaces.Dict({"intuition": gym.spaces.Discrete(2)})
        self._rng: torch.Generator = torch.Generator()  # type: ignore[reportPrivateImportUsage]
        if seed is not None:
            self._rng.manual_seed(seed)
        self._d = 0
        self._b = 0

    def _draw_confounders(self) -> None:
        self._d = int(torch.randint(0, 2, (1,), generator=self._rng).item())  # type: ignore[reportPrivateImportUsage]
        self._b = int(torch.randint(0, 2, (1,), generator=self._rng).item())  # type: ignore[reportPrivateImportUsage]

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, int], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng.manual_seed(seed)
        self._draw_confounders()
        intuition = self._d ^ self._b
        return {"intuition": intuition}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, int], float, bool, bool, dict[str, Any]]:
        lucky = self._d ^ self._b
        p = _P_WIN if action == lucky else _P_LOSE
        reward = float(torch.rand(1, generator=self._rng).item() < p)  # type: ignore[reportPrivateImportUsage]
        intuition = self._d ^ self._b
        return {"intuition": intuition}, reward, True, False, {"lucky": lucky}

causalrl.agents.baselines.NaiveOffline

Bases: BatchAgent

Fits E[R|s,a] from the logs as if unconfounded and acts greedily. Provably biased under confounding — the cautionary baseline.

Source code in src/causalrl/agents/baselines.py

class NaiveOffline(BatchAgent):
    """Fits E[R|s,a] from the logs as if unconfounded and acts greedily. Provably biased
    under confounding — the cautionary baseline."""

    def __init__(self, n_states: int, n_actions: int) -> None:
        self.n_states = n_states
        self.n_actions = n_actions
        self._mean = np.zeros((n_states, n_actions))

    def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
        for s in range(self.n_states):
            for a in range(self.n_actions):
                self._mean[s, a] = dataset.mean_reward(s, a)

    def act(self, observation: dict[str, Any]) -> int:
        s = int(observation["state"])
        return int(np.argmax(np.asarray(self._mean[s], dtype=np.float64)))

causalrl.agents.scbandit.NaivePOMISThompsonSampling

Bases: _ArmSubsetThompsonSampling

Naive baseline that ignores manipulability: it computes the UNCONSTRAINED POMIS and keeps only the arms that happen to exist. When the optimal lever lies outside the unconstrained POMIS — as when a non-manipulable variable dominates — it cannot reach the optimum (r40 Prop. 1: filtering the unconstrained POMIS is insufficient).

Source code in src/causalrl/agents/scbandit.py

class NaivePOMISThompsonSampling(_ArmSubsetThompsonSampling):
    """Naive baseline that ignores manipulability: it computes the UNCONSTRAINED POMIS and
    keeps only the arms that happen to exist. When the optimal lever lies outside the
    unconstrained POMIS — as when a non-manipulable variable dominates — it cannot reach the
    optimum (r40 Prop. 1: filtering the unconstrained POMIS is insufficient)."""

    def __init__(
        self, graph: CausalGraph, reward: str, arms: list[dict[str, int]], seed: int | None = None
    ) -> None:
        pomis_sets = set(pomis(graph, reward))  # unconstrained — ignores manipulability
        allowed = [i for i, arm in enumerate(arms) if frozenset(arm.keys()) in pomis_sets]
        super().__init__(allowed, seed)

causalrl.agents.bandits.NaiveThompsonSampling

Bases: Agent

Beta-Bernoulli Thompson sampling, one posterior per arm. Ignores intuition, so it cannot distinguish the arms (their interventional means are equal).

Sampling uses a per-instance numpy.random.Generator seeded from seed, so two agents with different seeds draw independent action sequences and a fixed seed is reproducible regardless of global RNG state.

Source code in src/causalrl/agents/bandits.py

class NaiveThompsonSampling(Agent):
    """Beta-Bernoulli Thompson sampling, one posterior per arm. Ignores intuition,
    so it cannot distinguish the arms (their interventional means are equal).

    Sampling uses a per-instance ``numpy.random.Generator`` seeded from ``seed``, so two
    agents with different seeds draw independent action sequences and a fixed seed is
    reproducible regardless of global RNG state.
    """

    def __init__(self, n_arms: int, seed: int | None = None) -> None:
        self.n_arms = n_arms
        self._alpha = np.ones(n_arms)
        self._beta = np.ones(n_arms)
        self._rng = np.random.default_rng(seed)

    def act(self, observation: dict[str, Any]) -> int:
        samples = np.asarray(self._rng.beta(self._alpha, self._beta), dtype=np.float64)
        return int(np.argmax(samples))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        if reward > 0:
            self._alpha[action] += 1.0
        else:
            self._beta[action] += 1.0

causalrl.agents.online_causal_mbrl.OnlineCausalMBRL

An agent that learns its SCM from its own interventions plus an off-policy log.

variables are the columns every buffered row must carry; treatment is the node the agent intervenes on, outcome the node it is scored by, and actions the values it may set the treatment to. Data arrives through :meth:ingest (bulk, off-policy) or :meth:observe (one transition, on-policy), and :meth:refit turns the buffers into the belief:

.. code-block:: text

discover_interventional(observational, interventional, variables)  ->  CPDAG (the I-MEC)
fit_scm_mec(all rows, cpdag=..., max_members=...)  ->  one fitted SCM per member

:meth:act turns the belief into an action from actions -- policy decides how the members' verdicts are combined and n_rollout how many samples each verdict costs -- and :meth:probe returns the intervention target the members disagree about most. The loop those two are meant for branches on :meth:structure_uncertain::

target = agent.probe() if agent.structure_uncertain() else None
action = agent.act() if target is None else None

refit_every is the cadence that loop should call :meth:refit at -- rerunning PC every step is wasteful and makes the belief flicker on sampling noise -- and is not self-triggering.

max_members is handed to :func:~causalrl.scm.fit.fit_scm_mec, whose refusal above the cap propagates out of :meth:refit deliberately: an equivalence class too large to enumerate is information the caller needs, and a silently truncated belief is a misreported one.

Source code in src/causalrl/agents/online_causal_mbrl.py

class OnlineCausalMBRL:
    """An agent that learns its SCM from its own interventions plus an off-policy log.

    ``variables`` are the columns every buffered row must carry; ``treatment`` is the node the
    agent intervenes on, ``outcome`` the node it is scored by, and ``actions`` the values it may
    set the treatment to. Data arrives through :meth:`ingest` (bulk, off-policy) or :meth:`observe`
    (one transition, on-policy), and :meth:`refit` turns the buffers into the belief:

    .. code-block:: text

        discover_interventional(observational, interventional, variables)  ->  CPDAG (the I-MEC)
        fit_scm_mec(all rows, cpdag=..., max_members=...)  ->  one fitted SCM per member

    :meth:`act` turns the belief into an action from ``actions`` -- ``policy`` decides how the
    members' verdicts are combined and ``n_rollout`` how many samples each verdict costs -- and
    :meth:`probe` returns the intervention target the members disagree about most. The loop those
    two are meant for branches on :meth:`structure_uncertain`::

        target = agent.probe() if agent.structure_uncertain() else None
        action = agent.act() if target is None else None

    ``refit_every`` is the cadence that loop should call :meth:`refit` at -- rerunning PC every
    step is wasteful and makes the belief flicker on sampling noise -- and is not self-triggering.

    ``max_members`` is handed to :func:`~causalrl.scm.fit.fit_scm_mec`, whose refusal above the cap
    propagates out of :meth:`refit` deliberately: an equivalence class too large to enumerate is
    information the caller needs, and a silently truncated belief is a misreported one.
    """

    def __init__(
        self,
        variables: Sequence[str],
        *,
        treatment: str,
        outcome: str,
        actions: Sequence[float],
        policy: Policy = "thompson",
        max_members: int = 32,
        refit_every: int = 8,
        n_rollout: int = 512,
        min_interventional_samples: int = 20,
        seed: int = 0,
    ) -> None:
        self.variables = tuple(variables)
        if treatment not in self.variables:
            raise ValueError(f"treatment={treatment!r} is not one of variables={self.variables}")
        if outcome not in self.variables:
            raise ValueError(f"outcome={outcome!r} is not one of variables={self.variables}")
        if not actions:
            raise ValueError("actions is empty: there is nothing for the agent to choose between")
        if policy not in _POLICIES:
            raise ValueError(f"policy={policy!r} is not one of {list(_POLICIES)}")
        if n_rollout < 1:
            # Every action value is the mean of a `see(n_rollout)` sample. At zero that mean is
            # nan, nan comparisons are all False, and argmax silently returns the first action --
            # a decision that looks like a choice and is not one.
            raise ValueError(
                f"n_rollout={n_rollout} must be at least 1: action values are means over "
                f"see(n_rollout) draws, and an empty draw makes every value nan, which argmax "
                f"resolves to the first action rather than reporting that nothing was compared."
            )
        self.treatment = treatment
        self.outcome = outcome
        self.actions = tuple(actions)
        self.policy: Policy = policy
        self.max_members = max_members
        self.refit_every = refit_every
        self.n_rollout = n_rollout
        self.min_interventional_samples = min_interventional_samples
        self.seed = seed

        # Columnar buffers: one list per variable, appended to in place. Rebuilding columns from
        # per-row dicts on every refit is the same data at a much worse cost.
        self._observational: dict[str, list[float]] = {v: [] for v in self.variables}
        self._interventional: dict[str, dict[str, list[float]]] = {}
        self._belief: tuple[StructuralCausalModel, ...] = ()
        self._cpdag: CPDAG | None = None
        self._history: list[RefitRecord] = []
        self._steps = 0
        self._draws = 0

    # -- data in ------------------------------------------------------------------------------
    def ingest(
        self,
        data: Mapping[str, np.ndarray],
        *,
        source: Source = "observational",
        target: str | None = None,
    ) -> None:
        """Append a bulk log to the observational buffer, or to ``target``'s do-buffer.

        ``source="interventional"`` requires ``target`` (the node that was set) and
        ``source="observational"`` forbids it. Every variable must be present with equal lengths:
        :func:`~causalrl.discovery.discover_interventional` compares the *marginal of every
        endpoint* between regimes, so a do-sample holding only the reward cannot orient anything.
        """
        if source == "interventional":
            if target is None:
                raise ValueError(
                    "ingest(source='interventional') requires target=<variable>: "
                    "discover_interventional orients the edges incident to each intervention "
                    "target separately, so a do-sample must name the node that was set."
                )
            buffer = self._buffer_for(target)
        elif source == "observational":
            if target is not None:
                raise ValueError(
                    f"ingest(source='observational') forbids target={target!r}: an observational "
                    "sample has no intervention target. Pass source='interventional' to route it "
                    "to that target's buffer."
                )
            buffer = self._observational
        else:
            raise ValueError(f"source={source!r} is not one of {list(_SOURCES)}")
        self._steps += _extend(buffer, data, self.variables)

    def observe(
        self,
        row: Mapping[str, float],
        *,
        intervention: Mapping[str, float] | None = None,
    ) -> None:
        """Append one transition -- to the observational buffer, or to its target's do-buffer.

        With ``intervention``, the row joins the buffer for that target and the do-value is what is
        recorded for it: under a perfect intervention the target was *set*, so ``row`` may omit it.
        Exactly one target is supported per row, because
        :func:`~causalrl.discovery.discover_interventional` orients per single target.
        """
        values = dict(row)
        if intervention is None:
            buffer = self._observational
        else:
            if len(intervention) != 1:
                raise ValueError(
                    f"observe() supports exactly one target per row, got "
                    f"{sorted(intervention)}: discover_interventional orients the edges incident "
                    "to a single intervention target, so a multi-node do() has no buffer to join."
                )
            target = next(iter(intervention))
            buffer = self._buffer_for(target)
            values.update(intervention)
        missing = [v for v in self.variables if v not in values]
        if missing:
            raise ValueError(
                f"row is missing variable(s) {missing}: discover_interventional needs every "
                "variable in every dataset it is handed, so buffers hold full rows."
            )
        for variable in self.variables:
            buffer[variable].append(float(values[variable]))
        self._steps += 1

    def _rows(self, buffer: Mapping[str, list[float]]) -> int:
        """How many transitions a buffer holds (its columns are appended to in lockstep)."""
        return len(buffer[self.variables[0]])

    def _buffer_for(self, target: str) -> dict[str, list[float]]:
        if target not in self.variables:
            raise ValueError(
                f"intervention target {target!r} is not a variable of this agent: {self.variables}"
            )
        if target not in self._interventional:
            self._interventional[target] = {v: [] for v in self.variables}
        return self._interventional[target]

    # -- learning -----------------------------------------------------------------------------
    def refit(self) -> tuple[StructuralCausalModel, ...]:
        """Recompute the I-MEC from the buffers and fit one SCM per member; return the belief.

        Two roles for the same rows, and they are not interchangeable. **Orientation** comes from
        the do-data: :func:`~causalrl.discovery.discover_interventional` runs PC on the
        observational sample for the skeleton and then orients the edges incident to each target by
        the invariance principle, which only the interventional marginals can decide.
        **Mechanisms** are fitted from every row the agent holds, observational and interventional
        concatenated, because a mechanism is invariant across regimes (Peters et al. 2016) and the
        do-rows are extra samples of every node that was not intervened on. The intervened node's
        own mechanism is fitted from rows where it was overridden, which biases *its* marginal --
        harmless here, since planning replaces that mechanism with ``do(treatment=a)`` anyway.

        A target holding fewer than ``min_interventional_samples`` rows is dropped rather than
        passed on. The invariance test compares empirical marginals, so a small do-sample clears
        any shift threshold by chance and an empty one clears it unconditionally — orienting every
        edge incident to that target from evidence the agent does not have.
        :func:`~causalrl.discovery.discover_interventional` refuses such a sample outright; the
        agent filters first so that an early round, when a target has been probed once or twice, is
        an ordinary under-informed refit rather than a raise.

        Raises ``ValueError`` if no observational rows have been buffered (PC needs them for the
        base CPDAG), and lets :func:`~causalrl.scm.fit.fit_scm_mec`'s over-size refusal propagate.
        """
        if not self._rows(self._observational):
            raise ValueError(
                "refit() needs observational rows: discover_interventional runs PC on the "
                "observational sample to build the base CPDAG before the invariance principle "
                "orients edges incident to each intervention target. Ingest an observational log "
                "(or call observe() without intervention=...) first."
            )
        observational = _columns(self._observational)
        interventions = {
            target: _columns(buffer)
            for target, buffer in sorted(self._interventional.items())
            if self._rows(buffer) >= self.min_interventional_samples
        }
        cpdag = discover_interventional(
            observational,
            interventions,
            self.variables,
            min_interventional_samples=self.min_interventional_samples,
        )
        pooled = {
            variable: np.concatenate(
                [observational[variable], *(data[variable] for data in interventions.values())]
            )
            for variable in self.variables
        }
        self._belief = tuple(
            fit_scm_mec(pooled, cpdag=cpdag, max_members=self.max_members, seed=self.seed)
        )
        self._cpdag = cpdag
        self._history.append(RefitRecord(step=self._steps, belief_size=len(self._belief)))
        return self._belief

    # -- acting -------------------------------------------------------------------------------
    def act(self) -> float:
        """Choose an action from ``actions`` by planning inside the belief.

        Every policy scores an action the same way -- ``member.do({treatment: action})``, then
        ``see(n_rollout)`` and the mean of ``outcome`` -- and they differ only in how the members'
        verdicts are combined:

        - ``"thompson"`` draws one member and returns its argmax. Thompson sampling over
          structure (Ortega & Braun, `arXiv:1303.4431 <https://arxiv.org/abs/1303.4431>`_): the
          spread of the belief supplies the exploration, so the policy sharpens exactly as the
          I-MEC collapses and needs no separate schedule.
        - ``"average"`` takes the mean value across members -- marginalising the structure out.
        - ``"robust"`` takes the maximin: score each action by its *worst* member, then take the
          best of those. The one to reach for when a wrong action is expensive, since it is the
          only one of the three that never bets on a particular member being the true DAG.

        Deterministic given ``seed``. The thompson draw comes from an RNG derived from ``seed``,
        the transition count and the number of draws already taken, so a fixed seed replays the
        whole decision sequence, successive calls draw successive members rather than repeating
        one, and two seeds diverge. Rollouts are seeded per member, which also hands every action
        of a member the same exogenous draws (common random numbers), so the gap between two
        actions is not partly the gap between two samples. Ties go to the earliest action in
        ``actions``.

        **This costs rollouts:** ``n_rollout`` samples per action per member scored, so
        ``len(actions) * n_rollout`` under ``"thompson"`` and
        ``len(actions) * belief_size() * n_rollout`` under ``"average"`` and ``"robust"``, every
        call. That is what planning in a model costs instead of reading a value table.

        Raises ``ValueError`` before the first :meth:`refit`, when there is no belief to plan in.
        """
        belief = self._require_belief("act")
        if self.policy == "thompson":
            rng = np.random.default_rng((self.seed, self._steps, self._draws))
            self._draws += 1
            drawn = int(rng.integers(len(belief)))
            values = [
                self._action_value(belief[drawn], action, seed=self.seed + drawn)
                for action in self.actions
            ]
        else:
            by_member = [
                [
                    self._action_value(member, action, seed=self.seed + index)
                    for action in self.actions
                ]
                for index, member in enumerate(belief)
            ]
            across = list(zip(*by_member, strict=True))
            values = (
                [float(np.mean(column)) for column in across]
                if self.policy == "average"
                else [min(column) for column in across]
            )
        return self.actions[max(range(len(self.actions)), key=lambda index: values[index])]

    def probe(self) -> str:
        """Return the intervention target the belief's members disagree about most.

        A **target**, not an action: the answer to "what experiment should I run next", which is
        the other half of the loop :meth:`act` serves. Each candidate is scored by the mean
        pairwise total variation between the members' outcome distributions under ``do(candidate)``
        -- predictive disagreement standing in for expected information gain, as in Scherrer et
        al. (`arXiv:2109.02429 <https://arxiv.org/abs/2109.02429>`_) and Zhang et al. (Nature
        Machine Intelligence 2023). It concentrates experiments on the edges the I-MEC still leaves
        unoriented, because those are the only ones whose ``do()`` distribution the members can
        disagree about; a target every member already predicts identically scores zero and is
        excluded, and if *every* candidate scores zero this raises rather than returning an
        arbitrary one.

        Each candidate is intervened at one representative value: the largest it takes anywhere in
        the agent's buffers. That keeps the do-value inside the observed support, so no fitted
        mechanism is asked to extrapolate, while sitting as far as the data allows from the bottom
        of the range -- a do-value near the middle of a variable's spread perturbs its descendants
        least and makes members that disagree look alike. One value cannot certify agreement,
        though: a zero score means "no disagreement detected at this value", not "these members
        agree everywhere".

        **This costs rollouts:** ``n_rollout`` samples per candidate per member, i.e.
        ``len(variables) * belief_size() * n_rollout`` every call. All members share one exogenous
        stream per candidate (common random numbers), because the score is a difference *between*
        members and independent draws would put a finite-sample floor under it.

        Raises ``ValueError`` before the first :meth:`refit`, on a single-member belief (nothing to
        disambiguate -- :meth:`structure_uncertain` is the guard to branch on), and when no
        candidate scores any disagreement at all.
        """
        belief = self._require_belief("probe")
        if len(belief) == 1:
            raise ValueError(
                "probe() has nothing to disambiguate: the belief holds a single I-MEC member, so "
                "every candidate target has exactly one implied do() distribution and no "
                "experiment can separate anything. structure_uncertain() is the guard to branch "
                "on -- probe() while it reports True, act() once it reports False."
            )
        scores = {
            candidate: self._disagreement(candidate, belief) for candidate in sorted(self.variables)
        }
        best = max(scores, key=lambda candidate: scores[candidate])
        if scores[best] <= 0.0:
            raise ValueError(
                f"probe() found no target the belief's {len(belief)} members disagree about: every "
                f"candidate in {sorted(self.variables)} scored zero total variation between the "
                "members' outcome distributions. Intervening on any of them would be an experiment "
                "whose result every member already predicts identically, so returning one would "
                "report an arbitrary pick as an informative choice. The members differ somewhere "
                "else -- in a mechanism, or at a do-value other than the one probe() uses -- and "
                "that is not something this score can locate."
            )
        return best

    def _require_belief(self, caller: str) -> tuple[StructuralCausalModel, ...]:
        """The belief, or a ``ValueError`` naming ``refit`` when none has been fitted."""
        if not self._belief:
            raise ValueError(
                f"{caller}() needs a belief and none has been fitted yet: call refit() first. An "
                "empty belief is not a structure the agent is certain about, it is one it knows "
                "nothing about, and the documented loop branches on structure_uncertain() -- so "
                "answering from an empty belief would route a caller who never fitted anything "
                "straight into act(), planning in a model that does not exist."
            )
        return self._belief

    def _action_value(self, member: StructuralCausalModel, action: float, *, seed: int) -> float:
        """``member``'s mean ``outcome`` under ``do(treatment=action)``, over ``n_rollout`` rows."""
        return float(np.mean(self._rollout(member, {self.treatment: action}, seed=seed)))

    def _disagreement(self, candidate: str, belief: tuple[StructuralCausalModel, ...]) -> float:
        """Mean pairwise total variation of the members' ``outcome`` under ``do(candidate)``."""
        value = self._probe_value(candidate)
        seed = self.seed + self.variables.index(candidate)
        samples = [self._rollout(member, {candidate: value}, seed=seed) for member in belief]
        return float(np.mean([_total_variation(a, b) for a, b in combinations(samples, 2)]))

    def _probe_value(self, candidate: str) -> float:
        """The largest value ``candidate`` takes across every buffer -- see :meth:`probe`."""
        values = list(self._observational[candidate])
        for buffer in self._interventional.values():
            values.extend(buffer[candidate])
        if not values:
            raise ValueError(
                f"probe() has no buffered row to take a do-value for {candidate!r} from: it "
                "intervenes at the largest value each candidate is observed to take, so that the "
                "do-value stays inside the support the mechanisms were fitted on. Ingest data "
                "before probing."
            )
        return max(values)

    def _rollout(
        self, member: StructuralCausalModel, intervention: Mapping[str, float], *, seed: int
    ) -> np.ndarray:
        """``n_rollout`` draws of ``outcome`` from ``member`` under ``do(intervention)``."""
        sample = member.do(intervention).see(self.n_rollout, seed=seed)
        return np.asarray(sample[self.outcome].detach().numpy(), dtype=float)

    # -- what the agent believes --------------------------------------------------------------
    def belief(self) -> tuple[StructuralCausalModel, ...]:
        """The fitted I-MEC members -- empty until the first :meth:`refit`."""
        return self._belief

    def belief_size(self) -> int:
        """How many DAGs the data still cannot choose between."""
        return len(self._belief)

    def structure_uncertain(self) -> bool:
        """Whether more than one I-MEC member survives -- a decision may need the certificate layer.

        Raises ``ValueError`` before the first :meth:`refit` rather than answering. With an empty
        belief neither answer is honest: there is nothing to be uncertain *between* yet, and
        nothing that says the structure is known either. This is the guard the documented loop
        branches on -- :meth:`probe` while it is ``True``, :meth:`act` once it is ``False`` -- so
        reporting ``False`` from an empty belief would quietly route a caller who never fitted
        anything into the exploit branch.
        """
        return len(self._require_belief("structure_uncertain")) > 1

    def history(self) -> tuple[RefitRecord, ...]:
        """An immutable snapshot of the belief trajectory, one :class:`RefitRecord` per refit."""
        return tuple(self._history)

    @property
    def cpdag(self) -> CPDAG | None:
        """The I-essential graph from the last :meth:`refit`; ``None`` before the first."""
        return self._cpdag

    @property
    def steps(self) -> int:
        """Transitions buffered so far, observational and interventional together."""
        return self._steps

cpdag property

The I-essential graph from the last :meth:refit; None before the first.

steps property

Transitions buffered so far, observational and interventional together.

ingest(data, *, source='observational', target=None)

Append a bulk log to the observational buffer, or to target's do-buffer.

source="interventional" requires target (the node that was set) and source="observational" forbids it. Every variable must be present with equal lengths: :func:~causalrl.discovery.discover_interventional compares the marginal of every endpoint between regimes, so a do-sample holding only the reward cannot orient anything.

Source code in src/causalrl/agents/online_causal_mbrl.py

def ingest(
    self,
    data: Mapping[str, np.ndarray],
    *,
    source: Source = "observational",
    target: str | None = None,
) -> None:
    """Append a bulk log to the observational buffer, or to ``target``'s do-buffer.

    ``source="interventional"`` requires ``target`` (the node that was set) and
    ``source="observational"`` forbids it. Every variable must be present with equal lengths:
    :func:`~causalrl.discovery.discover_interventional` compares the *marginal of every
    endpoint* between regimes, so a do-sample holding only the reward cannot orient anything.
    """
    if source == "interventional":
        if target is None:
            raise ValueError(
                "ingest(source='interventional') requires target=<variable>: "
                "discover_interventional orients the edges incident to each intervention "
                "target separately, so a do-sample must name the node that was set."
            )
        buffer = self._buffer_for(target)
    elif source == "observational":
        if target is not None:
            raise ValueError(
                f"ingest(source='observational') forbids target={target!r}: an observational "
                "sample has no intervention target. Pass source='interventional' to route it "
                "to that target's buffer."
            )
        buffer = self._observational
    else:
        raise ValueError(f"source={source!r} is not one of {list(_SOURCES)}")
    self._steps += _extend(buffer, data, self.variables)

observe(row, *, intervention=None)

Append one transition -- to the observational buffer, or to its target's do-buffer.

With intervention, the row joins the buffer for that target and the do-value is what is recorded for it: under a perfect intervention the target was set, so row may omit it. Exactly one target is supported per row, because :func:~causalrl.discovery.discover_interventional orients per single target.

Source code in src/causalrl/agents/online_causal_mbrl.py

def observe(
    self,
    row: Mapping[str, float],
    *,
    intervention: Mapping[str, float] | None = None,
) -> None:
    """Append one transition -- to the observational buffer, or to its target's do-buffer.

    With ``intervention``, the row joins the buffer for that target and the do-value is what is
    recorded for it: under a perfect intervention the target was *set*, so ``row`` may omit it.
    Exactly one target is supported per row, because
    :func:`~causalrl.discovery.discover_interventional` orients per single target.
    """
    values = dict(row)
    if intervention is None:
        buffer = self._observational
    else:
        if len(intervention) != 1:
            raise ValueError(
                f"observe() supports exactly one target per row, got "
                f"{sorted(intervention)}: discover_interventional orients the edges incident "
                "to a single intervention target, so a multi-node do() has no buffer to join."
            )
        target = next(iter(intervention))
        buffer = self._buffer_for(target)
        values.update(intervention)
    missing = [v for v in self.variables if v not in values]
    if missing:
        raise ValueError(
            f"row is missing variable(s) {missing}: discover_interventional needs every "
            "variable in every dataset it is handed, so buffers hold full rows."
        )
    for variable in self.variables:
        buffer[variable].append(float(values[variable]))
    self._steps += 1

refit()

Recompute the I-MEC from the buffers and fit one SCM per member; return the belief.

Two roles for the same rows, and they are not interchangeable. Orientation comes from the do-data: :func:~causalrl.discovery.discover_interventional runs PC on the observational sample for the skeleton and then orients the edges incident to each target by the invariance principle, which only the interventional marginals can decide. Mechanisms are fitted from every row the agent holds, observational and interventional concatenated, because a mechanism is invariant across regimes (Peters et al. 2016) and the do-rows are extra samples of every node that was not intervened on. The intervened node's own mechanism is fitted from rows where it was overridden, which biases its marginal -- harmless here, since planning replaces that mechanism with do(treatment=a) anyway.

A target holding fewer than min_interventional_samples rows is dropped rather than passed on. The invariance test compares empirical marginals, so a small do-sample clears any shift threshold by chance and an empty one clears it unconditionally — orienting every edge incident to that target from evidence the agent does not have. :func:~causalrl.discovery.discover_interventional refuses such a sample outright; the agent filters first so that an early round, when a target has been probed once or twice, is an ordinary under-informed refit rather than a raise.

Raises ValueError if no observational rows have been buffered (PC needs them for the base CPDAG), and lets :func:~causalrl.scm.fit.fit_scm_mec's over-size refusal propagate.

Source code in src/causalrl/agents/online_causal_mbrl.py

def refit(self) -> tuple[StructuralCausalModel, ...]:
    """Recompute the I-MEC from the buffers and fit one SCM per member; return the belief.

    Two roles for the same rows, and they are not interchangeable. **Orientation** comes from
    the do-data: :func:`~causalrl.discovery.discover_interventional` runs PC on the
    observational sample for the skeleton and then orients the edges incident to each target by
    the invariance principle, which only the interventional marginals can decide.
    **Mechanisms** are fitted from every row the agent holds, observational and interventional
    concatenated, because a mechanism is invariant across regimes (Peters et al. 2016) and the
    do-rows are extra samples of every node that was not intervened on. The intervened node's
    own mechanism is fitted from rows where it was overridden, which biases *its* marginal --
    harmless here, since planning replaces that mechanism with ``do(treatment=a)`` anyway.

    A target holding fewer than ``min_interventional_samples`` rows is dropped rather than
    passed on. The invariance test compares empirical marginals, so a small do-sample clears
    any shift threshold by chance and an empty one clears it unconditionally — orienting every
    edge incident to that target from evidence the agent does not have.
    :func:`~causalrl.discovery.discover_interventional` refuses such a sample outright; the
    agent filters first so that an early round, when a target has been probed once or twice, is
    an ordinary under-informed refit rather than a raise.

    Raises ``ValueError`` if no observational rows have been buffered (PC needs them for the
    base CPDAG), and lets :func:`~causalrl.scm.fit.fit_scm_mec`'s over-size refusal propagate.
    """
    if not self._rows(self._observational):
        raise ValueError(
            "refit() needs observational rows: discover_interventional runs PC on the "
            "observational sample to build the base CPDAG before the invariance principle "
            "orients edges incident to each intervention target. Ingest an observational log "
            "(or call observe() without intervention=...) first."
        )
    observational = _columns(self._observational)
    interventions = {
        target: _columns(buffer)
        for target, buffer in sorted(self._interventional.items())
        if self._rows(buffer) >= self.min_interventional_samples
    }
    cpdag = discover_interventional(
        observational,
        interventions,
        self.variables,
        min_interventional_samples=self.min_interventional_samples,
    )
    pooled = {
        variable: np.concatenate(
            [observational[variable], *(data[variable] for data in interventions.values())]
        )
        for variable in self.variables
    }
    self._belief = tuple(
        fit_scm_mec(pooled, cpdag=cpdag, max_members=self.max_members, seed=self.seed)
    )
    self._cpdag = cpdag
    self._history.append(RefitRecord(step=self._steps, belief_size=len(self._belief)))
    return self._belief

act()

Choose an action from actions by planning inside the belief.

Every policy scores an action the same way -- member.do({treatment: action}), then see(n_rollout) and the mean of outcome -- and they differ only in how the members' verdicts are combined:

• "thompson" draws one member and returns its argmax. Thompson sampling over structure (Ortega & Braun, arXiv:1303.4431 <https://arxiv.org/abs/1303.4431>_): the spread of the belief supplies the exploration, so the policy sharpens exactly as the I-MEC collapses and needs no separate schedule.
• "average" takes the mean value across members -- marginalising the structure out.
• "robust" takes the maximin: score each action by its worst member, then take the best of those. The one to reach for when a wrong action is expensive, since it is the only one of the three that never bets on a particular member being the true DAG.

Deterministic given seed. The thompson draw comes from an RNG derived from seed, the transition count and the number of draws already taken, so a fixed seed replays the whole decision sequence, successive calls draw successive members rather than repeating one, and two seeds diverge. Rollouts are seeded per member, which also hands every action of a member the same exogenous draws (common random numbers), so the gap between two actions is not partly the gap between two samples. Ties go to the earliest action in actions.

This costs rollouts: n_rollout samples per action per member scored, so len(actions) * n_rollout under "thompson" and len(actions) * belief_size() * n_rollout under "average" and "robust", every call. That is what planning in a model costs instead of reading a value table.

Raises ValueError before the first :meth:refit, when there is no belief to plan in.

Source code in src/causalrl/agents/online_causal_mbrl.py

def act(self) -> float:
    """Choose an action from ``actions`` by planning inside the belief.

    Every policy scores an action the same way -- ``member.do({treatment: action})``, then
    ``see(n_rollout)`` and the mean of ``outcome`` -- and they differ only in how the members'
    verdicts are combined:

    - ``"thompson"`` draws one member and returns its argmax. Thompson sampling over
      structure (Ortega & Braun, `arXiv:1303.4431 <https://arxiv.org/abs/1303.4431>`_): the
      spread of the belief supplies the exploration, so the policy sharpens exactly as the
      I-MEC collapses and needs no separate schedule.
    - ``"average"`` takes the mean value across members -- marginalising the structure out.
    - ``"robust"`` takes the maximin: score each action by its *worst* member, then take the
      best of those. The one to reach for when a wrong action is expensive, since it is the
      only one of the three that never bets on a particular member being the true DAG.

    Deterministic given ``seed``. The thompson draw comes from an RNG derived from ``seed``,
    the transition count and the number of draws already taken, so a fixed seed replays the
    whole decision sequence, successive calls draw successive members rather than repeating
    one, and two seeds diverge. Rollouts are seeded per member, which also hands every action
    of a member the same exogenous draws (common random numbers), so the gap between two
    actions is not partly the gap between two samples. Ties go to the earliest action in
    ``actions``.

    **This costs rollouts:** ``n_rollout`` samples per action per member scored, so
    ``len(actions) * n_rollout`` under ``"thompson"`` and
    ``len(actions) * belief_size() * n_rollout`` under ``"average"`` and ``"robust"``, every
    call. That is what planning in a model costs instead of reading a value table.

    Raises ``ValueError`` before the first :meth:`refit`, when there is no belief to plan in.
    """
    belief = self._require_belief("act")
    if self.policy == "thompson":
        rng = np.random.default_rng((self.seed, self._steps, self._draws))
        self._draws += 1
        drawn = int(rng.integers(len(belief)))
        values = [
            self._action_value(belief[drawn], action, seed=self.seed + drawn)
            for action in self.actions
        ]
    else:
        by_member = [
            [
                self._action_value(member, action, seed=self.seed + index)
                for action in self.actions
            ]
            for index, member in enumerate(belief)
        ]
        across = list(zip(*by_member, strict=True))
        values = (
            [float(np.mean(column)) for column in across]
            if self.policy == "average"
            else [min(column) for column in across]
        )
    return self.actions[max(range(len(self.actions)), key=lambda index: values[index])]

probe()

Return the intervention target the belief's members disagree about most.

A target, not an action: the answer to "what experiment should I run next", which is the other half of the loop :meth:act serves. Each candidate is scored by the mean pairwise total variation between the members' outcome distributions under do(candidate) -- predictive disagreement standing in for expected information gain, as in Scherrer et al. (arXiv:2109.02429 <https://arxiv.org/abs/2109.02429>_) and Zhang et al. (Nature Machine Intelligence 2023). It concentrates experiments on the edges the I-MEC still leaves unoriented, because those are the only ones whose do() distribution the members can disagree about; a target every member already predicts identically scores zero and is excluded, and if every candidate scores zero this raises rather than returning an arbitrary one.

Each candidate is intervened at one representative value: the largest it takes anywhere in the agent's buffers. That keeps the do-value inside the observed support, so no fitted mechanism is asked to extrapolate, while sitting as far as the data allows from the bottom of the range -- a do-value near the middle of a variable's spread perturbs its descendants least and makes members that disagree look alike. One value cannot certify agreement, though: a zero score means "no disagreement detected at this value", not "these members agree everywhere".

This costs rollouts: n_rollout samples per candidate per member, i.e. len(variables) * belief_size() * n_rollout every call. All members share one exogenous stream per candidate (common random numbers), because the score is a difference between members and independent draws would put a finite-sample floor under it.

Raises ValueError before the first :meth:refit, on a single-member belief (nothing to disambiguate -- :meth:structure_uncertain is the guard to branch on), and when no candidate scores any disagreement at all.

Source code in src/causalrl/agents/online_causal_mbrl.py

def probe(self) -> str:
    """Return the intervention target the belief's members disagree about most.

    A **target**, not an action: the answer to "what experiment should I run next", which is
    the other half of the loop :meth:`act` serves. Each candidate is scored by the mean
    pairwise total variation between the members' outcome distributions under ``do(candidate)``
    -- predictive disagreement standing in for expected information gain, as in Scherrer et
    al. (`arXiv:2109.02429 <https://arxiv.org/abs/2109.02429>`_) and Zhang et al. (Nature
    Machine Intelligence 2023). It concentrates experiments on the edges the I-MEC still leaves
    unoriented, because those are the only ones whose ``do()`` distribution the members can
    disagree about; a target every member already predicts identically scores zero and is
    excluded, and if *every* candidate scores zero this raises rather than returning an
    arbitrary one.

    Each candidate is intervened at one representative value: the largest it takes anywhere in
    the agent's buffers. That keeps the do-value inside the observed support, so no fitted
    mechanism is asked to extrapolate, while sitting as far as the data allows from the bottom
    of the range -- a do-value near the middle of a variable's spread perturbs its descendants
    least and makes members that disagree look alike. One value cannot certify agreement,
    though: a zero score means "no disagreement detected at this value", not "these members
    agree everywhere".

    **This costs rollouts:** ``n_rollout`` samples per candidate per member, i.e.
    ``len(variables) * belief_size() * n_rollout`` every call. All members share one exogenous
    stream per candidate (common random numbers), because the score is a difference *between*
    members and independent draws would put a finite-sample floor under it.

    Raises ``ValueError`` before the first :meth:`refit`, on a single-member belief (nothing to
    disambiguate -- :meth:`structure_uncertain` is the guard to branch on), and when no
    candidate scores any disagreement at all.
    """
    belief = self._require_belief("probe")
    if len(belief) == 1:
        raise ValueError(
            "probe() has nothing to disambiguate: the belief holds a single I-MEC member, so "
            "every candidate target has exactly one implied do() distribution and no "
            "experiment can separate anything. structure_uncertain() is the guard to branch "
            "on -- probe() while it reports True, act() once it reports False."
        )
    scores = {
        candidate: self._disagreement(candidate, belief) for candidate in sorted(self.variables)
    }
    best = max(scores, key=lambda candidate: scores[candidate])
    if scores[best] <= 0.0:
        raise ValueError(
            f"probe() found no target the belief's {len(belief)} members disagree about: every "
            f"candidate in {sorted(self.variables)} scored zero total variation between the "
            "members' outcome distributions. Intervening on any of them would be an experiment "
            "whose result every member already predicts identically, so returning one would "
            "report an arbitrary pick as an informative choice. The members differ somewhere "
            "else -- in a mechanism, or at a do-value other than the one probe() uses -- and "
            "that is not something this score can locate."
        )
    return best

belief()

The fitted I-MEC members -- empty until the first :meth:refit.

Source code in src/causalrl/agents/online_causal_mbrl.py

def belief(self) -> tuple[StructuralCausalModel, ...]:
    """The fitted I-MEC members -- empty until the first :meth:`refit`."""
    return self._belief

belief_size()

How many DAGs the data still cannot choose between.

Source code in src/causalrl/agents/online_causal_mbrl.py

def belief_size(self) -> int:
    """How many DAGs the data still cannot choose between."""
    return len(self._belief)

structure_uncertain()

Whether more than one I-MEC member survives -- a decision may need the certificate layer.

Raises ValueError before the first :meth:refit rather than answering. With an empty belief neither answer is honest: there is nothing to be uncertain between yet, and nothing that says the structure is known either. This is the guard the documented loop branches on -- :meth:probe while it is True, :meth:act once it is False -- so reporting False from an empty belief would quietly route a caller who never fitted anything into the exploit branch.

Source code in src/causalrl/agents/online_causal_mbrl.py

def structure_uncertain(self) -> bool:
    """Whether more than one I-MEC member survives -- a decision may need the certificate layer.

    Raises ``ValueError`` before the first :meth:`refit` rather than answering. With an empty
    belief neither answer is honest: there is nothing to be uncertain *between* yet, and
    nothing that says the structure is known either. This is the guard the documented loop
    branches on -- :meth:`probe` while it is ``True``, :meth:`act` once it is ``False`` -- so
    reporting ``False`` from an empty belief would quietly route a caller who never fitted
    anything into the exploit branch.
    """
    return len(self._require_belief("structure_uncertain")) > 1

history()

An immutable snapshot of the belief trajectory, one :class:RefitRecord per refit.

Source code in src/causalrl/agents/online_causal_mbrl.py

def history(self) -> tuple[RefitRecord, ...]:
    """An immutable snapshot of the belief trajectory, one :class:`RefitRecord` per refit."""
    return tuple(self._history)

causalrl.agents.baselines.OnlineOnlyUCB

Bases: Agent

UCB1 per state, ignoring any offline data. The 'learn from scratch online' baseline.

Source code in src/causalrl/agents/baselines.py

class OnlineOnlyUCB(Agent):
    """UCB1 per state, ignoring any offline data. The 'learn from scratch online' baseline."""

    def __init__(self, n_states: int, n_actions: int, seed: int | None = None) -> None:
        self.n_states = n_states
        self.n_actions = n_actions
        self._counts = np.zeros((n_states, n_actions))
        self._sums = np.zeros((n_states, n_actions))
        self._t = 1
        self._rng = np.random.default_rng(seed)

    def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
        """No-op: this baseline ignores offline data by design."""

    def act(self, observation: dict[str, Any]) -> int:
        s = int(observation["state"])
        untried = [a for a in range(self.n_actions) if self._counts[s, a] == 0]
        if untried:
            return int(self._rng.choice(untried))
        means = self._sums[s] / self._counts[s]
        bonus = np.sqrt(2.0 * math.log(self._t) / self._counts[s])
        return int(np.argmax(np.asarray(means + bonus, dtype=np.float64)))

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        s = int(observation["state"])
        self._counts[s, action] += 1.0
        self._sums[s, action] += reward
        self._t += 1

ingest_offline(dataset)

No-op: this baseline ignores offline data by design.

Source code in src/causalrl/agents/baselines.py

def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
    """No-op: this baseline ignores offline data by design."""

causalrl.envs.suite.seq_dtr.SequentialDTREnv

Bases: ConfoundedMDP

Multi-stage confounded dynamic treatment regime — genuinely confounded AND sequential.

Per episode a persistent hidden comorbidity U ~ Bernoulli(0.5) is drawn. At each of H stages the agent observes a subtype Z (the state, encoded stage*2 + Z so the stage is baked into the index; terminal = 2H), picks a treatment a in {0,1} and receives reward ~ Bernoulli(p(Z, a, U)). The subtype then transitions deterministically by Z' = a — a function of the chosen action and observed state only, NOT of U — so the transition dynamics are unconfounded and DOVI's value-iteration bound stays certified.

The reward table is tuned so the immediately-greedy action differs from the lookahead- optimal one (a genuine foresight gap), and U adds +0.2 to every cell when U = 1 while the logging clinician plays a = U (prob 0.85): so naive-offline over-prescribes treatment 1 in the logs AND, being myopic, also misorders the sequential trade-off.

Source code in src/causalrl/envs/suite/seq_dtr.py

class SequentialDTREnv(ConfoundedMDP):
    """Multi-stage confounded dynamic treatment regime — genuinely confounded AND sequential.

    Per episode a persistent hidden comorbidity ``U ~ Bernoulli(0.5)`` is drawn. At each of
    ``H`` stages the agent observes a subtype ``Z`` (the state, encoded ``stage*2 + Z`` so the
    stage is baked into the index; terminal = ``2H``), picks a treatment ``a in {0,1}`` and
    receives ``reward ~ Bernoulli(p(Z, a, U))``. The subtype then transitions deterministically
    by ``Z' = a`` — a function of the chosen action and observed state only, NOT of ``U`` — so
    the transition dynamics are unconfounded and DOVI's value-iteration bound stays certified.

    The reward table is tuned so the immediately-greedy action differs from the lookahead-
    optimal one (a genuine foresight gap), and ``U`` adds +0.2 to every cell when ``U = 1``
    while the logging clinician plays ``a = U`` (prob 0.85): so naive-offline over-prescribes
    treatment 1 in the logs AND, being myopic, also misorders the sequential trade-off.
    """

    n_actions = 2

    # P(R = 1 | Z, a, U). dv(Z,a) = mean over U: dv(0,0)=.5 dv(0,1)=.65 dv(1,0)=.3 dv(1,1)=.35.
    _REWARD_PROB: ClassVar[dict[tuple[int, int, int], float]] = {
        (0, 0, 0): 0.30,
        (0, 0, 1): 0.70,
        (0, 1, 0): 0.45,
        (0, 1, 1): 0.85,
        (1, 0, 0): 0.10,
        (1, 0, 1): 0.50,
        (1, 1, 0): 0.15,
        (1, 1, 1): 0.55,
    }
    _BEHAVIOR_FOLLOW_U = 0.85

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]  # gymnasium Env.metadata is an instance var in the base

    def __init__(self, horizon: int = 2, seed: int | None = None) -> None:
        if horizon < 1:
            raise ValueError(f"horizon must be >= 1, got {horizon}")
        super().__init__()
        self.horizon = horizon
        self.n_states = horizon * 2 + 1
        self.action_space = gym.spaces.Discrete(2)
        self.observation_space = gym.spaces.Dict(
            {"state": gym.spaces.Discrete(self.n_states), "t": gym.spaces.Discrete(horizon + 1)}
        )
        self._rng = np.random.default_rng(seed)
        self._terminal = horizon * 2
        self._t = 0
        self._z = 0
        self._u = 0

    def dv(self, z: int, a: int) -> float:
        """Per-stage interventional value E[R | do(a), Z=z], averaged over U ~ Bernoulli(0.5)."""
        return 0.5 * self._REWARD_PROB[(z, a, 0)] + 0.5 * self._REWARD_PROB[(z, a, 1)]

    def do_value(self, z: int, a: int, stage: int) -> float:
        """Interventional value of taking ``a`` at ``stage`` in subtype ``z``, then optimal."""
        if stage >= self.horizon - 1:
            return self.dv(z, a)
        z_next = a  # transition Z' = a
        v_next = max(self.do_value(z_next, a2, stage + 1) for a2 in (0, 1))
        return self.dv(z, a) + v_next

    @property
    def optimal_value(self) -> float:
        """Expected return of the do-optimal policy from a random initial subtype Z."""
        return 0.5 * sum(max(self.do_value(z, a, 0) for a in (0, 1)) for z in (0, 1))

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, int], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        self._t = 0
        self._u = int(self._rng.integers(0, 2))
        self._z = int(self._rng.integers(0, 2))
        return {"state": self._z, "t": 0}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, int], float, bool, bool, dict[str, Any]]:
        p = self._REWARD_PROB[(self._z, action, self._u)]
        reward = float(self._rng.random() < p)
        self._t += 1
        if self._t >= self.horizon:
            return {"state": self._terminal, "t": self._t}, reward, True, False, {"u": self._u}
        self._z = action  # transition Z' = a (U-independent)
        state = self._t * 2 + self._z
        return {"state": state, "t": self._t}, reward, False, False, {}

    def behavior_policy(self, observation: dict[str, int]) -> int:
        """Confounded logging: the clinician follows the hidden comorbidity U (plays a = U
        with prob 0.85), tying the logged action to U, which also drives the reward."""
        if self._rng.random() < self._BEHAVIOR_FOLLOW_U:
            return self._u
        return 1 - self._u

optimal_value property

Expected return of the do-optimal policy from a random initial subtype Z.

dv(z, a)

Per-stage interventional value E[R | do(a), Z=z], averaged over U ~ Bernoulli(0.5).

Source code in src/causalrl/envs/suite/seq_dtr.py

def dv(self, z: int, a: int) -> float:
    """Per-stage interventional value E[R | do(a), Z=z], averaged over U ~ Bernoulli(0.5)."""
    return 0.5 * self._REWARD_PROB[(z, a, 0)] + 0.5 * self._REWARD_PROB[(z, a, 1)]

do_value(z, a, stage)

Interventional value of taking a at stage in subtype z, then optimal.

Source code in src/causalrl/envs/suite/seq_dtr.py

def do_value(self, z: int, a: int, stage: int) -> float:
    """Interventional value of taking ``a`` at ``stage`` in subtype ``z``, then optimal."""
    if stage >= self.horizon - 1:
        return self.dv(z, a)
    z_next = a  # transition Z' = a
    v_next = max(self.do_value(z_next, a2, stage + 1) for a2 in (0, 1))
    return self.dv(z, a) + v_next

behavior_policy(observation)

Confounded logging: the clinician follows the hidden comorbidity U (plays a = U with prob 0.85), tying the logged action to U, which also drives the reward.

Source code in src/causalrl/envs/suite/seq_dtr.py

def behavior_policy(self, observation: dict[str, int]) -> int:
    """Confounded logging: the clinician follows the hidden comorbidity U (plays a = U
    with prob 0.85), tying the logged action to U, which also drives the reward."""
    if self._rng.random() < self._BEHAVIOR_FOLLOW_U:
        return self._u
    return 1 - self._u

causalrl.envs.suite.seq_mabuc.SequentialMABUCEnv

Bases: ConfoundedMDP

Horizon-H contextual bandit with an unobserved confounder (a genuine MABUC).

Each step is an independent confounded contextual-bandit decision: an observed context C ~ Bernoulli(0.5) (the state, stage*2 + C; terminal = 2H) and a HIDDEN U ~ Bernoulli(0.5) are drawn, then reward ~ Bernoulli(p(C, a, U)). The logging policy follows the hidden U (plays a = U with prob 0.85), so conditional on the observed context the logged action is tied to U — which also drives the reward. Naive-offline therefore over-prescribes treatment 1 in every context (its apparent reward is inflated by U), exactly as in DTREnv, but here repeated across the horizon.

There are no state transitions: this is the bandit member of the suite. The sequential-credit-assignment lesson lives in SequentialDTREnv.

Source code in src/causalrl/envs/suite/seq_mabuc.py

class SequentialMABUCEnv(ConfoundedMDP):
    """Horizon-H contextual bandit with an unobserved confounder (a genuine MABUC).

    Each step is an independent confounded contextual-bandit decision: an observed context
    ``C ~ Bernoulli(0.5)`` (the state, ``stage*2 + C``; terminal = ``2H``) and a HIDDEN
    ``U ~ Bernoulli(0.5)`` are drawn, then ``reward ~ Bernoulli(p(C, a, U))``. The logging
    policy follows the hidden ``U`` (plays ``a = U`` with prob 0.85), so conditional on the
    observed context the logged action is tied to ``U`` — which also drives the reward.
    Naive-offline therefore over-prescribes treatment 1 in every context (its apparent reward
    is inflated by ``U``), exactly as in ``DTREnv``, but here repeated across the horizon.

    There are no state transitions: this is the bandit member of the suite. The
    sequential-credit-assignment lesson lives in ``SequentialDTREnv``.
    """

    n_actions = 2

    # P(R = 1 | C, a, U). do-optimal a = C (value 0.75/step); U inflates treatment 1.
    _REWARD_PROB: ClassVar[dict[tuple[int, int, int], float]] = {
        (0, 0, 0): 0.70,
        (0, 0, 1): 0.70,
        (0, 1, 0): 0.20,
        (0, 1, 1): 0.90,
        (1, 0, 0): 0.20,
        (1, 0, 1): 0.20,
        (1, 1, 0): 0.70,
        (1, 1, 1): 0.90,
    }
    _BEHAVIOR_FOLLOW_U = 0.85

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]  # gymnasium Env.metadata is an instance var in the base

    def __init__(self, horizon: int = 3, seed: int | None = None) -> None:
        if horizon < 1:
            raise ValueError(f"horizon must be >= 1, got {horizon}")
        super().__init__()
        self.horizon = horizon
        self.n_states = horizon * 2 + 1
        self.action_space = gym.spaces.Discrete(2)
        self.observation_space = gym.spaces.Dict(
            {"state": gym.spaces.Discrete(self.n_states), "t": gym.spaces.Discrete(horizon + 1)}
        )
        self._rng = np.random.default_rng(seed)
        self._terminal = horizon * 2
        self._t = 0
        self._c = 0
        self._u = 0

    def do_value(self, c: int, a: int) -> float:
        """Per-step interventional value E[R | do(a), C=c], averaged over U ~ Bernoulli(0.5)."""
        return 0.5 * self._REWARD_PROB[(c, a, 0)] + 0.5 * self._REWARD_PROB[(c, a, 1)]

    @property
    def optimal_value(self) -> float:
        """Per-step value of the do-optimal policy a*(C) = argmax_a do_value(C, a)."""
        return 0.5 * sum(max(self.do_value(c, 0), self.do_value(c, 1)) for c in (0, 1))

    def _draw(self) -> None:
        self._c = int(self._rng.integers(0, 2))
        self._u = int(self._rng.integers(0, 2))

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, int], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        self._t = 0
        self._draw()
        return {"state": self._c, "t": 0}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, int], float, bool, bool, dict[str, Any]]:
        p = self._REWARD_PROB[(self._c, action, self._u)]
        reward = float(self._rng.random() < p)
        self._t += 1
        if self._t >= self.horizon:
            return {"state": self._terminal, "t": self._t}, reward, True, False, {}
        self._draw()
        state = self._t * 2 + self._c
        return {"state": state, "t": self._t}, reward, False, False, {}

    def behavior_policy(self, observation: dict[str, int]) -> int:
        """Confounded logging: follows the hidden U (plays a = U with prob 0.85)."""
        if self._rng.random() < self._BEHAVIOR_FOLLOW_U:
            return self._u
        return 1 - self._u

optimal_value property

Per-step value of the do-optimal policy a*(C) = argmax_a do_value(C, a).

do_value(c, a)

Per-step interventional value E[R | do(a), C=c], averaged over U ~ Bernoulli(0.5).

Source code in src/causalrl/envs/suite/seq_mabuc.py

def do_value(self, c: int, a: int) -> float:
    """Per-step interventional value E[R | do(a), C=c], averaged over U ~ Bernoulli(0.5)."""
    return 0.5 * self._REWARD_PROB[(c, a, 0)] + 0.5 * self._REWARD_PROB[(c, a, 1)]

behavior_policy(observation)

Confounded logging: follows the hidden U (plays a = U with prob 0.85).

Source code in src/causalrl/envs/suite/seq_mabuc.py

def behavior_policy(self, observation: dict[str, int]) -> int:
    """Confounded logging: follows the hidden U (plays a = U with prob 0.85)."""
    if self._rng.random() < self._BEHAVIOR_FOLLOW_U:
        return self._u
    return 1 - self._u

causalrl.envs.suite.simpson_bandit.SimpsonBandit

Single-decision bandit with an observed confounder Z on the back-door path A <- Z -> Y.

Source code in src/causalrl/envs/suite/simpson_bandit.py

class SimpsonBandit:
    """Single-decision bandit with an observed confounder Z on the back-door path A <- Z -> Y."""

    n_states: int = 1
    n_actions: int = 2

    def __init__(self, *, seed: int | None = None) -> None:
        self.graph = CausalGraph(directed_edges=[("Z", "A"), ("Z", "Y"), ("A", "Y")])
        self._rng = np.random.default_rng(seed)

    def sample(self, n: int, *, seed: int | None = None) -> dict[str, np.ndarray]:
        """Draw ``n`` confounded observations as columnar arrays ``{"Z", "A", "Y"}``."""
        return self.sample_do({}, n, seed=seed)

    def sample_do(
        self, do: Mapping[str, int], n: int, *, seed: int | None = None
    ) -> dict[str, np.ndarray]:
        """Draw ``n`` samples under the mutilated model ``do`` (``{}`` = observational).

        Supports interventions on ``Z`` and/or ``A``; unset variables follow their mechanism, so
        interventional causal discovery can orient the graph by the invariance principle.
        """
        rng = np.random.default_rng(seed) if seed is not None else self._rng
        z = (rng.random(n) < 0.5).astype(int)
        if "Z" in do:
            z = np.full(n, do["Z"], dtype=int)
        a = (rng.random(n) < _P_TREAT[z]).astype(int)
        if "A" in do:
            a = np.full(n, do["A"], dtype=int)
        y = (rng.random(n) < _P_REWARD[a, z]).astype(float)
        return {"Z": z, "A": a, "Y": y}

    def true_action_value(self, action: int) -> float:
        """Exact interventional value ``E[Y | do(A = action)] = Σ_z P(z) · P(Y=1 | action, z)``."""
        return float(0.5 * _P_REWARD[action, 0] + 0.5 * _P_REWARD[action, 1])

    @property
    def optimal_value(self) -> float:
        return max(self.true_action_value(a) for a in range(self.n_actions))

sample(n, *, seed=None)

Draw n confounded observations as columnar arrays {"Z", "A", "Y"}.

Source code in src/causalrl/envs/suite/simpson_bandit.py

def sample(self, n: int, *, seed: int | None = None) -> dict[str, np.ndarray]:
    """Draw ``n`` confounded observations as columnar arrays ``{"Z", "A", "Y"}``."""
    return self.sample_do({}, n, seed=seed)

sample_do(do, n, *, seed=None)

Draw n samples under the mutilated model do ({} = observational).

Supports interventions on Z and/or A; unset variables follow their mechanism, so interventional causal discovery can orient the graph by the invariance principle.

Source code in src/causalrl/envs/suite/simpson_bandit.py

def sample_do(
    self, do: Mapping[str, int], n: int, *, seed: int | None = None
) -> dict[str, np.ndarray]:
    """Draw ``n`` samples under the mutilated model ``do`` (``{}`` = observational).

    Supports interventions on ``Z`` and/or ``A``; unset variables follow their mechanism, so
    interventional causal discovery can orient the graph by the invariance principle.
    """
    rng = np.random.default_rng(seed) if seed is not None else self._rng
    z = (rng.random(n) < 0.5).astype(int)
    if "Z" in do:
        z = np.full(n, do["Z"], dtype=int)
    a = (rng.random(n) < _P_TREAT[z]).astype(int)
    if "A" in do:
        a = np.full(n, do["A"], dtype=int)
    y = (rng.random(n) < _P_REWARD[a, z]).astype(float)
    return {"Z": z, "A": a, "Y": y}

true_action_value(action)

Exact interventional value E[Y | do(A = action)] = Σ_z P(z) · P(Y=1 | action, z).

Source code in src/causalrl/envs/suite/simpson_bandit.py

def true_action_value(self, action: int) -> float:
    """Exact interventional value ``E[Y | do(A = action)] = Σ_z P(z) · P(Y=1 | action, z)``."""
    return float(0.5 * _P_REWARD[action, 0] + 0.5 * _P_REWARD[action, 1])

causalrl.envs.suite.scbandit.StructuralCausalBanditEnv

Bases: CausalEnv

One-step bandit whose arms are interventions on an SCM. step returns a single stochastic reward draw under do(arm) (not its expected value); arm_values and optimal_value — Monte-Carlo estimates computed once at construction — are the ground truth for regret. graph is the abstracted ADMG given to agents; arms and manipulable support POMIS arm pruning.

Source code in src/causalrl/envs/suite/scbandit.py

class StructuralCausalBanditEnv(CausalEnv):
    """One-step bandit whose arms are interventions on an SCM. `step` returns a single
    stochastic reward draw under do(arm) (not its expected value); `arm_values` and
    `optimal_value` — Monte-Carlo estimates computed once at construction — are the ground
    truth for regret. `graph` is the abstracted ADMG given to agents; `arms` and
    `manipulable` support POMIS arm pruning."""

    metadata: ClassVar[dict[str, list[str]]] = {"render_modes": []}  # type: ignore[misc]

    def __init__(
        self,
        scm: StructuralCausalModel,
        admg: CausalGraph,
        reward: str,
        manipulable: list[str],
        domains: dict[str, list[int]],
        *,
        n_mc: int = 2000,
        seed: int | None = None,
    ) -> None:
        super().__init__(scm)
        self.graph = admg
        self.reward = reward
        self.manipulable = manipulable
        self.domains = domains
        self.arms = enumerate_arms(manipulable, domains)
        self.action_space = gym.spaces.Discrete(len(self.arms))
        self.observation_space = gym.spaces.Dict({"constant": gym.spaces.Discrete(1)})
        self._mc_seed = 12345 if seed is None else seed + 12345
        self._rng = np.random.default_rng(seed)
        self.arm_values = self._estimate_arm_values(n_mc)
        self.optimal_value = max(self.arm_values)

    def _arm_reward_mean(self, arm: dict[str, int], n: int, *, seed: int) -> float:
        model = self.scm.do({k: float(v) for k, v in arm.items()}) if arm else self.scm
        samples = model.see(n, seed=seed)[self.reward]
        return float(samples.mean().item())

    def _estimate_arm_values(self, n_mc: int) -> list[float]:
        rng = np.random.default_rng(self._mc_seed)
        return [
            self._arm_reward_mean(arm, n_mc, seed=int(rng.integers(0, 2**31))) for arm in self.arms
        ]

    def reset(  # type: ignore[override]
        self, *, seed: int | None = None, options: dict[str, Any] | None = None
    ) -> tuple[dict[str, Any], dict[str, Any]]:
        super().reset(seed=seed)
        if seed is not None:
            self._rng = np.random.default_rng(seed)
        return {"constant": 0}, {}

    def step(  # type: ignore[override]
        self, action: int
    ) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
        arm = self.arms[action]
        reward = self._arm_reward_mean(arm, 1, seed=int(self._rng.integers(0, 2**31)))
        return (
            {"constant": 0},
            reward,
            True,
            False,
            {"optimal_value": self.optimal_value, "arm": arm},
        )

causalrl.agents.mbrl.TransportBackdoorAgent

Bases: BatchAgent

Deconfound + transport: back-door-adjust for the confounders (source weighting) and reweight the selection variables by the target distribution, then ship the argmax of the transported interventional value E_target[Y | do(A=a)]. The M2 upgrade of :class:BackdoorAdjustedAgent that carries a policy across a covariate shift where a correlational agent cannot.

The back-door set comes from graph via :func:~causalrl.backdoor_adjustment_set (minus the declared transport variables); the estimand's transportability is confirmed up front with :func:~causalrl.is_transportable_effect (dogfooding the identification layer). Fitted on columnar source logs plus unlabeled target draws of the transport variables via :meth:fit.

:meth:act is a policy over adjustment + transport: the transported marginal can say "never play arm 1" while arm 1 still wins in a particular cell.

Source code in src/causalrl/agents/mbrl.py

class TransportBackdoorAgent(BatchAgent):
    """Deconfound + transport: back-door-adjust for the confounders (source weighting) and reweight
    the selection variables by the target distribution, then ship the argmax of the transported
    interventional value ``E_target[Y | do(A=a)]``. The M2 upgrade of :class:`BackdoorAdjustedAgent`
    that carries a policy across a covariate shift where a correlational agent cannot.

    The back-door set comes from ``graph`` via :func:`~causalrl.backdoor_adjustment_set` (minus the
    declared ``transport`` variables); the estimand's transportability is confirmed up front with
    :func:`~causalrl.is_transportable_effect` (dogfooding the identification layer). Fitted on
    columnar source logs plus unlabeled target draws of the transport variables via :meth:`fit`.

    :meth:`act` is a *policy* over ``adjustment + transport``: the transported marginal can say
    "never play arm 1" while arm 1 still wins in a particular cell.
    """

    def __init__(
        self,
        n_actions: int,
        *,
        graph: CausalGraph,
        treatment: str = "A",
        outcome: str = "Y",
        transport: Sequence[str] = ("W",),
    ) -> None:
        self.n_actions = n_actions
        self.treatment = treatment
        self.outcome = outcome
        self.transport = tuple(transport)
        adjust = set(backdoor_adjustment_set(graph, treatment, outcome)) - set(self.transport)
        self.adjustment: tuple[str, ...] = tuple(sorted(adjust))
        self.transportable = is_transportable_effect(
            graph, {treatment}, {outcome}, set(self.transport)
        )
        self._context: tuple[str, ...] = self.adjustment + self.transport
        self._best_action = 0
        self.values: list[float] = [0.0] * n_actions
        self._stratum_values: dict[tuple[Any, ...], tuple[float, ...]] = {}

    def fit(
        self,
        source: Mapping[str, np.ndarray],
        *,
        target_covariates: Mapping[str, np.ndarray],
    ) -> None:
        """Estimate each action's transported interventional value, cache the per-cell action
        values :meth:`act` conditions on, and record the marginal argmax.

        ``source`` is columnar ``{treatment, outcome, *adjustment, *transport}``;
        ``target_covariates`` supplies unlabeled target draws of the ``transport`` variables.
        """
        self._stratum_values = _stratum_action_values(
            source, self.treatment, self.outcome, self._context, self.n_actions
        )
        self.values = [
            _transport_value(
                a,
                source,
                self.treatment,
                self.outcome,
                self.adjustment,
                self.transport,
                target_covariates,
            )
            for a in range(self.n_actions)
        ]
        self._best_action = int(np.argmax(np.asarray(self.values)))

    def act(self, observation: dict[str, Any]) -> int:
        """Contextual policy ``argmax_a E_target[Y | do(A=a), Z=z, W=w]`` at the adjustment and
        transport values in ``observation`` (supplied by name, e.g. ``{"Z": 1, "W": 0}``).

        The target reweighting that :meth:`fit` applies is a statement about *P(W)*, not about the
        outcome mechanism: under the S-admissibility this agent checks up front,
        ``E_target[Y | do(a), z, w] = E_source[Y | a, z, w]``, so the per-unit decision carries
        across the shift with no reweighting at all — only the marginal needs one. With none of the
        conditioning variables supplied, or a cell never logged in the source, the transported
        marginal decision is returned; a partial context raises.
        """
        return _stratum_action(self._stratum_values, observation, self._context, self._best_action)

fit(source, *, target_covariates)

Estimate each action's transported interventional value, cache the per-cell action values :meth:act conditions on, and record the marginal argmax.

source is columnar {treatment, outcome, *adjustment, *transport}; target_covariates supplies unlabeled target draws of the transport variables.

Source code in src/causalrl/agents/mbrl.py

def fit(
    self,
    source: Mapping[str, np.ndarray],
    *,
    target_covariates: Mapping[str, np.ndarray],
) -> None:
    """Estimate each action's transported interventional value, cache the per-cell action
    values :meth:`act` conditions on, and record the marginal argmax.

    ``source`` is columnar ``{treatment, outcome, *adjustment, *transport}``;
    ``target_covariates`` supplies unlabeled target draws of the ``transport`` variables.
    """
    self._stratum_values = _stratum_action_values(
        source, self.treatment, self.outcome, self._context, self.n_actions
    )
    self.values = [
        _transport_value(
            a,
            source,
            self.treatment,
            self.outcome,
            self.adjustment,
            self.transport,
            target_covariates,
        )
        for a in range(self.n_actions)
    ]
    self._best_action = int(np.argmax(np.asarray(self.values)))

act(observation)

Contextual policy argmax_a E_target[Y | do(A=a), Z=z, W=w] at the adjustment and transport values in observation (supplied by name, e.g. {"Z": 1, "W": 0}).

The target reweighting that :meth:fit applies is a statement about P(W), not about the outcome mechanism: under the S-admissibility this agent checks up front, E_target[Y | do(a), z, w] = E_source[Y | a, z, w], so the per-unit decision carries across the shift with no reweighting at all — only the marginal needs one. With none of the conditioning variables supplied, or a cell never logged in the source, the transported marginal decision is returned; a partial context raises.

Source code in src/causalrl/agents/mbrl.py

def act(self, observation: dict[str, Any]) -> int:
    """Contextual policy ``argmax_a E_target[Y | do(A=a), Z=z, W=w]`` at the adjustment and
    transport values in ``observation`` (supplied by name, e.g. ``{"Z": 1, "W": 0}``).

    The target reweighting that :meth:`fit` applies is a statement about *P(W)*, not about the
    outcome mechanism: under the S-admissibility this agent checks up front,
    ``E_target[Y | do(a), z, w] = E_source[Y | a, z, w]``, so the per-unit decision carries
    across the shift with no reweighting at all — only the marginal needs one. With none of the
    conditioning variables supplied, or a cell never logged in the source, the transported
    marginal decision is returned; a partial context raises.
    """
    return _stratum_action(self._stratum_values, observation, self._context, self._best_action)

causalrl.envs.suite.transport_bandit.TransportableConfoundedBandit

One-decision bandit with an orthogonal confounder Z and shift variable W.

Graph Z -> A, Z -> Y, W -> Y, A -> Y. Z is the back-door confounder; W is the selection variable whose distribution differs between source and target.

Source code in src/causalrl/envs/suite/transport_bandit.py

class TransportableConfoundedBandit:
    """One-decision bandit with an orthogonal confounder ``Z`` and shift variable ``W``.

    Graph ``Z -> A, Z -> Y, W -> Y, A -> Y``. ``Z`` is the back-door confounder; ``W`` is the
    selection variable whose distribution differs between source and target.
    """

    n_states: int = 1
    n_actions: int = 2

    def __init__(self, *, gamma: float, shift: float, seed: int | None = None) -> None:
        if not 0.0 <= gamma <= 1.0:
            raise ValueError("gamma must be in [0, 1]")
        if not 0.0 <= shift <= 1.0:
            raise ValueError("shift must be in [0, 1]")
        self.gamma = float(gamma)
        self.shift = float(shift)
        self.w_source = 0.5 - self.shift / 2.0
        self.w_target = 0.5 + self.shift / 2.0
        self.graph = CausalGraph(directed_edges=[("Z", "A"), ("Z", "Y"), ("W", "Y"), ("A", "Y")])
        self._rng = np.random.default_rng(seed)

    def _reward_mean(self, a: np.ndarray, z: np.ndarray, w: np.ndarray) -> np.ndarray:
        """Vectorized ``E[Y | A=a, Z=z, W=w]``. Safe arm 0 is flat 0.5; arm 1 adds the Z/W terms."""
        arm1 = 0.5 + _DC * np.where(z == 1, 1.0, -1.0) + np.where(w == 1, -_DT, 0.0)
        return np.where(a == 1, arm1, 0.5)

    def sample(
        self, n: int, *, domain: str = "source", seed: int | None = None
    ) -> dict[str, np.ndarray]:
        """Draw ``n`` confounded observations ``{"Z", "W", "A", "Y"}`` from ``domain``.

        ``domain`` is ``"source"`` (log-collection) or ``"target"`` (deployment); they differ only
        in ``P(W=1)``.
        """
        if domain not in ("source", "target"):
            raise ValueError("domain must be 'source' or 'target'")
        rng = np.random.default_rng(seed) if seed is not None else self._rng
        p_w = self.w_source if domain == "source" else self.w_target
        z = (rng.random(n) < _P_Z1).astype(int)
        w = (rng.random(n) < p_w).astype(int)
        a = (rng.random(n) < 0.5 + _C * self.gamma * (2 * z - 1)).astype(int)
        y = (rng.random(n) < self._reward_mean(a, z, w)).astype(float)
        return {"Z": z, "W": w, "A": a, "Y": y}

    def true_action_value(self, action: int, *, domain: str = "target") -> float:
        """Exact ``E[Y | do(A=action)] = Σ_{z,w} P(z) P(w) E[Y | action, z, w]`` in ``domain``."""
        p_w = self.w_target if domain == "target" else self.w_source
        z = np.array([0, 0, 1, 1])
        w = np.array([0, 1, 0, 1])
        means = self._reward_mean(np.full(4, action), z, w)
        p_z = np.where(z == 1, _P_Z1, 1.0 - _P_Z1)
        p_wv = np.where(w == 1, p_w, 1.0 - p_w)
        return float((p_z * p_wv * means).sum())

    def optimal_action(self, *, domain: str = "target") -> int:
        return max(range(self.n_actions), key=lambda a: self.true_action_value(a, domain=domain))

    def optimal_value(self, *, domain: str = "target") -> float:
        return max(self.true_action_value(a, domain=domain) for a in range(self.n_actions))

sample(n, *, domain='source', seed=None)

Draw n confounded observations {"Z", "W", "A", "Y"} from domain.

domain is "source" (log-collection) or "target" (deployment); they differ only in P(W=1).

Source code in src/causalrl/envs/suite/transport_bandit.py

def sample(
    self, n: int, *, domain: str = "source", seed: int | None = None
) -> dict[str, np.ndarray]:
    """Draw ``n`` confounded observations ``{"Z", "W", "A", "Y"}`` from ``domain``.

    ``domain`` is ``"source"`` (log-collection) or ``"target"`` (deployment); they differ only
    in ``P(W=1)``.
    """
    if domain not in ("source", "target"):
        raise ValueError("domain must be 'source' or 'target'")
    rng = np.random.default_rng(seed) if seed is not None else self._rng
    p_w = self.w_source if domain == "source" else self.w_target
    z = (rng.random(n) < _P_Z1).astype(int)
    w = (rng.random(n) < p_w).astype(int)
    a = (rng.random(n) < 0.5 + _C * self.gamma * (2 * z - 1)).astype(int)
    y = (rng.random(n) < self._reward_mean(a, z, w)).astype(float)
    return {"Z": z, "W": w, "A": a, "Y": y}

true_action_value(action, *, domain='target')

Exact E[Y | do(A=action)] = Σ_{z,w} P(z) P(w) E[Y | action, z, w] in domain.

Source code in src/causalrl/envs/suite/transport_bandit.py

def true_action_value(self, action: int, *, domain: str = "target") -> float:
    """Exact ``E[Y | do(A=action)] = Σ_{z,w} P(z) P(w) E[Y | action, z, w]`` in ``domain``."""
    p_w = self.w_target if domain == "target" else self.w_source
    z = np.array([0, 0, 1, 1])
    w = np.array([0, 1, 0, 1])
    means = self._reward_mean(np.full(4, action), z, w)
    p_z = np.where(z == 1, _P_Z1, 1.0 - _P_Z1)
    p_wv = np.where(w == 1, p_w, 1.0 - p_w)
    return float((p_z * p_wv * means).sum())

causalrl.agents.offline_online.UCDTR

Bases: Agent

Causal offline-to-online learner (Zhang & Bareinboim 2019).

Offline: prune each state's action set to the non-dominated actions implied by the Manski causal bounds (a safety mechanism; with natural bounds it is a no-op). Online: run UCB1 over only the surviving actions. Unlike NaiveOffline, it never trusts the confounded offline point estimates, so it corrects to the true optimum online.

Source code in src/causalrl/agents/offline_online.py

class UCDTR(Agent):
    """Causal offline-to-online learner (Zhang & Bareinboim 2019).

    Offline: prune each state's action set to the non-dominated actions implied by the
    Manski causal bounds (a safety mechanism; with natural bounds it is a no-op). Online:
    run UCB1 over only the surviving actions. Unlike NaiveOffline, it never trusts the
    confounded offline point estimates, so it corrects to the true optimum online.
    """

    def __init__(
        self,
        n_states: int,
        n_actions: int,
        seed: int | None = None,
        *,
        require_identified: bool = False,
    ) -> None:
        self.n_states = n_states
        self.n_actions = n_actions
        self._counts = np.zeros((n_states, n_actions))
        self._sums = np.zeros((n_states, n_actions))
        self._t = 1
        self._rng = np.random.default_rng(seed)
        self._allowed: dict[int, list[int]] = {s: list(range(n_actions)) for s in range(n_states)}
        self._require_identified = require_identified

    def ingest_offline(self, dataset: ConfoundedTrajectoryDataset) -> None:
        if self._require_identified:
            from causalrl.ope.bounds import causal_q_bounds

            for s in range(self.n_states):
                for a in range(self.n_actions):
                    causal_q_bounds(dataset, s, a, require_identified=True)
        for s in range(self.n_states):
            self._allowed[s] = non_dominated_actions(dataset, s) or list(range(self.n_actions))

    def act(self, observation: dict[str, Any]) -> int:
        s = int(observation["state"])
        allowed = self._allowed[s]
        untried = [a for a in allowed if self._counts[s, a] == 0]
        if untried:
            return int(self._rng.choice(untried))
        means = self._sums[s] / np.maximum(self._counts[s], 1.0)
        bonus = np.sqrt(2.0 * math.log(self._t) / np.maximum(self._counts[s], 1.0))
        scores = np.asarray(means + bonus, dtype=np.float64)
        best = max(allowed, key=lambda a: float(scores[a]))
        return int(best)

    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
        s = int(observation["state"])
        self._counts[s, action] += 1.0
        self._sums[s, action] += reward
        self._t += 1

causalrl.magames.cce_bounds(game, functional, *, do=None, epsilon=0.0)

Min/max of functional's expectation over the epsilon-CCE polytope under do.

functional maps a joint profile ({agent: action}) to a real value; epsilon is a uniform or per-agent relaxation of the no-deviation constraints — pass a measured realized regret (:func:cce_regret) for the finite-time form.

Source code in src/causalrl/magames/cce.py

def cce_bounds(
    game: CausalGame,
    functional: Functional,
    *,
    do: Mapping[str, int] | None = None,
    epsilon: Epsilon = 0.0,
) -> Interval:
    """Min/max of ``functional``'s expectation over the ``epsilon``-CCE polytope under ``do``.

    ``functional`` maps a joint profile (``{agent: action}``) to a real value; ``epsilon`` is a
    uniform or per-agent relaxation of the no-deviation constraints — pass a *measured* realized
    regret (:func:`cce_regret`) for the finite-time form.
    """
    polytope = cce_polytope(game, do=do)
    interval, _ = _polytope_bounds(polytope, _functional_vector(polytope, functional), epsilon)
    return interval

causalrl.magames.cce_polytope(game, *, do=None)

Deviation-gain representation of the CCE polytope of game under intervention do.

do pins the given agents' actions: profiles are restricted to the consistent ones and only the free agents contribute no-deviation constraints (the intervened game).

Source code in src/causalrl/magames/cce.py

def cce_polytope(game: CausalGame, *, do: Mapping[str, int] | None = None) -> CCEPolytope:
    """Deviation-gain representation of the CCE polytope of ``game`` under intervention ``do``.

    ``do`` pins the given agents' actions: profiles are restricted to the consistent ones and only
    the free agents contribute no-deviation constraints (the intervened game).
    """
    do = dict(do or {})
    for name, action in do.items():
        if name not in game.agents:
            raise KeyError(f"unknown agent: {name!r}")
        if action not in game.actions[name]:
            raise ValueError(f"action {action!r} not available to agent {name!r}")
    profiles = tuple(
        p
        for p in product(*(game.actions[a] for a in game.agents))
        if all(p[game.agents.index(a)] == do[a] for a in do)
    )
    free = [a for a in game.agents if a not in do]
    rows: list[list[float]] = []
    labels: list[tuple[str, int]] = []
    for agent in free:
        idx = game.agents.index(agent)
        table = game.utilities[agent]
        for deviation in game.actions[agent]:
            rows.append([table[(*p[:idx], deviation, *p[idx + 1 :])] - table[p] for p in profiles])
            labels.append((agent, deviation))
    gains = np.asarray(rows, dtype=np.float64) if rows else np.zeros((0, len(profiles)))
    return CCEPolytope(tuple(game.agents), profiles, gains, tuple(labels))

causalrl.magames.cce_regret(game, weights, *, do=None)

Maximal deviation gain of a joint distribution over profiles: the measured realized regret.

weights is either a mapping from joint-action tuples (in game.agents order) to probabilities or a sequence aligned with cce_polytope(game, do=do).profiles. The result is the smallest epsilon for which weights is epsilon-CCE-feasible (0 for an exact CCE); it is what a no-regret population drives to 0 and what :func:certify_cce_do accepts as the measured epsilon.

Source code in src/causalrl/magames/cce.py

def cce_regret(
    game: CausalGame,
    weights: Mapping[tuple[int, ...], float] | Sequence[float] | FloatArray,
    *,
    do: Mapping[str, int] | None = None,
) -> float:
    """Maximal deviation gain of a joint distribution over profiles: the measured realized regret.

    ``weights`` is either a mapping from joint-action tuples (in ``game.agents`` order) to
    probabilities or a sequence aligned with ``cce_polytope(game, do=do).profiles``. The result is
    the smallest ``epsilon`` for which ``weights`` is ``epsilon``-CCE-feasible (0 for an exact CCE);
    it is what a no-regret population drives to 0 and what :func:`certify_cce_do` accepts as the
    measured ``epsilon``.
    """
    polytope = cce_polytope(game, do=do)
    if isinstance(weights, Mapping):
        mu = np.array([float(weights.get(p, 0.0)) for p in polytope.profiles])
    else:
        mu = np.asarray(weights, dtype=np.float64)
        if mu.shape != (len(polytope.profiles),):
            raise ValueError(
                f"weights must align with the {len(polytope.profiles)} polytope profiles"
            )
    if np.any(mu < 0) or abs(float(mu.sum()) - 1.0) > 1e-6:
        raise ValueError("weights must be nonnegative and sum to 1 over the polytope profiles")
    if not len(polytope.constraint_labels):
        return 0.0
    return float(np.max(polytope.deviation_gains @ mu))

causalrl.magames.certify_cce_do(game, functional, *, do=None, no_regret=True, epsilon=None, payoff_error=None, tol=1e-09, seed=0)

Certify what functional's time average can be for a learning population under do.

no_regret asserts the population's time-averaged play approaches the exact CCE set (the asymptotic route); epsilon instead supplies a measured realized regret at the horizon actually run (the finite-time route, which needs no assumption). See the module docstring for the resulting IDENTIFIED/BOUNDED/EMPIRICAL ladder.

payoff_error says the game's payoff table was measured rather than given — from a round-robin of simulated matches, say (:class:~causalrl.magames.empirical.EmpiricalGame produces one from its own standard errors). It never narrows a claim's arithmetic and it never widens value, which stays the partial-identification region of the table as handed over; it fills ci with the interval that also survives the table's error at level alpha, judges vacuity on that, and withholds IDENTIFIED — a functional constant over an estimated polytope is constant in the estimate, which is not the same claim.

Source code in src/causalrl/magames/cce.py

def certify_cce_do(
    game: CausalGame,
    functional: Functional,
    *,
    do: Mapping[str, int] | None = None,
    no_regret: bool = True,
    epsilon: Epsilon | None = None,
    payoff_error: PayoffError | None = None,
    tol: float = 1e-9,
    seed: int = 0,
) -> Certificate:
    """Certify what ``functional``'s time average can be for a learning population under ``do``.

    ``no_regret`` asserts the population's time-averaged play approaches the exact CCE set (the
    asymptotic route); ``epsilon`` instead supplies a *measured* realized regret at the horizon
    actually run (the finite-time route, which needs no assumption). See the module docstring for
    the resulting ``IDENTIFIED``/``BOUNDED``/``EMPIRICAL`` ladder.

    ``payoff_error`` says the game's payoff table was *measured* rather than given — from a
    round-robin of simulated matches, say (:class:`~causalrl.magames.empirical.EmpiricalGame`
    produces one from its own standard errors). It never narrows a claim's arithmetic and it never
    widens ``value``, which stays the partial-identification region of the table as handed over;
    it fills ``ci`` with the interval that also survives the table's error at level ``alpha``,
    judges vacuity on *that*, and withholds ``IDENTIFIED`` — a functional constant over an
    estimated polytope is constant in the estimate, which is not the same claim.
    """
    polytope = cce_polytope(game, do=do)
    values = _functional_vector(polytope, functional)
    measured_epsilon = epsilon if epsilon is not None else 0.0
    interval, sensitivity = _polytope_bounds(polytope, values, measured_epsilon)
    error = payoff_error if payoff_error is not None else PayoffError(0.0)
    outer = interval if error.exact else _outer_interval(polytope, values, measured_epsilon, error)
    width = interval.upper - interval.lower
    reported_width = outer.upper - outer.lower
    full = Interval(float(np.min(values)), float(np.max(values)))
    full_width = full.upper - full.lower
    vacuous = full_width > tol and reported_width >= full_width - tol
    licensed = no_regret or epsilon is not None
    degenerate = width <= tol

    measured = None if epsilon is None else epsilon
    detail: dict[str, object] = {
        "interval": [interval.lower, interval.upper],
        "width": width,
        "epsilon": measured if not isinstance(measured, Mapping) else dict(measured),
        "epsilon_sensitivity": sensitivity,
        "n_profiles": len(polytope.profiles),
        "n_constraints": len(polytope.constraint_labels),
        "do": dict(do or {}),
    }
    if not error.exact:
        detail["payoff_error"] = {
            "utility": error.utility,
            "functional": error.functional,
            "alpha": error.alpha,
        }
        detail["ci"] = [outer.lower, outer.upper]
    witness = Witness("cce-interval", detail)
    assumptions = (
        Assumption(name="finite-game", params={"agents": list(game.agents)}, checkable=True),
        Assumption(
            name="no-regret",
            params={
                "assumed": no_regret,
                "measured_epsilon": None if isinstance(measured, Mapping) else measured,
            },
            checkable=True,
            diagnostic={"check": "cce_regret(game, realized_joint_distribution, do=do)"},
        ),
    )
    if not error.exact:
        assumptions = (
            *assumptions,
            Assumption(
                name="payoff-estimate",
                params={
                    "utility": error.utility,
                    "functional": error.functional,
                    "alpha": error.alpha,
                },
                checkable=True,
                diagnostic={
                    "check": "re-measure the payoff table and confirm every cell moves by less "
                    "than `utility`"
                },
            ),
        )
    reported_alpha = None if error.exact else error.alpha
    reported_ci = None if error.exact else outer
    provenance = Provenance.create(seeds=(seed,), graph_hash=graph_hash(game.graph))
    estimand = EstimandSpec(query="equilibrium", target="mean")
    intervened = f" | do={dict(do)}" if do else ""

    if vacuous:
        return Certificate(
            claim=f"CCE interval is vacuous at this regret level{intervened}: abstaining",
            estimand=estimand,
            kind=Kind.EMPIRICAL,
            value=interval,
            alpha=reported_alpha,
            assumptions=assumptions,
            method=_METHOD,
            witness=witness,
            hedge=Hedge(
                reason=(
                    "the epsilon-CCE interval spans the functional's full range — vacuous; "
                    "abstaining rather than certifying an uninformative bound"
                ),
                detail={"width": reported_width, "full_range": [full.lower, full.upper]},
            ),
            provenance=provenance,
            ci=reported_ci,
        )
    if licensed and degenerate and error.exact:
        midpoint = 0.5 * (interval.lower + interval.upper)
        return Certificate(
            claim=(
                "functional is constant over the CCE polytope: the equilibrium point prediction "
                f"is valid for every no-regret population{intervened}"
            ),
            estimand=estimand,
            kind=Kind.IDENTIFIED,
            value=midpoint,
            alpha=reported_alpha,
            assumptions=assumptions,
            method=_METHOD,
            witness=witness,
            hedge=None,
            provenance=provenance,
            ci=reported_ci,
        )
    degenerate_hedge = (
        Hedge(
            reason=(
                "the functional is constant over the polytope of the *estimated* payoff table, "
                "not demonstrably over the game's: an identification claim is withheld"
            ),
            detail={"width": width, "payoff_error": error.utility},
        )
        if degenerate and not error.exact
        else None
    )
    if licensed:
        finite_time = epsilon is not None
        route = (
            "measured realized regret (finite-time, no asymptotic assumption)"
            if finite_time
            else "no-regret learning (time-averaged play approaches the CCE set)"
        )
        return Certificate(
            claim=(
                f"time-averaged functional of the learning population lies in the CCE interval"
                f"{intervened} [{route}]"
            ),
            estimand=estimand,
            kind=Kind.BOUNDED,
            value=interval,
            alpha=reported_alpha,
            assumptions=assumptions,
            method=_METHOD,
            witness=witness,
            hedge=degenerate_hedge,
            provenance=provenance,
            ci=reported_ci,
        )
    return Certificate(
        claim=f"CCE interval reported as evidence only{intervened}",
        estimand=estimand,
        kind=Kind.EMPIRICAL,
        value=interval,
        alpha=reported_alpha,
        assumptions=assumptions,
        method=_METHOD,
        witness=witness,
        hedge=Hedge(
            reason=(
                "no-regret neither assumed nor measured for this population: the CCE interval "
                "is evidence, not a guarantee"
            ),
            detail={"width": width},
        ),
        provenance=provenance,
        ci=reported_ci,
    )

causalrl.magames.topology_max_kind(topology)

The strongest epistemic :class:~causalrl.certify.certificate.Kind topology licenses.

Source code in src/causalrl/magames/population.py

def topology_max_kind(topology: LearnerTopology) -> Kind:
    """The strongest epistemic :class:`~causalrl.certify.certificate.Kind` ``topology`` licenses."""
    return _TOPOLOGY_MAX_KIND[topology]

Estimation, Bounds & Certificates — Complete Reference

causalrl.certify.Assumption dataclass

A consumed assumption, optionally with a falsification/diagnostic result.

Source code in src/causalrl/certify/certificate.py

@dataclass(frozen=True)
class Assumption:
    """A consumed assumption, optionally with a falsification/diagnostic result."""

    name: str  # e.g. "MSM", "backdoor", "overlap", "mi-cap", "moment-condition"
    params: dict[str, Any] = field(default_factory=_empty_dict)
    checkable: bool = False
    diagnostic: dict[str, Any] | None = None

causalrl.estimate.estimators.EffectEstimate dataclass

A point estimate of an ATE with an (asymptotic) confidence interval and overlap diagnostics.

overlap records min_propensity / max_propensity of the estimated propensity score (nan for the propensity-free plugin estimator); a certifier downgrades to a hedge when positivity is destroyed (I3). std_error/ci are influence-function based for aipw and dml (the estimators with a valid asymptotic distribution); for plugin they are a coarse outcome-spread approximation and for ipw a self-normalised-weight approximation.

Source code in src/causalrl/estimate/estimators.py

@dataclass(frozen=True)
class EffectEstimate:
    """A point estimate of an ATE with an (asymptotic) confidence interval and overlap diagnostics.

    ``overlap`` records ``min_propensity`` / ``max_propensity`` of the estimated propensity score
    (``nan`` for the propensity-free ``plugin`` estimator); a certifier downgrades to a hedge when
    positivity is destroyed (I3). ``std_error``/``ci`` are influence-function based for ``aipw`` and
    ``dml`` (the estimators with a valid asymptotic distribution); for ``plugin`` they are a coarse
    outcome-spread approximation and for ``ipw`` a self-normalised-weight approximation.
    """

    value: float
    std_error: float
    ci: Interval
    alpha: float
    n: int
    method: str
    overlap: dict[str, float]
    n_folds: int | None = None

causalrl.certify.EstimandSpec dataclass

What a certificate is about.

Named EstimandSpec to avoid colliding with the shipped identification :class:causalrl.identification.id_algorithm.Estimand. target is the functional (mean/quantile/tail); only mean is used before Phase 1.

Source code in src/causalrl/certify/certificate.py

@dataclass(frozen=True)
class EstimandSpec:
    """What a certificate is about.

    Named ``EstimandSpec`` to avoid colliding with the shipped identification
    :class:`causalrl.identification.id_algorithm.Estimand`. ``target`` is the functional
    (``mean``/``quantile``/``tail``); only ``mean`` is used before Phase 1.
    """

    query: str  # see | do | counterfactual | policy_value | transport | equilibrium
    target: str = "mean"
    policy: str | None = None
    domains: tuple[str, ...] = ()

causalrl.certify.Hedge dataclass

Why a claim was refused or weakened (subsumes the shipped act/abstain verdict; I3).

Source code in src/causalrl/certify/certificate.py

@dataclass(frozen=True)
class Hedge:
    """Why a claim was refused or weakened (subsumes the shipped act/abstain verdict; I3)."""

    reason: str
    detail: dict[str, Any] | None = None
    downgraded_from: str | None = None  # original target when downgraded (e.g. "mean" -> quantile)

causalrl.certify.Kind

Bases: Enum

Epistemic status of a certificate's claim (I2). Never conflate the three.

Source code in src/causalrl/certify/certificate.py

class Kind(enum.Enum):
    """Epistemic status of a certificate's claim (I2). Never conflate the three."""

    IDENTIFIED = "identified"  # point-identified under the stated graph/assumptions
    BOUNDED = "bounded"  # partial identification under an explicit sensitivity budget
    EMPIRICAL = "empirical"  # simulation/sample evidence only; no identification guarantee

causalrl.certify.Provenance dataclass

Reproducibility record (I8): library version, seeds, data/graph fingerprints, timestamp.

Source code in src/causalrl/certify/certificate.py

@dataclass(frozen=True)
class Provenance:
    """Reproducibility record (I8): library version, seeds, data/graph fingerprints, timestamp."""

    library_version: str
    seeds: tuple[int, ...] = ()
    data_fingerprint: str | None = None
    graph_hash: str | None = None
    timestamp: str = ""  # ISO-8601 UTC

    @staticmethod
    def create(
        *,
        seeds: tuple[int, ...] = (),
        data_fingerprint: str | None = None,
        graph_hash: str | None = None,
    ) -> Provenance:
        """Fill ``library_version`` and a UTC ``timestamp`` automatically."""
        return Provenance(
            library_version=_library_version(),
            seeds=tuple(seeds),
            data_fingerprint=data_fingerprint,
            graph_hash=graph_hash,
            timestamp=_dt.datetime.now(_dt.UTC).isoformat(),
        )

create(*, seeds=(), data_fingerprint=None, graph_hash=None) staticmethod

Fill library_version and a UTC timestamp automatically.

Source code in src/causalrl/certify/certificate.py

@staticmethod
def create(
    *,
    seeds: tuple[int, ...] = (),
    data_fingerprint: str | None = None,
    graph_hash: str | None = None,
) -> Provenance:
    """Fill ``library_version`` and a UTC ``timestamp`` automatically."""
    return Provenance(
        library_version=_library_version(),
        seeds=tuple(seeds),
        data_fingerprint=data_fingerprint,
        graph_hash=graph_hash,
        timestamp=_dt.datetime.now(_dt.UTC).isoformat(),
    )

causalrl.ope.sequential.SequentialValueEstimate dataclass

A finite-horizon policy-value estimate with an (influence-function) CI and per-stage overlap.

overlap_by_step[t] records min_propensity / max_propensity of the stage-t estimated propensity (nan for the propensity-free gcomp estimator); a certifier downgrades to a hedge when any stage's positivity is destroyed (I3). std_error / ci are influence-function based for dr (cross-fitted) and a coarse pseudo-outcome-spread approximation for gcomp.

Source code in src/causalrl/ope/sequential.py

@dataclass(frozen=True)
class SequentialValueEstimate:
    """A finite-horizon policy-value estimate with an (influence-function) CI and per-stage overlap.

    ``overlap_by_step[t]`` records ``min_propensity`` / ``max_propensity`` of the stage-``t``
    estimated propensity (``nan`` for the propensity-free ``gcomp`` estimator); a certifier
    downgrades to a hedge when any stage's positivity is destroyed (I3). ``std_error`` / ``ci`` are
    influence-function based for ``dr`` (cross-fitted) and a coarse pseudo-outcome-spread
    approximation for ``gcomp``.
    """

    value: float
    std_error: float
    ci: Interval
    alpha: float
    n: int
    horizon: int
    method: str
    overlap_by_step: list[dict[str, float]]
    n_folds: int | None = None

causalrl.certify.Witness dataclass

Why an identified/bounded claim holds: adjustment set, transport formula, ID trace, ...

Source code in src/causalrl/certify/certificate.py

@dataclass(frozen=True)
class Witness:
    """Why an identified/bounded claim holds: adjustment set, transport formula, ID trace, ..."""

    kind: str
    detail: dict[str, Any] = field(default_factory=_empty_dict)

causalrl.bounds.continuous.certify_sensitivity_bounds(outcomes, *, gamma, propensities=None, treatment=None, covariates=None, propensity_model=None)

A kind=BOUNDED certificate wrapping :func:msm_sensitivity_bounds.

Source code in src/causalrl/bounds/continuous.py

def certify_sensitivity_bounds(
    outcomes: Sequence[float],
    *,
    gamma: float,
    propensities: Sequence[float] | None = None,
    treatment: Sequence[float] | None = None,
    covariates: Any | None = None,
    propensity_model: Classifier | None = None,
) -> Certificate:
    """A ``kind=BOUNDED`` certificate wrapping :func:`msm_sensitivity_bounds`."""
    interval = msm_sensitivity_bounds(
        outcomes,
        gamma=gamma,
        propensities=propensities,
        treatment=treatment,
        covariates=covariates,
        propensity_model=propensity_model,
    )
    known = propensities is not None
    return Certificate(
        claim=f"E[Y(1)] in [{interval.lower:.4g}, {interval.upper:.4g}] under MSM(Gamma={gamma})",
        estimand=EstimandSpec(query="do", target="mean"),
        kind=Kind.BOUNDED,
        value=interval,
        alpha=None,
        assumptions=(
            Assumption(name="MSM", params={"gamma": gamma}, checkable=False),
            Assumption(
                name="known-propensity" if known else "estimated-propensity",
                params={},
                checkable=True,
            ),
        ),
        method="msm-sensitivity" if known else "msm-sensitivity-estimated",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
    )

causalrl.ope.sequential.certify_sequential_value(histories, actions, target_actions, reward, *, method='dr', alpha=0.05, n_folds=5, seed=0, outcome_model=None, propensity_model=None, overlap_eps=0.01, clip=0.001, policy='pi')

Certify a deterministic policy's finite-horizon value under sequential ignorability.

Returns a kind=IDENTIFIED :class:Certificate with the point value (value), its influence-function confidence interval (ci), a per-stage sequential-ignorability witness, and provenance — or a hedged certificate when any stage's estimated positivity falls below overlap_eps (I3). Sequential ignorability is a non-checkable assumption the caller asserts (identification from a graph is out of scope); it is recorded explicitly, never assumed silently.

Source code in src/causalrl/ope/sequential.py

def certify_sequential_value(
    histories: Sequence[object],
    actions: Sequence[object],
    target_actions: Sequence[object],
    reward: object,
    *,
    method: str = "dr",
    alpha: float = 0.05,
    n_folds: int = 5,
    seed: int = 0,
    outcome_model: OutcomeFactory | None = None,
    propensity_model: PropensityFactory | None = None,
    overlap_eps: float = 0.01,
    clip: float = 1e-3,
    policy: str = "pi",
) -> Certificate:
    """Certify a deterministic policy's finite-horizon value under sequential ignorability.

    Returns a ``kind=IDENTIFIED`` :class:`Certificate` with the point value (``value``), its
    influence-function confidence interval (``ci``), a per-stage ``sequential-ignorability``
    witness, and provenance — or a hedged certificate when any stage's estimated positivity falls
    below ``overlap_eps`` (I3). Sequential ignorability is a non-checkable assumption the caller
    asserts (identification from a graph is out of scope); it is recorded explicitly, never assumed
    silently.
    """
    est = estimate_sequential_value(
        histories,
        actions,
        target_actions,
        reward,
        method=method,
        alpha=alpha,
        n_folds=n_folds,
        seed=seed,
        outcome_model=outcome_model,
        propensity_model=propensity_model,
        clip=clip,
    )
    act = [np.asarray(a, dtype=np.float64) for a in actions]
    y = np.asarray(reward, dtype=np.float64)
    fingerprint = _fingerprint(y, act)

    for t, ov in enumerate(est.overlap_by_step):
        min_e = ov.get("min_propensity", float("nan"))
        max_e = ov.get("max_propensity", float("nan"))
        if np.isfinite(min_e) and (min_e < overlap_eps or max_e > 1.0 - overlap_eps):
            return Certificate(
                claim=f"V({policy}) over horizon {est.horizon} refused: overlap-violation",
                estimand=EstimandSpec(query="policy_value", target="mean", policy=policy),
                kind=Kind.IDENTIFIED,
                value=None,
                alpha=alpha,
                assumptions=(),
                method="refused",
                witness=None,
                hedge=Hedge(
                    reason="overlap-violation",
                    detail={"step": t, "overlap_eps": overlap_eps, **ov},
                ),
                provenance=Provenance.create(seeds=(seed,), data_fingerprint=fingerprint),
            )

    xfit = f" (cross-fit K={est.n_folds})" if est.n_folds else ""
    return Certificate(
        claim=f"V({policy}) over horizon {est.horizon} = {est.value:.4g}",
        estimand=EstimandSpec(query="policy_value", target="mean", policy=policy),
        kind=Kind.IDENTIFIED,
        value=est.value,
        alpha=alpha,
        assumptions=(
            Assumption(
                name="sequential-ignorability",
                params={"horizon": est.horizon, "adjustment": "per-stage H_t"},
                checkable=False,
            ),
            Assumption(
                name="overlap",
                params={"eps": overlap_eps},
                checkable=True,
                diagnostic={"by_step": est.overlap_by_step},
            ),
        ),
        method=f"sequential-{est.method}{xfit}",
        witness=Witness(
            kind="sequential-adjustment",
            detail={"horizon": est.horizon, "stages": [f"H_{t}" for t in range(est.horizon)]},
        ),
        hedge=None,
        provenance=Provenance.create(seeds=(seed,), data_fingerprint=fingerprint),
        ci=est.ci,
    )

causalrl.eval.metrics.cumulative_regret(rewards, optimal_per_step)

Total regret = sum over steps of (optimal expected reward - realized reward).

Source code in src/causalrl/eval/metrics.py

def cumulative_regret(rewards: list[float], optimal_per_step: float) -> float:
    """Total regret = sum over steps of (optimal expected reward - realized reward)."""
    return sum(optimal_per_step - r for r in rewards)

causalrl.estimate.estimators.estimate_ate(data, treatment, outcome, adjustment, *, method='dml', alpha=0.05, n_folds=5, seed=0, outcome_model=None, propensity_model=None, clip=0.001)

Estimate the back-door ATE of binary treatment on outcome given adjustment.

data maps variable name -> 1-D array (covariates may be name -> (n,) or (n, k)). method selects the estimator (plugin/ipw/aipw/dml). outcome_model and propensity_model are zero-argument factories returning fresh sklearn-style learners (defaults: ridge outcome, logistic propensity). Raises ValueError for a non-binary treatment or a missing adjustment variable.

Source code in src/causalrl/estimate/estimators.py

def estimate_ate(
    data: Mapping[str, Any],
    treatment: str,
    outcome: str,
    adjustment: tuple[str, ...],
    *,
    method: str = "dml",
    alpha: float = 0.05,
    n_folds: int = 5,
    seed: int = 0,
    outcome_model: OutcomeFactory | None = None,
    propensity_model: PropensityFactory | None = None,
    clip: float = 1e-3,
) -> EffectEstimate:
    """Estimate the back-door ATE of binary ``treatment`` on ``outcome`` given ``adjustment``.

    ``data`` maps variable name -> 1-D array (covariates may be name -> ``(n,)`` or ``(n, k)``).
    ``method`` selects the estimator (``plugin``/``ipw``/``aipw``/``dml``). ``outcome_model`` and
    ``propensity_model`` are zero-argument factories returning fresh sklearn-style learners
    (defaults: ridge outcome, logistic propensity). Raises ``ValueError`` for a non-binary treatment
    or a missing adjustment variable.
    """
    if method not in _METHODS:
        raise ValueError(f"method must be one of {_METHODS}, got {method!r}")
    x = np.asarray(data[treatment], dtype=np.float64)
    y = np.asarray(data[outcome], dtype=np.float64)
    uniq = set(np.unique(x).tolist())
    if not uniq <= {0.0, 1.0}:
        raise ValueError(f"treatment {treatment!r} must be binary 0/1; got values {sorted(uniq)}")
    n = len(x)
    z = _assemble_z(data, adjustment, n)
    of = outcome_model or _default_outcome
    pf = propensity_model or _default_propensity

    n_folds_out: int | None = None
    if method == "plugin":
        value, se, ov = _plugin(x, z, y, outcome=of)
    elif method == "ipw":
        value, se, ov = _ipw(x, z, y, propensity=pf, clip=clip)
    elif method == "aipw":
        value, se, ov = _aipw(x, z, y, outcome=of, propensity=pf, clip=clip)
    else:  # dml
        value, se, ov = _dml(
            x, z, y, outcome=of, propensity=pf, clip=clip, n_folds=n_folds, seed=seed
        )
        n_folds_out = n_folds

    return EffectEstimate(
        value=value,
        std_error=se,
        ci=_ci(value, se, alpha),
        alpha=alpha,
        n=n,
        method=method,
        overlap=ov,
        n_folds=n_folds_out,
    )

causalrl.ope.sequential.estimate_sequential_value(histories, actions, target_actions, reward, *, method='dr', alpha=0.05, n_folds=5, seed=0, outcome_model=None, propensity_model=None, clip=0.001)

Estimate a deterministic target policy's finite-horizon value under sequential ignorability.

histories[t] is the stage-t adjustment history H_t (shape (n,) or (n, d_t)); actions[t] / target_actions[t] are binary (n,) observed / target actions. method is "gcomp" (g-computation) or "dr" (cross-fitted sequentially doubly-robust). Nuisances are pluggable sklearn-style factories (defaults: ridge outcome, logistic propensity).

Source code in src/causalrl/ope/sequential.py

def estimate_sequential_value(
    histories: Sequence[object],
    actions: Sequence[object],
    target_actions: Sequence[object],
    reward: object,
    *,
    method: str = "dr",
    alpha: float = 0.05,
    n_folds: int = 5,
    seed: int = 0,
    outcome_model: OutcomeFactory | None = None,
    propensity_model: PropensityFactory | None = None,
    clip: float = 1e-3,
) -> SequentialValueEstimate:
    """Estimate a deterministic target policy's finite-horizon value under sequential ignorability.

    ``histories[t]`` is the stage-``t`` adjustment history ``H_t`` (shape ``(n,)`` or ``(n, d_t)``);
    ``actions[t]`` / ``target_actions[t]`` are binary ``(n,)`` observed / target actions.
    ``method`` is ``"gcomp"`` (g-computation) or ``"dr"`` (cross-fitted sequentially doubly-robust).
    Nuisances are pluggable sklearn-style factories (defaults: ridge outcome, logistic propensity).
    """
    if method not in _METHODS:
        raise ValueError(f"method must be one of {_METHODS}, got {method!r}")
    hist = [_as2d(h) for h in histories]
    act = [np.asarray(a, dtype=np.float64) for a in actions]
    target = [np.asarray(a, dtype=np.float64) for a in target_actions]
    y = np.asarray(reward, dtype=np.float64)
    horizon, n = _check_shapes(hist, act, target, y)
    of = outcome_model or _default_outcome
    pf = propensity_model or _default_propensity

    if method == "gcomp":
        q1 = _gcomp(hist, act, target, y, outcome_factory=of)
        value = float(q1.mean())
        se = float(q1.std(ddof=1) / np.sqrt(n)) if n > 1 else float("nan")
        overlap_by_step = [
            {"min_propensity": float("nan"), "max_propensity": float("nan")} for _ in range(horizon)
        ]
        return SequentialValueEstimate(
            value, se, _ci(value, se, alpha), alpha, n, horizon, method, overlap_by_step, None
        )

    idx: IntArray = np.arange(n, dtype=np.intp)
    rng = np.random.default_rng(seed)
    rng.shuffle(idx)
    q1 = np.empty(n, dtype=np.float64)
    overlaps: list[list[FloatArray]] = [[] for _ in range(horizon)]
    for te_raw in np.array_split(idx, n_folds):
        te = np.asarray(te_raw, dtype=np.intp)
        tr = np.asarray(np.setdiff1d(idx, te, assume_unique=True), dtype=np.intp)
        q1[te] = _dr_fold(
            tr,
            te,
            hist,
            act,
            target,
            y,
            outcome_factory=of,
            propensity_factory=pf,
            clip=clip,
            overlaps=overlaps,
        )
    value = float(q1.mean())
    se = float(q1.std(ddof=1) / np.sqrt(n)) if n > 1 else float("nan")
    overlap_by_step = [
        {
            "min_propensity": float(np.concatenate(o).min()),
            "max_propensity": float(np.concatenate(o).max()),
        }
        for o in overlaps
    ]
    return SequentialValueEstimate(
        value, se, _ci(value, se, alpha), alpha, n, horizon, method, overlap_by_step, n_folds
    )

causalrl.eval.metrics.finite_horizon_regret(returns, optimal_return)

Total regret over episodes = sum of (optimal episode return - realized return).

Source code in src/causalrl/eval/metrics.py

def finite_horizon_regret(returns: list[float], optimal_return: float) -> float:
    """Total regret over episodes = sum of (optimal episode return - realized return)."""
    return sum(optimal_return - r for r in returns)

causalrl.certify.identify_effect_certified(graph, treatment, outcome)

identify_effect as an IDENTIFIED certificate; raises NotIdentifiableError if not.

The witness carries the do-free identification formula. Identification is symbolic, so value is None (there is no numeric estimate to report — estimation is Phase 1).

Source code in src/causalrl/certify/routines.py

def identify_effect_certified(
    graph: CausalGraph, treatment: Iterable[str], outcome: Iterable[str]
) -> Certificate:
    """`identify_effect` as an ``IDENTIFIED`` certificate; raises ``NotIdentifiableError`` if not.

    The witness carries the do-free identification formula. Identification is symbolic, so ``value``
    is ``None`` (there is no numeric estimate to report — estimation is Phase 1).
    """
    treat = sorted(treatment)
    out = sorted(outcome)
    # raises NotIdentifiableError (witnessing hedge); return_certificate=False avoids the flip warn
    estimand = identify_effect(graph, treat, out, return_certificate=False)
    # An identification claim that rests on an unfalsifiable assertion must say so: the graph's
    # complete-parent declarations are not derivable from data, and dropping them silently would
    # let a certificate look exactly like one earned by the graph alone.
    declared = graph.complete_parents
    assumptions = (
        (
            Assumption(
                name="parents-complete-by-design",
                params=dict(declared),
                checkable=False,
            ),
        )
        if declared
        else ()
    )
    return Certificate(
        claim=f"P({','.join(out)} | do({','.join(treat)})) is identified",
        estimand=EstimandSpec(query="do", target="mean", domains=tuple(out)),
        kind=Kind.IDENTIFIED,
        value=None,
        alpha=None,
        assumptions=assumptions,
        method="identify_effect",
        witness=Witness(
            "id-formula", {"formula": estimand.render(), "treatment": treat, "outcome": out}
        ),
        hedge=None,
        provenance=Provenance.create(graph_hash=graph_hash(graph)),
    )

causalrl.certify.ipw_sensitivity_bounds_certified(outcomes, propensities, *, gamma)

ipw_sensitivity_bounds as a BOUNDED certificate (Tan MSM on E[Y(1)]).

Source code in src/causalrl/certify/routines.py

def ipw_sensitivity_bounds_certified(
    outcomes: Sequence[float], propensities: Sequence[float], *, gamma: float
) -> Certificate:
    """`ipw_sensitivity_bounds` as a ``BOUNDED`` certificate (Tan MSM on ``E[Y(1)]``)."""
    interval: Interval = ipw_sensitivity_bounds(
        outcomes, propensities, gamma=gamma, return_certificate=False
    )
    return Certificate(
        claim=f"E[Y(1)] bounded under the marginal sensitivity model (gamma={gamma})",
        estimand=EstimandSpec(query="do", target="mean"),
        kind=Kind.BOUNDED,
        value=interval,
        alpha=None,
        assumptions=(Assumption("MSM", {"gamma": gamma}),),
        method="ipw_sensitivity_bounds",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
    )

causalrl.ope.ipw.ipw_value(actions, rewards, behavior_probs, target_probs)

Inverse-propensity-weighted off-policy value estimate.

Samples with zero behavior propensity (b == 0) contribute 0 to the sum but are still counted in the n denominator, which biases the estimate toward 0 — pass only samples with positive behavior propensity for an unbiased estimate.

Source code in src/causalrl/ope/ipw.py

def ipw_value(
    actions: list[int],
    rewards: list[float],
    behavior_probs: list[float],
    target_probs: list[float],
) -> float:
    """Inverse-propensity-weighted off-policy value estimate.

    Samples with zero behavior propensity (``b == 0``) contribute 0 to the sum but are
    still counted in the ``n`` denominator, which biases the estimate toward 0 — pass only
    samples with positive behavior propensity for an unbiased estimate.
    """
    n = len(actions)
    total = 0.0
    for _a, r, b, t in zip(actions, rewards, behavior_probs, target_probs, strict=True):
        total += (t / b) * r if b > 0 else 0.0
    return total / n

causalrl.certify.msm_policy_value_bounds_certified(outcomes, logging_propensities, target_propensities, *, gamma)

msm_policy_value_bounds as a BOUNDED certificate (off-policy value under Tan MSM).

Source code in src/causalrl/certify/routines.py

def msm_policy_value_bounds_certified(
    outcomes: Sequence[float],
    logging_propensities: Sequence[float],
    target_propensities: Sequence[float],
    *,
    gamma: float,
) -> Certificate:
    """`msm_policy_value_bounds` as a ``BOUNDED`` certificate (off-policy value under Tan MSM)."""
    interval: Interval = msm_policy_value_bounds(
        outcomes, logging_propensities, target_propensities, gamma=gamma, return_certificate=False
    )
    return Certificate(
        claim=f"V(pi_t) bounded under the marginal sensitivity model (gamma={gamma})",
        estimand=EstimandSpec(query="policy_value", target="mean", policy="pi_t"),
        kind=Kind.BOUNDED,
        value=interval,
        alpha=None,
        assumptions=(Assumption("MSM", {"gamma": gamma}),),
        method="msm_policy_value_bounds",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
    )

causalrl.eval.benchmark.report_to_dict(report)

Convert a report into JSON-serializable dictionaries.

Source code in src/causalrl/eval/benchmark.py

def report_to_dict(report: dict[str, BenchmarkEstimate]) -> dict[str, dict[str, object]]:
    """Convert a report into JSON-serializable dictionaries."""
    return {name: estimate.to_dict() for name, estimate in report.items()}

causalrl.eval.harness.run_episodes(agent, env, n_episodes, seed)

Run an agent on a finite-horizon env for n_episodes, returning per-episode returns.

The agent acts each step, is updated with the step reward, and is shown the transition via observe_step -- the representation-neutral hook, so an agent whose states are feature vectors can be driven by this harness too; the tabular default forwards to observe_transition unchanged. The per-episode return is the sum of step rewards (terminal-reward envs put it all on the last step).

Source code in src/causalrl/eval/harness.py

def run_episodes(agent: Agent, env: Any, n_episodes: int, seed: int) -> list[float]:
    """Run an agent on a finite-horizon env for n_episodes, returning per-episode returns.

    The agent acts each step, is updated with the step reward, and is shown the transition via
    ``observe_step`` -- the representation-neutral hook, so an agent whose states are feature
    vectors can be driven by this harness too; the tabular default forwards to
    ``observe_transition`` unchanged. The per-episode return is the sum of step rewards
    (terminal-reward envs put it all on the last step).
    """
    returns: list[float] = []
    obs, _ = env.reset(seed=seed)
    for ep in range(n_episodes):
        if ep > 0:
            obs, _ = env.reset()
        done = False
        ep_return = 0.0
        while not done:
            action = agent.act(obs)
            next_obs, reward, terminated, truncated, _info = env.step(action)
            done = bool(terminated or truncated)
            agent.update(obs, action, float(reward))
            agent.observe_step(obs, action, float(reward), next_obs, done)
            obs = next_obs
            ep_return += float(reward)
        returns.append(ep_return)
    return returns

causalrl.eval.mbrl_probe.run_m0_kill_gate(*, seeds=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), n=5000)

Back-door-adjusted vs naive-marginal on the Simpson bandit; report interventional values.

Keys: causal (back-door-adjusted agent), naive (marginal E[Y|A] agent), optimal (oracle). GO iff causal.mean > naive.mean.

Source code in src/causalrl/eval/mbrl_probe.py

def run_m0_kill_gate(
    *,
    seeds: Sequence[int] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
    n: int = 5000,
) -> dict[str, BenchmarkEstimate]:
    """Back-door-adjusted vs naive-marginal on the Simpson bandit; report interventional values.

    Keys: ``causal`` (back-door-adjusted agent), ``naive`` (marginal ``E[Y|A]`` agent), ``optimal``
    (oracle). GO iff ``causal.mean > naive.mean``.
    """
    rows: dict[str, list[float]] = {"causal": [], "naive": [], "optimal": []}
    for seed in seeds:
        env = SimpsonBandit(seed=seed)
        data = env.sample(n, seed=seed)

        # Naive marginal agent: sees only (action, reward) in a single context -> E[Y | A].
        transitions = [
            Transition(0, int(a), float(y), 0, True)
            for a, y in zip(data["A"], data["Y"], strict=True)
        ]
        dataset = ConfoundedTrajectoryDataset(transitions, n_states=1, n_actions=2)
        naive = NaiveOffline(env.n_states, env.n_actions)
        naive.ingest_offline(dataset)

        # Causal agent: back-door-adjusts for the observed confounder Z.
        causal = BackdoorAdjustedAgent(env.n_actions, graph=env.graph)
        causal.fit(data)

        rows["causal"].append(env.true_action_value(causal.act({"state": 0})))
        rows["naive"].append(env.true_action_value(naive.act({"state": 0})))
        rows["optimal"].append(env.optimal_value)

    return {
        name: BenchmarkEstimate.from_values(name, seeds=seeds, values=values)
        for name, values in rows.items()
    }

causalrl.eval.mbrl_probe.run_m1_discovery_gate(*, seeds=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), n=5000)

M1: an agent that learns the structure (interventional discovery) then adjusts, vs naive.

Keys: discovery (discover -> back-door-adjust), naive (marginal), optimal (oracle). GO iff discovery.mean > naive.mean (the recovered adjustment set is checked in the tests).

Source code in src/causalrl/eval/mbrl_probe.py

def run_m1_discovery_gate(
    *,
    seeds: Sequence[int] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
    n: int = 5000,
) -> dict[str, BenchmarkEstimate]:
    """M1: an agent that *learns* the structure (interventional discovery) then adjusts, vs naive.

    Keys: ``discovery`` (discover -> back-door-adjust), ``naive`` (marginal), ``optimal`` (oracle).
    GO iff ``discovery.mean > naive.mean`` (the recovered adjustment set is checked in the tests).
    """
    rows: dict[str, list[float]] = {"discovery": [], "naive": [], "optimal": []}
    for seed in seeds:
        env = SimpsonBandit(seed=seed)
        observational = env.sample(n, seed=seed)
        transitions = [
            Transition(0, int(a), float(y), 0, True)
            for a, y in zip(observational["A"], observational["Y"], strict=True)
        ]
        dataset = ConfoundedTrajectoryDataset(transitions, n_states=1, n_actions=2)
        naive = NaiveOffline(env.n_states, env.n_actions)
        naive.ingest_offline(dataset)

        agent = DiscoveryBackdoorAgent(env.n_actions, variables=("Z", "A", "Y"))
        agent.discover_and_fit(observational, tiers=(("Z",), ("A",), ("Y",)))

        rows["discovery"].append(env.true_action_value(agent.act({"state": 0})))
        rows["naive"].append(env.true_action_value(naive.act({"state": 0})))
        rows["optimal"].append(env.optimal_value)

    return {
        name: BenchmarkEstimate.from_values(name, seeds=seeds, values=values)
        for name, values in rows.items()
    }

causalrl.eval.mbrl_probe.run_m1b_dtr_gate(*, seeds=(0, 1, 2, 3, 4), horizon=2, n_episodes=8000)

M1b: deconfounded sequential planning vs naive on the confounded medicine DTR.

On SequentialDTREnv (a hidden comorbidity U, a confounded clinician who plays a = U, and a foresight gap), the deconfounded value-iteration agent DOVI (an existing library agent) is trained on confounded logs and rolled out; NaiveOffline is the confounded baseline. Keys: causal (DOVI late-window mean return), naive, optimal. GO iff causal > naive.

Source code in src/causalrl/eval/mbrl_probe.py

def run_m1b_dtr_gate(
    *,
    seeds: Sequence[int] = (0, 1, 2, 3, 4),
    horizon: int = 2,
    n_episodes: int = 8000,
) -> dict[str, BenchmarkEstimate]:
    """M1b: deconfounded sequential planning vs naive on the confounded medicine DTR.

    On ``SequentialDTREnv`` (a hidden comorbidity U, a confounded clinician who plays ``a = U``, and
    a foresight gap), the deconfounded value-iteration agent ``DOVI`` (an existing library agent) is
    trained on confounded logs and rolled out; ``NaiveOffline`` is the confounded baseline. Keys:
    ``causal`` (DOVI late-window mean return), ``naive``, ``optimal``. GO iff ``causal > naive``.
    """
    n_states = SequentialDTREnv(horizon=horizon).n_states
    optimal = SequentialDTREnv(horizon=horizon).optimal_value
    tail = max(1, n_episodes // 4)
    rows: dict[str, list[float]] = {"causal": [], "naive": [], "optimal": []}
    for seed in seeds:
        logs = generate_logs(
            SequentialDTREnv(horizon=horizon, seed=seed + 11), n_episodes=n_episodes, seed=seed + 11
        )
        dovi = DOVI(
            n_states=n_states,
            n_actions=2,
            horizon=horizon,
            seed=seed,
            transition_assumption="unconfounded",
        )
        dovi.ingest_offline(logs)
        dovi_returns = run_episodes(
            dovi, SequentialDTREnv(horizon=horizon, seed=seed), n_episodes=n_episodes, seed=seed
        )
        naive = NaiveOffline(n_states=n_states, n_actions=2)
        naive.ingest_offline(logs)
        naive_returns = run_episodes(
            naive, SequentialDTREnv(horizon=horizon, seed=seed), n_episodes=n_episodes, seed=seed
        )
        rows["causal"].append(sum(dovi_returns[-tail:]) / tail)
        rows["naive"].append(sum(naive_returns) / len(naive_returns))
        rows["optimal"].append(optimal)
    return {
        name: BenchmarkEstimate.from_values(name, seeds=seeds, values=values)
        for name, values in rows.items()
    }

causalrl.eval.mbrl_probe.run_m2_phase_diagram(*, gammas=(0.0, 0.25, 0.5, 0.75, 1.0), shifts=(0.0, 0.25, 0.5, 0.75, 1.0), seeds=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), n=6000, mono_tol=0.02)

M2: 2-D phase diagram of the causal-minus-naive post-shift gap over gamma x shift.

Per cell, a :class:~causalrl.agents.mbrl.TransportBackdoorAgent (deconfound + transport) and a NaiveOffline marginal agent are fit on SOURCE logs and evaluated at the TARGET; the gap is causal_target - naive_target. Returns per-cell :class:BenchmarkEstimates keyed by (gamma, shift) plus the mean-gap grid and whether it is monotone nondecreasing (within mono_tol) in each axis -- the "confounding bites where theory predicts" signature.

Source code in src/causalrl/eval/mbrl_probe.py

def run_m2_phase_diagram(
    *,
    gammas: Sequence[float] = (0.0, 0.25, 0.5, 0.75, 1.0),
    shifts: Sequence[float] = (0.0, 0.25, 0.5, 0.75, 1.0),
    seeds: Sequence[int] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
    n: int = 6000,
    mono_tol: float = 0.02,
) -> PhaseDiagram:
    """M2: 2-D phase diagram of the causal-minus-naive post-shift gap over ``gamma`` x ``shift``.

    Per cell, a :class:`~causalrl.agents.mbrl.TransportBackdoorAgent` (deconfound + transport) and
    a ``NaiveOffline`` marginal agent are fit on SOURCE logs and evaluated at the TARGET; the gap is
    ``causal_target - naive_target``. Returns per-cell :class:`BenchmarkEstimate`s keyed by
    ``(gamma, shift)`` plus the mean-gap grid and whether it is monotone nondecreasing (within
    ``mono_tol``) in each axis -- the "confounding bites where theory predicts" signature.
    """
    gammas = tuple(gammas)
    shifts = tuple(shifts)
    seeds = tuple(seeds)
    gap: dict[tuple[float, float], BenchmarkEstimate] = {}
    causal: dict[tuple[float, float], BenchmarkEstimate] = {}
    naive: dict[tuple[float, float], BenchmarkEstimate] = {}
    for g in gammas:
        for s in shifts:
            causal_vals: list[float] = []
            naive_vals: list[float] = []
            gap_vals: list[float] = []
            for seed in seeds:
                env = TransportableConfoundedBandit(gamma=g, shift=s, seed=seed)
                source = env.sample(n, domain="source", seed=seed)
                target_w = env.sample(n, domain="target", seed=seed + 10_000)["W"]
                agent = TransportBackdoorAgent(env.n_actions, graph=env.graph, transport=("W",))
                agent.fit(source, target_covariates={"W": target_w})
                causal_target = env.true_action_value(agent.act({"state": 0}), domain="target")

                transitions = [
                    Transition(0, int(a), float(y), 0, True)
                    for a, y in zip(source["A"], source["Y"], strict=True)
                ]
                dataset = ConfoundedTrajectoryDataset(transitions, n_states=1, n_actions=2)
                marginal = NaiveOffline(env.n_states, env.n_actions)
                marginal.ingest_offline(dataset)
                naive_target = env.true_action_value(marginal.act({"state": 0}), domain="target")

                causal_vals.append(causal_target)
                naive_vals.append(naive_target)
                gap_vals.append(causal_target - naive_target)
            key = (g, s)
            causal[key] = BenchmarkEstimate.from_values(
                f"causal_g{g}_s{s}", seeds=seeds, values=causal_vals
            )
            naive[key] = BenchmarkEstimate.from_values(
                f"naive_g{g}_s{s}", seeds=seeds, values=naive_vals
            )
            gap[key] = BenchmarkEstimate.from_values(f"gap_g{g}_s{s}", seeds=seeds, values=gap_vals)
    grid = np.array([[gap[(g, s)].mean for s in shifts] for g in gammas])
    mono_gamma = bool(np.all(np.diff(grid, axis=0) >= -mono_tol)) if len(gammas) > 1 else True
    mono_shift = bool(np.all(np.diff(grid, axis=1) >= -mono_tol)) if len(shifts) > 1 else True
    return PhaseDiagram(
        gammas=gammas,
        shifts=shifts,
        gap=gap,
        causal=causal,
        naive=naive,
        gap_grid=grid.tolist(),
        monotone_in_gamma=mono_gamma,
        monotone_in_shift=mono_shift,
    )

causalrl.eval.mbrl_probe.run_m3_function_approx_gate(*, seeds=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), gamma=1.0, n=6000)

M3: a continuous-confounder function-approximation agent vs naive on a nonlinear reward.

On ContinuousConfoundedBandit (continuous Z ~ Uniform, a nonlinear arm-1 bump, and behavior that over-samples arm 1 near the bump), a FunctionApproxBackdoorAgent (ridge on RBF features + back-door integration) recovers the true low value of arm 1 and keeps the safe arm 0; the NaiveOffline marginal is fooled into arm 1. Keys: causal, naive, optimal. GO iff causal.mean > naive.mean.

Source code in src/causalrl/eval/mbrl_probe.py

def run_m3_function_approx_gate(
    *,
    seeds: Sequence[int] = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9),
    gamma: float = 1.0,
    n: int = 6000,
) -> dict[str, BenchmarkEstimate]:
    """M3: a continuous-confounder function-approximation agent vs naive on a nonlinear reward.

    On ``ContinuousConfoundedBandit`` (continuous ``Z ~ Uniform``, a nonlinear arm-1 bump, and
    behavior that over-samples arm 1 near the bump), a ``FunctionApproxBackdoorAgent`` (ridge
    on RBF features + back-door integration) recovers the true low value of arm 1 and keeps the safe
    arm 0; the ``NaiveOffline`` marginal is fooled into arm 1. Keys: ``causal``, ``naive``,
    ``optimal``. GO iff ``causal.mean > naive.mean``.
    """
    rows: dict[str, list[float]] = {"causal": [], "naive": [], "optimal": []}
    for seed in seeds:
        env = ContinuousConfoundedBandit(gamma=gamma, seed=seed)
        data = env.sample(n, seed=seed)

        agent = FunctionApproxBackdoorAgent(env.n_actions, graph=env.graph)
        agent.fit(data)

        transitions = [
            Transition(0, int(a), float(y), 0, True)
            for a, y in zip(data["A"], data["Y"], strict=True)
        ]
        dataset = ConfoundedTrajectoryDataset(transitions, n_states=1, n_actions=2)
        naive = NaiveOffline(env.n_states, env.n_actions)
        naive.ingest_offline(dataset)

        rows["causal"].append(env.true_action_value(agent.act({"state": 0})))
        rows["naive"].append(env.true_action_value(naive.act({"state": 0})))
        rows["optimal"].append(env.optimal_value())
    return {
        name: BenchmarkEstimate.from_values(name, seeds=seeds, values=values)
        for name, values in rows.items()
    }

causalrl.ope.sequential.sequential_ice_values(histories, actions, target_actions, reward, *, outcome_model=None)

Per-unit iterated-conditional-expectation values q_1(H_1^i) under the target policy.

The g-computation backbone shared by :func:estimate_sequential_value (gcomp) and the identified sequential-transport subcase: the per-unit conditional policy value given the baseline history, before averaging — so a downstream caller may re-average it over a different (e.g. transported) baseline distribution.

Source code in src/causalrl/ope/sequential.py

def sequential_ice_values(
    histories: Sequence[object],
    actions: Sequence[object],
    target_actions: Sequence[object],
    reward: object,
    *,
    outcome_model: OutcomeFactory | None = None,
) -> FloatArray:
    """Per-unit iterated-conditional-expectation values ``q_1(H_1^i)`` under the target policy.

    The g-computation backbone shared by :func:`estimate_sequential_value` (``gcomp``) and the
    identified sequential-transport subcase: the per-unit conditional policy value given the
    baseline history, *before* averaging — so a downstream caller may re-average it over a
    different (e.g. transported) baseline distribution.
    """
    hist = [_as2d(h) for h in histories]
    act = [np.asarray(a, dtype=np.float64) for a in actions]
    target = [np.asarray(a, dtype=np.float64) for a in target_actions]
    y = np.asarray(reward, dtype=np.float64)
    _check_shapes(hist, act, target, y)
    of = outcome_model or _default_outcome
    return _gcomp(hist, act, target, y, outcome_factory=of)

Interop, Scale & Errors — Complete Reference

causalrl.interop.sbi_numpyro.PosteriorRegimeSampler

A reusable sampler over posterior regimes (a calibrated-configuration ensemble).

Holds the posterior matrix once; :meth:regimes materialises all regimes, :meth:sample draws a bootstrap subset, and :meth:mean_regime returns the single posterior-mean configuration.

Source code in src/causalrl/interop/sbi_numpyro.py

class PosteriorRegimeSampler:
    """A reusable sampler over posterior regimes (a calibrated-configuration ensemble).

    Holds the posterior matrix once; :meth:`regimes` materialises all regimes, :meth:`sample` draws
    a bootstrap subset, and :meth:`mean_regime` returns the single posterior-mean configuration.
    """

    def __init__(
        self,
        samples: PosteriorSamples,
        *,
        selection: Iterable[str] = (),
        name_prefix: str = "posterior",
    ) -> None:
        self._names, self._matrix = _as_named_matrix(samples)
        self._selection = frozenset(selection)
        self._name_prefix = name_prefix

    def __len__(self) -> int:
        return int(self._matrix.shape[0])

    @property
    def names(self) -> tuple[str, ...]:
        return tuple(self._names)

    def regimes(self) -> list[Regime]:
        """All posterior regimes."""
        return regimes_from_posterior(
            (self._matrix, self._names), selection=self._selection, name_prefix=self._name_prefix
        )

    def sample(self, n: int, *, seed: int | None = None) -> list[Regime]:
        """Draw ``n`` regimes uniformly with replacement (a bootstrap over the posterior)."""
        total = len(self)
        if total == 0:
            raise ValueError("empty posterior")
        rng = np.random.default_rng(seed)
        picks = rng.integers(0, total, size=n)
        out: list[Regime] = []
        for j in picks.tolist():
            params = {self._names[k]: float(self._matrix[j, k]) for k in range(len(self._names))}
            out.append(
                Regime.create(
                    f"{self._name_prefix}[{j}]", selection=self._selection, parameters=params
                )
            )
        return out

    def mean_regime(self, *, name: str = "posterior-mean") -> Regime:
        """The single regime at the posterior mean of every parameter."""
        if len(self) == 0:
            raise ValueError("empty posterior")
        means = self._matrix.mean(axis=0)
        params = {self._names[k]: float(means[k]) for k in range(len(self._names))}
        return Regime.create(name, selection=self._selection, parameters=params)

regimes()

All posterior regimes.

Source code in src/causalrl/interop/sbi_numpyro.py

def regimes(self) -> list[Regime]:
    """All posterior regimes."""
    return regimes_from_posterior(
        (self._matrix, self._names), selection=self._selection, name_prefix=self._name_prefix
    )

sample(n, *, seed=None)

Draw n regimes uniformly with replacement (a bootstrap over the posterior).

Source code in src/causalrl/interop/sbi_numpyro.py

def sample(self, n: int, *, seed: int | None = None) -> list[Regime]:
    """Draw ``n`` regimes uniformly with replacement (a bootstrap over the posterior)."""
    total = len(self)
    if total == 0:
        raise ValueError("empty posterior")
    rng = np.random.default_rng(seed)
    picks = rng.integers(0, total, size=n)
    out: list[Regime] = []
    for j in picks.tolist():
        params = {self._names[k]: float(self._matrix[j, k]) for k in range(len(self._names))}
        out.append(
            Regime.create(
                f"{self._name_prefix}[{j}]", selection=self._selection, parameters=params
            )
        )
    return out

mean_regime(*, name='posterior-mean')

The single regime at the posterior mean of every parameter.

Source code in src/causalrl/interop/sbi_numpyro.py

def mean_regime(self, *, name: str = "posterior-mean") -> Regime:
    """The single regime at the posterior mean of every parameter."""
    if len(self) == 0:
        raise ValueError("empty posterior")
    means = self._matrix.mean(axis=0)
    params = {self._names[k]: float(means[k]) for k in range(len(self._names))}
    return Regime.create(name, selection=self._selection, parameters=params)

causalrl.interop.columnar_sim.check_conformance(sim, *, n=16, seed=0)

Validate that sim emits a well-formed TrajectoryLog; return diagnostics or raise.

Calls sim.sample(n) and checks it returns a :class:TrajectoryLog, then probes do (a do-nothing intervention) to record whether interventions are supported. Returns the observed row count, the distinct value names, and whether do / a noise ledger are available — a quick conformance report for a new simulator integration.

Source code in src/causalrl/interop/columnar_sim.py

def check_conformance(sim: CausalEnvProtocol, *, n: int = 16, seed: int = 0) -> dict[str, Any]:
    """Validate that ``sim`` emits a well-formed ``TrajectoryLog``; return diagnostics or raise.

    Calls ``sim.sample(n)`` and checks it returns a :class:`TrajectoryLog`, then probes ``do`` (a
    do-nothing intervention) to record whether interventions are supported. Returns the observed row
    count, the distinct value names, and whether ``do`` / a noise ledger are available — a quick
    conformance report for a new simulator integration.
    """
    log = sim.sample(n, seed=seed)
    # Runtime guard: the protocol types this as TrajectoryLog; a non-conforming stand-in may not.
    if not isinstance(log, TrajectoryLog):  # pyright: ignore[reportUnnecessaryIsInstance]
        raise TypeError(f"sample() must return a TrajectoryLog, got {type(log).__name__}")
    names = sorted({str(x) for x in log.column("name")})
    supports_do = True
    try:
        sim.do({}, 1, seed=seed)
    except CausalInterfaceUnavailableError:
        supports_do = False
    return {
        "n_rows": len(log),
        "names": names,
        "supports_do": supports_do,
        "has_noise_ledger": sim.noise_ledger() is not None,
    }

causalrl.interop.causal_gym.from_causal_gym(env, behavior_policy, *, n_episodes=1000, max_steps=100, seed=0, n_states=None, n_actions=None, state_fn=None)

Roll env out under behavior_policy into a ConfoundedTrajectoryDataset.

behavior_policy(observation) -> action chooses the discrete action; state_fn maps an observation to a discrete state index (defaults to an int observation or obs['state']). n_states / n_actions are taken from the env's attributes when present, else from these arguments, else inferred from the logged data. Gymnasium reset(seed=...) -> (obs, info) and step(action) -> (obs, reward, terminated, truncated, info) are assumed (a legacy 4-tuple step is also tolerated).

Source code in src/causalrl/interop/causal_gym.py

def from_causal_gym(
    env: Any,
    behavior_policy: BehaviorPolicy,
    *,
    n_episodes: int = 1000,
    max_steps: int = 100,
    seed: int = 0,
    n_states: int | None = None,
    n_actions: int | None = None,
    state_fn: StateFn | None = None,
) -> ConfoundedTrajectoryDataset:
    """Roll ``env`` out under ``behavior_policy`` into a ``ConfoundedTrajectoryDataset``.

    ``behavior_policy(observation) -> action`` chooses the discrete action; ``state_fn`` maps
    an observation to a discrete state index (defaults to an int observation or ``obs['state']``).
    ``n_states`` / ``n_actions`` are taken from the env's attributes when present, else from these
    arguments, else inferred from the logged data. Gymnasium
    ``reset(seed=...) -> (obs, info)`` and ``step(action) -> (obs, reward, terminated, truncated,
    info)`` are assumed (a legacy 4-tuple ``step`` is also tolerated).
    """
    to_state = state_fn or _default_state_index
    transitions: list[Transition] = []
    max_state = 0
    max_action = 0
    for episode in range(n_episodes):
        obs = _reset(env, seed + episode)
        for _ in range(max_steps):
            state = to_state(obs)
            action = int(behavior_policy(obs))
            next_obs, reward, done = _step(env, action)
            next_state = to_state(next_obs)
            transitions.append(Transition(state, action, reward, next_state, done))
            max_state = max(max_state, state, next_state)
            max_action = max(max_action, action)
            obs = next_obs
            if done:
                break
    resolved_states = n_states or int(getattr(env, "n_states", 0)) or max_state + 1
    resolved_actions = n_actions or int(getattr(env, "n_actions", 0)) or max_action + 1
    return ConfoundedTrajectoryDataset(
        transitions, n_states=resolved_states, n_actions=resolved_actions
    )

causalrl.interop.pettingzoo.pettingzoo_to_trajectory_log(env, policies, *, n_episodes=1, max_steps=100, seed=0)

Roll env out under policies and log every agent-step to a TrajectoryLog.

policies[agent](observation) -> action drives each active agent. Each agent-step logs the observation columns, the chosen action, and the reward, keyed by entity_id (the agent's index), episode_id, and t. Raises KeyError if an active agent has no policy.

Source code in src/causalrl/interop/pettingzoo.py

def pettingzoo_to_trajectory_log(
    env: Any,
    policies: Mapping[str, Policy],
    *,
    n_episodes: int = 1,
    max_steps: int = 100,
    seed: int = 0,
) -> TrajectoryLog:
    """Roll ``env`` out under ``policies`` and log every agent-step to a ``TrajectoryLog``.

    ``policies[agent](observation) -> action`` drives each active agent. Each agent-step logs the
    observation columns, the chosen ``action``, and the ``reward``, keyed by ``entity_id`` (the
    agent's index), ``episode_id``, and ``t``. Raises ``KeyError`` if an active agent has no policy.
    """
    possible = list(getattr(env, "possible_agents", []))
    agent_index = {a: i for i, a in enumerate(possible)}
    rows: list[dict[str, Any]] = []
    for episode in range(n_episodes):
        obs = dict(_reset(env, seed + episode))
        for t in range(max_steps):
            active = list(getattr(env, "agents", []))
            if not active:
                break
            actions = {a: policies[a](obs.get(a)) for a in active}
            next_obs, rewards, terms, truncs = _step(env, actions)
            for a in active:
                ei = agent_index.get(a, len(agent_index))
                agent_index.setdefault(a, ei)
                cols = {**_obs_columns(obs.get(a)), "action": float(actions[a])}
                cols["reward"] = float(dict(rewards).get(a, 0.0))
                kinds = {"action": "action", "reward": "reward"}
                for name, value in cols.items():
                    rows.append(
                        {
                            "entity_id": ei,
                            "episode_id": episode,
                            "t": t,
                            "kind": kinds.get(name, "obs"),
                            "name": name,
                            "value": value,
                            "regime": "observed",
                            "observed": True,
                        }
                    )
            obs = dict(next_obs)
            terms_d, truncs_d = dict(terms), dict(truncs)
            if all(terms_d.get(a, False) or truncs_d.get(a, False) for a in active):
                break
    return TrajectoryLog.from_rows(rows)

causalrl.scm.bounded_fit.fit_scm_bounded(data, *, graph, value_ranges=None, families=None, holdout=0.2, seed=0)

Fit graph's identified mechanisms and bound the confounded ones.

The confounded-graph counterpart of :func:causalrl.fit_scm, which refuses this input. A node is bounded exactly when it shares a bidirected edge with one of its own parents; every other node is fitted point-wise by the same machinery, with the same families overrides and the same held-out scoring.

value_ranges gives (low, high) for each bounded node and is REQUIRED for them: a Manski bound is built out of the range the unobserved fraction could occupy, so a missing range is a missing bound rather than a default. Ranges for identified nodes are ignored.

Bounded nodes' confounded parents must be discrete with few levels -- their joint configuration indexes the arms of the bound.

Source code in src/causalrl/scm/bounded_fit.py

def fit_scm_bounded(
    data: Mapping[str, np.ndarray],
    *,
    graph: CausalGraph,
    value_ranges: Mapping[str, tuple[float, float]] | None = None,
    families: Mapping[str, MechanismFitter] | None = None,
    holdout: float = 0.2,
    seed: int = 0,
) -> BoundedSCMFit:
    """Fit ``graph``'s identified mechanisms and bound the confounded ones.

    The confounded-graph counterpart of :func:`causalrl.fit_scm`, which refuses this input. A node
    is bounded exactly when it shares a bidirected edge with one of its own parents; every other
    node is fitted point-wise by the same machinery, with the same ``families`` overrides and the
    same held-out scoring.

    ``value_ranges`` gives ``(low, high)`` for each bounded node and is REQUIRED for them: a Manski
    bound is built out of the range the unobserved fraction could occupy, so a missing range is a
    missing bound rather than a default. Ranges for identified nodes are ignored.

    Bounded nodes' confounded parents must be discrete with few levels -- their joint configuration
    indexes the arms of the bound.
    """
    ranges = dict(value_ranges or {})
    confounded: dict[str, tuple[str, ...]] = {}
    for node in graph.topological_order():
        parents = tuple(sorted(graph.parents(node)))
        confounded[node] = tuple(p for p in parents if graph.is_confounded(node, p))

    bounded_nodes = [n for n, ps in confounded.items() if ps]
    missing_ranges = sorted(n for n in bounded_nodes if n not in ranges)
    if missing_ranges:
        raise ValueError(
            f"value_ranges is missing bounded node(s) {missing_ranges}: each is confounded with "
            f"one of its own parents (e.g. {missing_ranges[0]} <-> "
            f"{confounded[missing_ranges[0]][0]}), so its mechanism is not identified and the "
            "only thing the data support is a Manski interval -- which is built from the range "
            "the node's unobserved values could occupy. Without that range there is no bound."
        )

    # Every node whose own parents are unconfounded is fitted exactly as fit_scm would: strip the
    # bidirected edges it refuses on, and fit that graph. Bounded nodes are refitted below, so
    # their point mechanism is computed and discarded -- cheap, and it keeps one fitting path.
    point_graph = CausalGraph(directed_edges=graph.directed_edges, nodes=graph.nodes)
    point_model = fit_scm(data, graph=point_graph, families=families, holdout=holdout, seed=seed)

    bounds: dict[str, FunctionalManskiBounds] = {}
    arm_index: dict[str, dict[tuple[float, ...], int]] = {}
    fits: list[BoundedNodeFit] = []
    n_samples = len(next(iter(data.values())))

    for node in graph.topological_order():
        parents = tuple(sorted(graph.parents(node)))
        confounders = confounded[node]
        if not confounders:
            fits.append(
                BoundedNodeFit(node=node, parents=parents, confounded_parents=(), identified=True)
            )
            continue

        level_sets = [_levels(data[p]) for p in confounders]
        n_arms = int(np.prod([len(levels) for levels in level_sets]))
        if n_arms > _MAX_ARMS:
            raise ValueError(
                f"node {node!r} has {n_arms} confounded-parent configurations "
                f"({', '.join(confounders)}), above the limit of {_MAX_ARMS}. Each is an arm "
                "needing its own outcome and propensity fit, so the bound would be vacuous at "
                "almost all of them. Coarsen the confounded parents first."
            )
        index = {combo: i for i, combo in enumerate(product(*level_sets))}
        covariates = [p for p in parents if p not in confounders]
        features = (
            np.column_stack([np.asarray(data[p], dtype=np.float64) for p in covariates])
            if covariates
            else np.zeros((n_samples, 1))
        )
        actions = np.array(
            [index[tuple(float(data[p][i]) for p in confounders)] for i in range(n_samples)],
            dtype=np.int_,
        )
        model = FunctionalManskiBounds(
            n_actions=len(index), reward_range=ranges[node], seed=seed
        ).fit(features, actions, np.asarray(data[node], dtype=np.float64))
        bounds[node] = model
        arm_index[node] = index
        fits.append(
            BoundedNodeFit(
                node=node,
                parents=parents,
                confounded_parents=confounders,
                identified=False,
                value_range=ranges[node],
                overlap=model.diagnostic(),
            )
        )

    return BoundedSCMFit(
        fits,
        n_samples=n_samples,
        point_model=point_model,
        bounds=bounds,
        arm_index=arm_index,
        seed=seed,
    )

causalrl.scm.bounded_fit.BoundedSCMFit

A fit over a confounded graph: point mechanisms where identified, intervals elsewhere.

Query it with :meth:interval. An identified node returns a degenerate interval (the point estimate, twice over) so that a caller can treat every node uniformly without first asking which kind it is -- and a caller that does care can ask :meth:is_identified.

Source code in src/causalrl/scm/bounded_fit.py

class BoundedSCMFit:
    """A fit over a confounded graph: point mechanisms where identified, intervals elsewhere.

    Query it with :meth:`interval`. An identified node returns a degenerate interval (the point
    estimate, twice over) so that a caller can treat every node uniformly without first asking
    which kind it is -- and a caller that *does* care can ask :meth:`is_identified`.
    """

    def __init__(
        self,
        nodes: Sequence[BoundedNodeFit],
        *,
        n_samples: int,
        point_model: StructuralCausalModel,
        bounds: Mapping[str, FunctionalManskiBounds],
        arm_index: Mapping[str, Mapping[tuple[float, ...], int]],
        seed: int,
    ) -> None:
        self._nodes = tuple(nodes)
        self.n_samples = int(n_samples)
        self._point_model = point_model
        self._bounds = dict(bounds)
        self._arm_index = {k: dict(v) for k, v in arm_index.items()}
        self._seed = int(seed)

    @property
    def nodes(self) -> tuple[BoundedNodeFit, ...]:
        """Per-node results, in the graph's topological order."""
        return self._nodes

    @property
    def bounded_nodes(self) -> tuple[str, ...]:
        """Nodes whose mechanism is not identified, and so is reported as an interval."""
        return tuple(f.node for f in self._nodes if not f.identified)

    @property
    def identified_nodes(self) -> tuple[str, ...]:
        """Nodes whose mechanism the data identify point-wise."""
        return tuple(f.node for f in self._nodes if f.identified)

    def is_identified(self, node: str) -> bool:
        """Whether ``node``'s mechanism was point-identified."""
        return self._fit_for(node).identified

    @property
    def point_model(self) -> StructuralCausalModel:
        """The fitted SCM, whose mechanisms are trustworthy at :attr:`identified_nodes` ONLY.

        Exposed because an identified node deserves the whole library -- ``do``/``see``,
        transport, counterfactuals -- rather than the single scalar :meth:`interval` returns. Its
        mechanisms at :attr:`bounded_nodes` are the regression the confounding invalidates, and are
        not to be used; they exist because one fitting path is fitted for every node at once.
        """
        return self._point_model

    def interval(
        self, node: str, assignment: Mapping[str, float], *, n_samples: int = 4096
    ) -> Interval:
        """Bound ``E[node | do(parents = assignment)]``.

        ``assignment`` must give a value for every parent of ``node``. An identified node returns
        a degenerate interval, estimated from :attr:`point_model` with ``n_samples`` draws; a
        confounded one returns the cross-fitted Manski bound, which is as wide as the logs'
        overlap at that arm requires and is vacuous where they carry no information at all.
        """
        fit = self._fit_for(node)
        missing = [p for p in fit.parents if p not in assignment]
        if missing:
            raise KeyError(
                f"assignment is missing parent(s) {sorted(missing)} of {node!r}: bounding "
                f"E[{node} | do(...)] needs a value for every parent ({list(fit.parents)})."
            )
        if fit.identified:
            value = self._point_value(node, assignment, n_samples)
            return Interval(value, value)

        arm_key = tuple(float(assignment[p]) for p in fit.confounded_parents)
        arms = self._arm_index[node]
        if arm_key not in arms:
            low, high = fit.value_range or (0.0, 1.0)
            return Interval(low, high)
        covariates = [p for p in fit.parents if p not in fit.confounded_parents]
        features = np.array(
            [[float(assignment[p]) for p in covariates]] if covariates else [[0.0]],
            dtype=np.float64,
        )
        low, high = self._bounds[node].bounds(features)
        arm = arms[arm_key]
        return Interval(float(low[0, arm]), float(high[0, arm]))

    def _point_value(self, node: str, assignment: Mapping[str, float], n: int) -> float:
        """``E[node | do(parents)]`` from the fitted SCM, by intervening and averaging.

        Estimated by sampling rather than read off the mechanism, so that every family works the
        same way -- an additive-noise mechanism, a tabular CPT and a pinned equation all answer
        this question, and none of them answers it through the same attribute. Deterministic given
        the fit's seed.
        """
        parents = {p: float(assignment[p]) for p in self._fit_for(node).parents}
        drawn = self._point_model.do(parents).see(n, seed=self._seed)
        return float(np.asarray(drawn[node], dtype=np.float64).mean())

    def _fit_for(self, node: str) -> BoundedNodeFit:
        for fit in self._nodes:
            if fit.node == node:
                return fit
        raise CausalGraphError(f"unknown node: {node!r}")

    def certificate(self) -> Certificate:
        """A ``BOUNDED`` certificate naming which nodes cost the model its point identification."""
        bounded = self.bounded_nodes
        return Certificate(
            claim=(
                f"fitted SCM over {len(self._nodes)} nodes; {len(self.identified_nodes)} "
                f"identified point-wise, {len(bounded)} bounded under latent confounding"
                + (f" ({', '.join(bounded)})" if bounded else "")
            ),
            estimand=EstimandSpec(query="do", target="mean"),
            kind=Kind.BOUNDED if bounded else Kind.IDENTIFIED,
            value=None,
            alpha=None,
            assumptions=(
                Assumption(
                    name="bounded-outcome-range",
                    params={"nodes": list(bounded)},
                    checkable=True,
                ),
                Assumption(name="correctly-specified-nuisances", params={}, checkable=False),
            ),
            method="manski-bounds-on-confounded-mechanisms",
            witness=None,
            hedge=None,
            provenance=Provenance.create(),
            ci=None,
        )

    def summary(self) -> str:
        lines = [f"BoundedSCMFit(n={self.n_samples})"]
        for fit in self._nodes:
            if fit.identified:
                lines.append(f"  {fit.node}: IDENTIFIED parents=[{', '.join(fit.parents)}]")
            else:
                overlap = f" {fit.overlap.summary()}" if fit.overlap is not None else ""
                lines.append(
                    f"  {fit.node}: BOUNDED confounded_with="
                    f"[{', '.join(fit.confounded_parents)}]{overlap}"
                )
        return "\n".join(lines)

nodes property

Per-node results, in the graph's topological order.

bounded_nodes property

Nodes whose mechanism is not identified, and so is reported as an interval.

identified_nodes property

Nodes whose mechanism the data identify point-wise.

point_model property

The fitted SCM, whose mechanisms are trustworthy at :attr:identified_nodes ONLY.

Exposed because an identified node deserves the whole library -- do/see, transport, counterfactuals -- rather than the single scalar :meth:interval returns. Its mechanisms at :attr:bounded_nodes are the regression the confounding invalidates, and are not to be used; they exist because one fitting path is fitted for every node at once.

is_identified(node)

Whether node's mechanism was point-identified.

Source code in src/causalrl/scm/bounded_fit.py

def is_identified(self, node: str) -> bool:
    """Whether ``node``'s mechanism was point-identified."""
    return self._fit_for(node).identified

interval(node, assignment, *, n_samples=4096)

Bound E[node | do(parents = assignment)].

assignment must give a value for every parent of node. An identified node returns a degenerate interval, estimated from :attr:point_model with n_samples draws; a confounded one returns the cross-fitted Manski bound, which is as wide as the logs' overlap at that arm requires and is vacuous where they carry no information at all.

Source code in src/causalrl/scm/bounded_fit.py

def interval(
    self, node: str, assignment: Mapping[str, float], *, n_samples: int = 4096
) -> Interval:
    """Bound ``E[node | do(parents = assignment)]``.

    ``assignment`` must give a value for every parent of ``node``. An identified node returns
    a degenerate interval, estimated from :attr:`point_model` with ``n_samples`` draws; a
    confounded one returns the cross-fitted Manski bound, which is as wide as the logs'
    overlap at that arm requires and is vacuous where they carry no information at all.
    """
    fit = self._fit_for(node)
    missing = [p for p in fit.parents if p not in assignment]
    if missing:
        raise KeyError(
            f"assignment is missing parent(s) {sorted(missing)} of {node!r}: bounding "
            f"E[{node} | do(...)] needs a value for every parent ({list(fit.parents)})."
        )
    if fit.identified:
        value = self._point_value(node, assignment, n_samples)
        return Interval(value, value)

    arm_key = tuple(float(assignment[p]) for p in fit.confounded_parents)
    arms = self._arm_index[node]
    if arm_key not in arms:
        low, high = fit.value_range or (0.0, 1.0)
        return Interval(low, high)
    covariates = [p for p in fit.parents if p not in fit.confounded_parents]
    features = np.array(
        [[float(assignment[p]) for p in covariates]] if covariates else [[0.0]],
        dtype=np.float64,
    )
    low, high = self._bounds[node].bounds(features)
    arm = arms[arm_key]
    return Interval(float(low[0, arm]), float(high[0, arm]))

certificate()

A BOUNDED certificate naming which nodes cost the model its point identification.

Source code in src/causalrl/scm/bounded_fit.py

def certificate(self) -> Certificate:
    """A ``BOUNDED`` certificate naming which nodes cost the model its point identification."""
    bounded = self.bounded_nodes
    return Certificate(
        claim=(
            f"fitted SCM over {len(self._nodes)} nodes; {len(self.identified_nodes)} "
            f"identified point-wise, {len(bounded)} bounded under latent confounding"
            + (f" ({', '.join(bounded)})" if bounded else "")
        ),
        estimand=EstimandSpec(query="do", target="mean"),
        kind=Kind.BOUNDED if bounded else Kind.IDENTIFIED,
        value=None,
        alpha=None,
        assumptions=(
            Assumption(
                name="bounded-outcome-range",
                params={"nodes": list(bounded)},
                checkable=True,
            ),
            Assumption(name="correctly-specified-nuisances", params={}, checkable=False),
        ),
        method="manski-bounds-on-confounded-mechanisms",
        witness=None,
        hedge=None,
        provenance=Provenance.create(),
        ci=None,
    )

causalrl.scm.bounded_fit.BoundedNodeFit dataclass

What one node's mechanism turned out to be: a point fit, or an interval.

identified is the field to branch on. When True the node was fitted exactly as :func:causalrl.fit_scm would have, and confounded_parents is empty. When False the node has at least one parent it shares a latent cause with, and every query about it comes back as an :class:~causalrl.Interval whose width is what that confounding costs.

Source code in src/causalrl/scm/bounded_fit.py

@dataclass(frozen=True)
class BoundedNodeFit:
    """What one node's mechanism turned out to be: a point fit, or an interval.

    ``identified`` is the field to branch on. When ``True`` the node was fitted exactly as
    :func:`causalrl.fit_scm` would have, and ``confounded_parents`` is empty. When ``False`` the
    node has at least one parent it shares a latent cause with, and every query about it comes back
    as an :class:`~causalrl.Interval` whose width is what that confounding costs.
    """

    node: str
    parents: tuple[str, ...]
    confounded_parents: tuple[str, ...]
    identified: bool
    value_range: tuple[float, float] | None = None
    overlap: OverlapDiagnostic | None = None

causalrl.transport.localize.localize_mechanism_shift(data_by_regime, *, graph, alpha=0.05)

Test each node's mechanism for invariance across regimes; emit the selection set.

data_by_regime maps a regime name to its columns -- at least two regimes, each carrying every node of graph. Columns must be integer-coded (the CI test is the discrete one); bin continuous variables first, and be aware that the test then asks about the binned mechanism.

Returns a :class:ShiftReport whose :attr:~ShiftReport.selection plugs straight into :func:~causalrl.identify_transport.

Source code in src/causalrl/transport/localize.py

def localize_mechanism_shift(
    data_by_regime: Mapping[str, Mapping[str, np.ndarray]],
    *,
    graph: CausalGraph,
    alpha: float = 0.05,
) -> ShiftReport:
    """Test each node's mechanism for invariance across regimes; emit the selection set.

    ``data_by_regime`` maps a regime name to its columns -- at least two regimes, each carrying
    every node of ``graph``. Columns must be integer-coded (the CI test is the discrete one); bin
    continuous variables first, and be aware that the test then asks about the binned mechanism.

    Returns a :class:`ShiftReport` whose :attr:`~ShiftReport.selection` plugs straight into
    :func:`~causalrl.identify_transport`.
    """
    if len(data_by_regime) < 2:
        raise ValueError(
            f"localize_mechanism_shift needs at least 2 regimes, got {len(data_by_regime)}: "
            "invariance is a statement about a comparison, and there is nothing to compare a "
            "single regime with."
        )
    if not 0.0 < alpha < 1.0:
        raise ValueError(f"alpha={alpha} must lie in (0, 1)")

    regimes = list(data_by_regime)
    nodes = graph.topological_order()
    for name, columns in data_by_regime.items():
        missing = [node for node in nodes if node not in columns]
        if missing:
            raise KeyError(f"regime {name!r} is missing column(s) for graph node(s): {missing}")

    pooled: dict[str, np.ndarray] = {}
    for node in nodes:
        pooled[node] = np.concatenate(
            [np.asarray(data_by_regime[r][node]).ravel() for r in regimes]
        )
    pooled[_REGIME_COLUMN] = np.concatenate(
        [
            np.full(len(np.asarray(data_by_regime[r][nodes[0]]).ravel()), i)
            for i, r in enumerate(regimes)
        ]
    )
    n = len(pooled[_REGIME_COLUMN])

    shifts: list[MechanismShift] = []
    for node in nodes:
        parents = tuple(sorted(graph.parents(node)))
        cmi = conditional_mutual_information(pooled, node, _REGIME_COLUMN, list(parents))
        statistic = 2.0 * n * cmi
        # Degrees of freedom for the conditional G-test: (levels-1)(regimes-1) per observed
        # parent configuration. Counting only OBSERVED configurations matters -- an unobserved
        # cell contributes nothing to the statistic, so charging df for it would only make the
        # test conservative in proportion to how sparse the parent space is.
        if parents:
            observed_configs = len({tuple(int(pooled[p][i]) for p in parents) for i in range(n)})
        else:
            observed_configs = 1
        df = max(1, (_levels(pooled[node]) - 1) * (len(regimes) - 1) * observed_configs)
        p_value = _chi2_sf(statistic, df)
        shifts.append(
            MechanismShift(
                node=node,
                parents=parents,
                cmi=cmi,
                statistic=statistic,
                df=df,
                p_value=p_value,
                shifted=p_value < alpha,
            )
        )
    return ShiftReport(shifts, regimes=regimes, alpha=alpha)

causalrl.transport.localize.ShiftReport

Which mechanisms moved between regimes, and the selection set that follows.

Source code in src/causalrl/transport/localize.py

class ShiftReport:
    """Which mechanisms moved between regimes, and the selection set that follows."""

    def __init__(self, shifts: Sequence[MechanismShift], *, regimes: Sequence[str], alpha: float):
        self._shifts = tuple(shifts)
        self.regimes = tuple(regimes)
        self.alpha = float(alpha)

    @property
    def shifts(self) -> tuple[MechanismShift, ...]:
        """Every node's verdict, in the graph's topological order -- not only the rejections."""
        return self._shifts

    @property
    def selection(self) -> frozenset[str]:
        """The shifted nodes, ready to hand to :func:`~causalrl.identify_transport`."""
        return frozenset(s.node for s in self._shifts if s.shifted)

    @property
    def invariant(self) -> frozenset[str]:
        """Nodes whose mechanism the data give no reason to think moved."""
        return frozenset(s.node for s in self._shifts if not s.shifted)

    def certificate(self) -> Certificate:
        """An ``EMPIRICAL`` certificate: this is a hypothesis test, not an identification result.

        Failing to reject invariance is not proof of it, and the certificate says so rather than
        reporting the invariant set as established. Anything built on top -- a transport formula
        keyed on :attr:`selection` -- inherits that.
        """
        shifted = sorted(self.selection)
        return Certificate(
            claim=(
                f"{len(shifted)} of {len(self._shifts)} mechanisms differ across regimes "
                f"{list(self.regimes)} at alpha={self.alpha:g}"
                + (f": {', '.join(shifted)}" if shifted else "")
            ),
            estimand=EstimandSpec(query="see", target="mean"),
            kind=Kind.EMPIRICAL,
            value=None,
            alpha=self.alpha,
            assumptions=(
                Assumption(
                    name="supplied-graph-is-correct",
                    params={"n_nodes": len(self._shifts)},
                    checkable=False,
                ),
                Assumption(
                    name="no-multiplicity-correction",
                    params={"n_tests": len(self._shifts)},
                    checkable=True,
                    diagnostic={s.node: s.p_value for s in self._shifts},
                ),
            ),
            method="conditional-invariance-test",
            witness=Witness("selection-set", {"selection": shifted}),
            hedge=None,
            provenance=Provenance.create(),
        )

    def summary(self) -> str:
        lines = [f"ShiftReport(regimes={list(self.regimes)}, alpha={self.alpha:g})"]
        for s in self._shifts:
            verdict = "SHIFTED " if s.shifted else "invariant"
            lines.append(
                f"  {verdict} {s.node}: p={s.p_value:.4f} (G={s.statistic:.2f}, df={s.df})"
            )
        return "\n".join(lines)

shifts property

Every node's verdict, in the graph's topological order -- not only the rejections.

selection property

The shifted nodes, ready to hand to :func:~causalrl.identify_transport.

invariant property

Nodes whose mechanism the data give no reason to think moved.

certificate()

An EMPIRICAL certificate: this is a hypothesis test, not an identification result.

Failing to reject invariance is not proof of it, and the certificate says so rather than reporting the invariant set as established. Anything built on top -- a transport formula keyed on :attr:selection -- inherits that.

Source code in src/causalrl/transport/localize.py

def certificate(self) -> Certificate:
    """An ``EMPIRICAL`` certificate: this is a hypothesis test, not an identification result.

    Failing to reject invariance is not proof of it, and the certificate says so rather than
    reporting the invariant set as established. Anything built on top -- a transport formula
    keyed on :attr:`selection` -- inherits that.
    """
    shifted = sorted(self.selection)
    return Certificate(
        claim=(
            f"{len(shifted)} of {len(self._shifts)} mechanisms differ across regimes "
            f"{list(self.regimes)} at alpha={self.alpha:g}"
            + (f": {', '.join(shifted)}" if shifted else "")
        ),
        estimand=EstimandSpec(query="see", target="mean"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=self.alpha,
        assumptions=(
            Assumption(
                name="supplied-graph-is-correct",
                params={"n_nodes": len(self._shifts)},
                checkable=False,
            ),
            Assumption(
                name="no-multiplicity-correction",
                params={"n_tests": len(self._shifts)},
                checkable=True,
                diagnostic={s.node: s.p_value for s in self._shifts},
            ),
        ),
        method="conditional-invariance-test",
        witness=Witness("selection-set", {"selection": shifted}),
        hedge=None,
        provenance=Provenance.create(),
    )

causalrl.transport.localize.MechanismShift dataclass

The invariance verdict at one node.

statistic is the G-test statistic 2 n I(V; R | Pa(V)) in nats-times-samples, which is asymptotically chi-squared under invariance; p_value is its upper tail. shifted is the verdict at the caller's alpha, and is what puts the node in the selection set.

Source code in src/causalrl/transport/localize.py

@dataclass(frozen=True)
class MechanismShift:
    """The invariance verdict at one node.

    ``statistic`` is the G-test statistic ``2 n I(V; R | Pa(V))`` in nats-times-samples, which is
    asymptotically chi-squared under invariance; ``p_value`` is its upper tail. ``shifted`` is the
    verdict at the caller's ``alpha``, and is what puts the node in the selection set.
    """

    node: str
    parents: tuple[str, ...]
    cmi: float
    statistic: float
    df: int
    p_value: float
    shifted: bool

causalrl.agents.anytime.Anytime

Bases: Protocol

An agent with a usable answer at every instant, refinable while time remains.

The two halves are inseparable: :meth:incumbent without :meth:refine is a fixed answer, and :meth:refine without :meth:incumbent is a computation you cannot interrupt. Together they are what lets a decision be handed a wall-clock budget instead of being run to convergence.

Source code in src/causalrl/agents/anytime.py

@runtime_checkable
class Anytime(Protocol):
    """An agent with a usable answer at every instant, refinable while time remains.

    The two halves are inseparable: :meth:`incumbent` without :meth:`refine` is a fixed answer, and
    :meth:`refine` without :meth:`incumbent` is a computation you cannot interrupt. Together they
    are what lets a decision be handed a wall-clock budget instead of being run to convergence.
    """

    def incumbent(self) -> Intervention:
        """The best answer found so far. Defined before any refinement has happened."""
        ...

    def refine(self, deadline: Deadline | None = None) -> bool:
        """Improve the incumbent within ``deadline``; return ``False`` once converged.

        ``False`` means further calls cannot help -- the caller should stop rather than spin. A
        return of ``True`` says only that more work remains, not that the last call used its whole
        budget.
        """
        ...

incumbent()

The best answer found so far. Defined before any refinement has happened.

Source code in src/causalrl/agents/anytime.py

def incumbent(self) -> Intervention:
    """The best answer found so far. Defined before any refinement has happened."""
    ...

refine(deadline=None)

Improve the incumbent within deadline; return False once converged.

False means further calls cannot help -- the caller should stop rather than spin. A return of True says only that more work remains, not that the last call used its whole budget.

Source code in src/causalrl/agents/anytime.py

def refine(self, deadline: Deadline | None = None) -> bool:
    """Improve the incumbent within ``deadline``; return ``False`` once converged.

    ``False`` means further calls cannot help -- the caller should stop rather than spin. A
    return of ``True`` says only that more work remains, not that the last call used its whole
    budget.
    """
    ...

causalrl.agents.anytime.AnytimeInterventionSearch

Bases: InterventionalAgent

Deadline-honouring random-shooting search over an :class:~causalrl.InterventionSpace.

Round 0 samples the whole space; each later round resamples around the incumbent inside a window that shrinks by shrink, so the answer improves monotonically and is usable after any number of rounds. value_fn(observation, intervention) scores a candidate -- typically a fitted outcome model, a BoundedSCMFit interval's midpoint, or an SCM's do expectation.

The deadline is checked between candidates, not only between rounds, so a budget that expires mid-round still returns the incumbent rather than overrunning to a round boundary. Nothing here interrupts value_fn itself: a single call that overruns the budget will overrun it, which is the cooperative contract :class:causalrl.Deadline documents.

Source code in src/causalrl/agents/anytime.py

class AnytimeInterventionSearch(InterventionalAgent):
    """Deadline-honouring random-shooting search over an :class:`~causalrl.InterventionSpace`.

    Round 0 samples the whole space; each later round resamples around the incumbent inside a
    window that shrinks by ``shrink``, so the answer improves monotonically and is usable after any
    number of rounds. ``value_fn(observation, intervention)`` scores a candidate -- typically a
    fitted outcome model, a ``BoundedSCMFit`` interval's midpoint, or an SCM's ``do`` expectation.

    The deadline is checked between candidates, not only between rounds, so a budget that expires
    mid-round still returns the incumbent rather than overrunning to a round boundary. Nothing here
    interrupts ``value_fn`` itself: a single call that overruns the budget will overrun it, which
    is the cooperative contract :class:`causalrl.Deadline` documents.
    """

    def __init__(
        self,
        value_fn: Callable[[Mapping[str, Any], Intervention], float],
        *,
        rounds: int = 8,
        candidates_per_round: int = 32,
        shrink: float = 0.5,
        seed: int | None = None,
    ) -> None:
        if rounds < 1:
            raise ValueError(f"rounds={rounds} must be at least 1")
        if candidates_per_round < 1:
            raise ValueError(f"candidates_per_round={candidates_per_round} must be at least 1")
        if not 0.0 < shrink <= 1.0:
            raise ValueError(
                f"shrink={shrink} must lie in (0, 1]: it multiplies the sampling window each "
                "round, so a value above 1 would widen the search around the incumbent and a "
                "non-positive one would collapse it to a point immediately."
            )
        self.value_fn = value_fn
        self.rounds = rounds
        self.candidates_per_round = candidates_per_round
        self.shrink = float(shrink)
        self._rng = np.random.default_rng(seed)
        self._incumbent: Intervention = {}
        self._best_value = -float("inf")
        self._report = SearchReport(0, rounds, 0, exhausted=False, best_value=-float("inf"))

    def incumbent(self) -> Intervention:
        """The best intervention found by the most recent :meth:`act`."""
        return dict(self._incumbent)

    @property
    def last_search(self) -> SearchReport:
        """How the most recent :meth:`act` terminated. Its certificate carries the hedge."""
        return self._report

    def act(
        self,
        observation: Mapping[str, Any],
        *,
        space: InterventionSpace,
        deadline: Deadline | None = None,
    ) -> Intervention:
        """Search ``space`` for the highest-value intervention within ``deadline``."""
        variables = sorted(space.variables)
        if not variables:
            self._incumbent = {}
            self._report = SearchReport(0, self.rounds, 0, exhausted=True, best_value=0.0)
            return {}

        self._incumbent = space.sample(variables, self._rng)
        self._best_value = float(self.value_fn(observation, self._incumbent))
        candidates = 1
        completed = 0

        for round_index in range(self.rounds):
            width = self.shrink**round_index
            for _ in range(self.candidates_per_round):
                if deadline is not None and deadline.expired():
                    self._report = SearchReport(
                        completed, self.rounds, candidates, False, self._best_value
                    )
                    return dict(self._incumbent)
                candidate = self._propose(space, variables, width)
                candidates += 1
                value = float(self.value_fn(observation, candidate))
                if value > self._best_value:
                    self._best_value, self._incumbent = value, candidate
            completed = round_index + 1

        self._report = SearchReport(completed, self.rounds, candidates, True, self._best_value)
        return dict(self._incumbent)

    def refine(self, deadline: Deadline | None = None) -> bool:
        """Not separately drivable: this agent's search runs inside :meth:`act`.

        Present because :class:`Anytime` requires it. It reports convergence (``False``) whenever
        the last :meth:`act` ran to completion, so a caller driving the protocol generically stops
        instead of spinning on a search that has nowhere left to go.
        """
        return not self._report.exhausted

    def _propose(
        self, space: InterventionSpace, variables: list[str], width: float
    ) -> Intervention:
        """One candidate: continuous variables jittered around the incumbent, discrete resampled.

        The window is a fraction ``width`` of each continuous domain, centred on the incumbent and
        clipped back into the domain, so every proposal is admissible by construction and no
        candidate is wasted on a rejection.
        """
        proposal: dict[str, Any] = {}
        for name in variables:
            domain = space.domain(name)
            if isinstance(domain, Continuous) and width < 1.0:
                span = (domain.high - domain.low) * width
                centre = float(self._incumbent[name])
                proposal[name] = domain.project(
                    self._rng.uniform(centre - span / 2.0, centre + span / 2.0)
                )
            else:
                proposal[name] = domain.sample(self._rng)
        return proposal

    def update(
        self, observation: Mapping[str, Any], intervention: Intervention, reward: float
    ) -> None:
        """No-op: this agent's policy is its ``value_fn``, which it does not fit.

        Learning the value model is a separate concern with its own entry points (``fit_scm``,
        ``FittedQIteration``, ``fit_scm_bounded``); this class is the search over one.
        """
        return None

last_search property

How the most recent :meth:act terminated. Its certificate carries the hedge.

incumbent()

The best intervention found by the most recent :meth:act.

Source code in src/causalrl/agents/anytime.py

def incumbent(self) -> Intervention:
    """The best intervention found by the most recent :meth:`act`."""
    return dict(self._incumbent)

act(observation, *, space, deadline=None)

Search space for the highest-value intervention within deadline.

Source code in src/causalrl/agents/anytime.py

def act(
    self,
    observation: Mapping[str, Any],
    *,
    space: InterventionSpace,
    deadline: Deadline | None = None,
) -> Intervention:
    """Search ``space`` for the highest-value intervention within ``deadline``."""
    variables = sorted(space.variables)
    if not variables:
        self._incumbent = {}
        self._report = SearchReport(0, self.rounds, 0, exhausted=True, best_value=0.0)
        return {}

    self._incumbent = space.sample(variables, self._rng)
    self._best_value = float(self.value_fn(observation, self._incumbent))
    candidates = 1
    completed = 0

    for round_index in range(self.rounds):
        width = self.shrink**round_index
        for _ in range(self.candidates_per_round):
            if deadline is not None and deadline.expired():
                self._report = SearchReport(
                    completed, self.rounds, candidates, False, self._best_value
                )
                return dict(self._incumbent)
            candidate = self._propose(space, variables, width)
            candidates += 1
            value = float(self.value_fn(observation, candidate))
            if value > self._best_value:
                self._best_value, self._incumbent = value, candidate
        completed = round_index + 1

    self._report = SearchReport(completed, self.rounds, candidates, True, self._best_value)
    return dict(self._incumbent)

refine(deadline=None)

Not separately drivable: this agent's search runs inside :meth:act.

Present because :class:Anytime requires it. It reports convergence (False) whenever the last :meth:act ran to completion, so a caller driving the protocol generically stops instead of spinning on a search that has nowhere left to go.

Source code in src/causalrl/agents/anytime.py

def refine(self, deadline: Deadline | None = None) -> bool:
    """Not separately drivable: this agent's search runs inside :meth:`act`.

    Present because :class:`Anytime` requires it. It reports convergence (``False``) whenever
    the last :meth:`act` ran to completion, so a caller driving the protocol generically stops
    instead of spinning on a search that has nowhere left to go.
    """
    return not self._report.exhausted

update(observation, intervention, reward)

No-op: this agent's policy is its value_fn, which it does not fit.

Learning the value model is a separate concern with its own entry points (fit_scm, FittedQIteration, fit_scm_bounded); this class is the search over one.

Source code in src/causalrl/agents/anytime.py

def update(
    self, observation: Mapping[str, Any], intervention: Intervention, reward: float
) -> None:
    """No-op: this agent's policy is its ``value_fn``, which it does not fit.

    Learning the value model is a separate concern with its own entry points (``fit_scm``,
    ``FittedQIteration``, ``fit_scm_bounded``); this class is the search over one.
    """
    return None

causalrl.agents.anytime.SearchReport dataclass

How much of the search actually ran, and whether the budget stopped it.

exhausted is the load-bearing field: True means the search finished the work it intended, False means a deadline cut it off and the answer is the best of a subset.

Source code in src/causalrl/agents/anytime.py

@dataclass(frozen=True)
class SearchReport:
    """How much of the search actually ran, and whether the budget stopped it.

    ``exhausted`` is the load-bearing field: ``True`` means the search finished the work it
    intended, ``False`` means a deadline cut it off and the answer is the best of a *subset*.
    """

    rounds: int
    rounds_planned: int
    candidates: int
    exhausted: bool
    best_value: float

    def certificate(self) -> Certificate:
        """An ``EMPIRICAL`` certificate, hedged when the budget truncated the search.

        Always ``EMPIRICAL``: a search over sampled candidates is sample evidence about which
        intervention looks best, and carries no identification guarantee whatever its budget. What
        the deadline changes is completeness, and that is what the hedge reports.
        """
        truncated = not self.exhausted
        return Certificate(
            claim=(
                f"best of {self.candidates} sampled interventions over {self.rounds} of "
                f"{self.rounds_planned} planned refinement rounds "
                f"(value {self.best_value:.6g})"
            ),
            estimand=EstimandSpec(query="do", target="mean"),
            kind=Kind.EMPIRICAL,
            value=None,
            alpha=None,
            assumptions=(
                Assumption(
                    name="search-completeness",
                    params={"rounds": self.rounds, "rounds_planned": self.rounds_planned},
                    checkable=True,
                    diagnostic={"exhausted": self.exhausted, "candidates": self.candidates},
                ),
            ),
            method="anytime-random-shooting",
            witness=None,
            hedge=(
                Hedge(
                    reason=(
                        "budget-truncated: the deadline expired after "
                        f"{self.rounds} of {self.rounds_planned} refinement rounds, so this is "
                        "the best of a subset of the intended search, not of all of it"
                    ),
                    detail={"rounds": self.rounds, "rounds_planned": self.rounds_planned},
                    downgraded_from="exhaustive-search",
                )
                if truncated
                else None
            ),
            provenance=Provenance.create(),
        )

certificate()

An EMPIRICAL certificate, hedged when the budget truncated the search.

Always EMPIRICAL: a search over sampled candidates is sample evidence about which intervention looks best, and carries no identification guarantee whatever its budget. What the deadline changes is completeness, and that is what the hedge reports.

Source code in src/causalrl/agents/anytime.py

def certificate(self) -> Certificate:
    """An ``EMPIRICAL`` certificate, hedged when the budget truncated the search.

    Always ``EMPIRICAL``: a search over sampled candidates is sample evidence about which
    intervention looks best, and carries no identification guarantee whatever its budget. What
    the deadline changes is completeness, and that is what the hedge reports.
    """
    truncated = not self.exhausted
    return Certificate(
        claim=(
            f"best of {self.candidates} sampled interventions over {self.rounds} of "
            f"{self.rounds_planned} planned refinement rounds "
            f"(value {self.best_value:.6g})"
        ),
        estimand=EstimandSpec(query="do", target="mean"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=None,
        assumptions=(
            Assumption(
                name="search-completeness",
                params={"rounds": self.rounds, "rounds_planned": self.rounds_planned},
                checkable=True,
                diagnostic={"exhausted": self.exhausted, "candidates": self.candidates},
            ),
        ),
        method="anytime-random-shooting",
        witness=None,
        hedge=(
            Hedge(
                reason=(
                    "budget-truncated: the deadline expired after "
                    f"{self.rounds} of {self.rounds_planned} refinement rounds, so this is "
                    "the best of a subset of the intended search, not of all of it"
                ),
                detail={"rounds": self.rounds, "rounds_planned": self.rounds_planned},
                downgraded_from="exhaustive-search",
            )
            if truncated
            else None
        ),
        provenance=Provenance.create(),
    )

causalrl.intervention.Continuous dataclass

A closed real interval [low, high] of admissible values.

The action-side counterpart of :mod:causalrl.state's feature vectors. The estimation core never needed a discretisation -- cross-fitted DML, additive-noise mechanisms, the continuous bounds and the RBF-encoded function approximators all take real-valued inputs -- so a continuous confounder was always expressible while a continuous treatment was not. A dose, a budget, a duration or any other quantity chosen on a scale is the ordinary case.

Deliberately NOT named Interval: :class:causalrl.Interval is a bound on an estimand's value, and reusing the name for an action domain would make every docstring mentioning one ambiguous about which it meant.

Source code in src/causalrl/intervention.py

@dataclass(frozen=True)
class Continuous:
    """A closed real interval ``[low, high]`` of admissible values.

    The action-side counterpart of :mod:`causalrl.state`'s feature vectors. The estimation core
    never needed a discretisation -- cross-fitted DML, additive-noise mechanisms, the continuous
    bounds and the RBF-encoded function approximators all take real-valued inputs -- so a
    continuous *confounder* was always expressible while a continuous *treatment* was not. A dose, a
    budget, a duration or any other quantity chosen on a scale is the ordinary case.

    Deliberately NOT named ``Interval``: :class:`causalrl.Interval` is a bound on an estimand's
    value, and reusing the name for an action domain would make every docstring mentioning one
    ambiguous about which it meant.
    """

    low: float
    high: float

    def __post_init__(self) -> None:
        low, high = float(self.low), float(self.high)
        if not low <= high:
            raise ValueError(
                f"Continuous(low={self.low}, high={self.high}) must satisfy low <= high: an "
                "inverted interval admits nothing, which makes every intervention on this "
                "variable inadmissible while the variable still counts as manipulable."
            )
        object.__setattr__(self, "low", low)
        object.__setattr__(self, "high", high)

    def contains(self, value: Any) -> bool:
        try:
            return self.low <= float(value) <= self.high
        except (TypeError, ValueError):
            return False

    def project(self, value: Any) -> float:
        """``value`` clipped into ``[low, high]``."""
        return min(self.high, max(self.low, float(value)))

    def sample(self, rng: np.random.Generator) -> float:
        return float(rng.uniform(self.low, self.high))

project(value)

value clipped into [low, high].

Source code in src/causalrl/intervention.py

def project(self, value: Any) -> float:
    """``value`` clipped into ``[low, high]``."""
    return min(self.high, max(self.low, float(value)))

causalrl.intervention.Discrete dataclass

A finite set of admissible values -- the only domain kind before continuous ones existed.

Source code in src/causalrl/intervention.py

@dataclass(frozen=True)
class Discrete:
    """A finite set of admissible values -- the only domain kind before continuous ones existed."""

    values: tuple[Any, ...]

    def contains(self, value: Any) -> bool:
        return value in self.values

    def project(self, value: Any) -> Any:
        """The admissible value closest to ``value``; ``value`` itself when already admissible.

        Closeness is numeric where the values are numeric and identity otherwise, so a
        non-admissible categorical projects to the first admissible value rather than raising --
        :meth:`InterventionSpace.project` is the "make this admissible" path, and a caller who
        wants a refusal should ask :meth:`InterventionSpace.permits` instead.
        """
        if self.contains(value):
            return value
        try:
            return min(self.values, key=lambda v: abs(float(v) - float(value)))
        except (TypeError, ValueError):
            return self.values[0]

    def sample(self, rng: np.random.Generator) -> Any:
        return self.values[int(rng.integers(len(self.values)))]

project(value)

The admissible value closest to value; value itself when already admissible.

Closeness is numeric where the values are numeric and identity otherwise, so a non-admissible categorical projects to the first admissible value rather than raising -- :meth:InterventionSpace.project is the "make this admissible" path, and a caller who wants a refusal should ask :meth:InterventionSpace.permits instead.

Source code in src/causalrl/intervention.py

def project(self, value: Any) -> Any:
    """The admissible value closest to ``value``; ``value`` itself when already admissible.

    Closeness is numeric where the values are numeric and identity otherwise, so a
    non-admissible categorical projects to the first admissible value rather than raising --
    :meth:`InterventionSpace.project` is the "make this admissible" path, and a caller who
    wants a refusal should ask :meth:`InterventionSpace.permits` instead.
    """
    if self.contains(value):
        return value
    try:
        return min(self.values, key=lambda v: abs(float(v) - float(value)))
    except (TypeError, ValueError):
        return self.values[0]

causalrl.intervention.InterventionDomain = Discrete | Continuous module-attribute

What a manipulable variable may be set to: a finite set, or a real interval.

Named in full because :class:causalrl.Domain is already taken by the transport domain -- a source or target population in a selection diagram. The two are unrelated, and one of them had to say which it was.

causalrl.ope.sequential_test.confidence_sequence(values, *, alpha=0.05, value_range, rho=None)

A band on E[value] valid at every sample size simultaneously.

value_range must genuinely contain every observation: it sets the sub-Gaussian variance proxy (high - low) / 2, and a range that is too narrow silently voids the guarantee rather than degrading it. rho tunes where the boundary is tightest and defaults to a value optimised near the sample size actually supplied; validity does not depend on it.

Source code in src/causalrl/ope/sequential_test.py

def confidence_sequence(
    values: Sequence[float],
    *,
    alpha: float = 0.05,
    value_range: tuple[float, float],
    rho: float | None = None,
) -> ConfidenceSequence:
    """A band on ``E[value]`` valid at every sample size simultaneously.

    ``value_range`` must genuinely contain every observation: it sets the sub-Gaussian variance
    proxy ``(high - low) / 2``, and a range that is too narrow silently voids the guarantee rather
    than degrading it. ``rho`` tunes where the boundary is tightest and defaults to a value
    optimised near the sample size actually supplied; validity does not depend on it.
    """
    if not 0.0 < alpha < 1.0:
        raise ValueError(f"alpha={alpha} must lie in (0, 1)")
    low, high = float(value_range[0]), float(value_range[1])
    if not low < high:
        raise ValueError(f"value_range={value_range} must satisfy low < high")
    x = np.asarray(values, dtype=np.float64).reshape(-1)
    n = int(x.size)
    if n == 0:
        raise ValueError(
            "confidence_sequence needs at least one observation: a band over nothing is the whole "
            "real line, which is true but not a claim anyone can act on."
        )
    outside = float(np.max(np.abs(x - (low + high) / 2.0))) if n else 0.0
    if outside > (high - low) / 2.0 + 1e-9:
        raise ValueError(
            f"value_range={value_range} does not contain every observation (found one "
            f"{outside - (high - low) / 2.0:.6g} outside). The range is the sub-Gaussian variance "
            "proxy, so a range that is too narrow voids the time-uniform guarantee outright "
            "rather than making the band merely optimistic."
        )
    sigma = (high - low) / 2.0
    # Tuned so the boundary is tightest near the sample size in hand; validity holds for any rho.
    tuned = rho if rho is not None else max(1.0 / max(n, 1), 1e-12)
    mean = float(x.mean())
    radius = _mixture_radius(n, sigma, alpha, tuned)
    return ConfidenceSequence(
        interval=Interval(mean - radius, mean + radius),
        n=n,
        mean=mean,
        alpha=alpha,
        sigma=sigma,
    )

causalrl.ope.sequential_test.ConfidenceSequence dataclass

A time-uniform band on a mean, and the record of what it took to get there.

interval is deliberately NOT reported through Certificate.ci's usual fixed-sample reading. It is a different object with a stronger guarantee -- valid at all sample sizes at once -- and conflating the two would let a reader assume the wrong one.

Source code in src/causalrl/ope/sequential_test.py

@dataclass(frozen=True)
class ConfidenceSequence:
    """A time-uniform band on a mean, and the record of what it took to get there.

    ``interval`` is deliberately NOT reported through ``Certificate.ci``'s usual fixed-sample
    reading. It is a different object with a stronger guarantee -- valid at all sample sizes at
    once -- and conflating the two would let a reader assume the wrong one.
    """

    interval: Interval
    n: int
    mean: float
    alpha: float
    sigma: float

    @property
    def excludes_zero(self) -> bool:
        """Whether the band has separated from zero -- a valid stopping condition, at any time."""
        return self.interval.lower > 0.0 or self.interval.upper < 0.0

    @property
    def width(self) -> float:
        return self.interval.upper - self.interval.lower

excludes_zero property

Whether the band has separated from zero -- a valid stopping condition, at any time.

causalrl.ope.sequential_test.sequential_policy_comparison(dataset, target_actions, *, alpha=0.05, reward_range, weight_cap=100.0)

Compare a target policy to the logging policy with a band valid under repeated looking.

The per-decision score is the clipped importance-weighted contrast (w_i - 1) * r_i with w_i = 1[a_i = pi(s_i)] / e0(a_i | s_i) capped at weight_cap, whose mean is the value improvement V(pi) - V(behaviour). Clipping is what makes the score bounded and hence sub-Gaussian; it also biases the estimate toward the behaviour policy, which is the conservative direction for a "should I ship this" question and is recorded as an assumption rather than left implicit.

Read :attr:SequentialVerdict.stop as the stopping rule: it may be consulted after every new batch of logs without spending any error budget.

Source code in src/causalrl/ope/sequential_test.py

def sequential_policy_comparison(
    dataset: LoggedDecisions[ActionT],
    target_actions: Sequence[ActionT],
    *,
    alpha: float = 0.05,
    reward_range: tuple[float, float],
    weight_cap: float = 100.0,
) -> SequentialVerdict:
    """Compare a target policy to the logging policy with a band valid under repeated looking.

    The per-decision score is the clipped importance-weighted contrast
    ``(w_i - 1) * r_i`` with ``w_i = 1[a_i = pi(s_i)] / e0(a_i | s_i)`` capped at ``weight_cap``,
    whose mean is the value improvement ``V(pi) - V(behaviour)``. Clipping is what makes the score
    bounded and hence sub-Gaussian; it also biases the estimate toward the behaviour policy, which
    is the conservative direction for a "should I ship this" question and is recorded as an
    assumption rather than left implicit.

    Read :attr:`SequentialVerdict.stop` as the stopping rule: it may be consulted after every new
    batch of logs without spending any error budget.
    """
    outcomes = list(dataset.outcomes())
    e0 = list(dataset.logging_propensities())
    matched = dataset.matches(target_actions)
    weights = [
        min(weight_cap, 1.0 / p) if m and p > 0.0 else 0.0 for m, p in zip(matched, e0, strict=True)
    ]
    scores = [(w - 1.0) * r for w, r in zip(weights, outcomes, strict=True)]

    low, high = float(reward_range[0]), float(reward_range[1])
    bound = max(abs(low), abs(high)) * max(weight_cap - 1.0, 1.0)
    sequence = confidence_sequence(scores, alpha=alpha, value_range=(-bound, bound))

    decided = sequence.excludes_zero
    certificate = Certificate(
        claim=(
            f"V(target) - V(behavior) in [{sequence.interval.lower:.4g}, "
            f"{sequence.interval.upper:.4g}] simultaneously at every sample size "
            f"(confidence sequence, alpha={alpha:g}, n={sequence.n})"
        ),
        estimand=EstimandSpec(query="policy_value", target="mean", policy="target"),
        kind=Kind.EMPIRICAL,
        value=None,
        alpha=alpha,
        assumptions=(
            Assumption(
                name="sub-gaussian-scores",
                params={"bound": bound, "weight_cap": weight_cap},
                checkable=True,
                diagnostic={"n": sequence.n},
            ),
            Assumption(name="no-unmeasured-confounding", params={}, checkable=False),
            Assumption(
                name="clipped-importance-weights",
                params={"weight_cap": weight_cap},
                checkable=True,
                diagnostic={"clipped": sum(1 for w in weights if w >= weight_cap)},
            ),
        ),
        method="normal-mixture-confidence-sequence",
        witness=None,
        hedge=(
            None
            if decided
            else Hedge(
                reason=(
                    "undecided: the band still contains zero, so stopping now would report a "
                    "difference the data have not separated from no difference"
                ),
                detail={"n": sequence.n, "width": sequence.width},
            )
        ),
        provenance=Provenance.create(),
    )
    return SequentialVerdict(sequence=sequence, certificate=certificate)

causalrl.ope.sequential_test.SequentialVerdict dataclass

The result of an anytime-valid policy comparison, and whether it is safe to stop.

Source code in src/causalrl/ope/sequential_test.py

@dataclass(frozen=True)
class SequentialVerdict:
    """The result of an anytime-valid policy comparison, and whether it is safe to stop."""

    sequence: ConfidenceSequence
    certificate: Certificate

    @property
    def stop(self) -> bool:
        """Whether the band has separated from zero, so stopping now is a valid decision."""
        return self.sequence.excludes_zero

    @property
    def better(self) -> str | None:
        """``"target"``, ``"behavior"``, or ``None`` while the comparison is undecided."""
        if not self.stop:
            return None
        return "target" if self.sequence.interval.lower > 0.0 else "behavior"

stop property

Whether the band has separated from zero, so stopping now is a valid decision.

better property

"target", "behavior", or None while the comparison is undecided.

causalrl.scm.fidelity.certify_fitted_query(model, data, *, intervention, outcome, tolerance=2.0, atol=1e-08, n_samples=256, seed=0)

Answer E[outcome | do(intervention)] and gate it on the model's fidelity in that regime.

Runs the query, then tests the outcome's mechanism against the rows of data where the intervened variables already take the queried values. A mechanism that mispredicts those rows by more than tolerance observed standard deviations gets a model-fidelity :class:~causalrl.certify.Hedge and a kind downgraded to EMPIRICAL; a regime with no matching rows at all gets a no-factual-support hedge, because a query there is pure extrapolation and nothing in the data speaks to it either way.

The returned certificate always carries the model's answer. Gating changes what the answer is claimed to be, not whether it is reported -- a caller who wants the number anyway can still read it, and one who is deciding on it can read the hedge.

Source code in src/causalrl/scm/fidelity.py

def certify_fitted_query(
    model: StructuralCausalModel,
    data: Mapping[str, np.ndarray],
    *,
    intervention: Mapping[str, float],
    outcome: str,
    tolerance: float = 2.0,
    atol: float = 1e-8,
    n_samples: int = 256,
    seed: int = 0,
) -> Certificate:
    """Answer ``E[outcome | do(intervention)]`` and gate it on the model's fidelity in that regime.

    Runs the query, then tests the outcome's mechanism against the rows of ``data`` where the
    intervened variables already take the queried values. A mechanism that mispredicts those rows
    by more than ``tolerance`` observed standard deviations gets a ``model-fidelity``
    :class:`~causalrl.certify.Hedge` and a ``kind`` downgraded to ``EMPIRICAL``; a regime with no
    matching rows at all gets a ``no-factual-support`` hedge, because a query there is pure
    extrapolation and nothing in the data speaks to it either way.

    The returned certificate always carries the model's answer. Gating changes what the answer is
    *claimed to be*, not whether it is reported -- a caller who wants the number anyway can still
    read it, and one who is deciding on it can read the hedge.
    """
    if not intervention:
        raise ValueError(
            "intervention must name at least one variable: with nothing intervened on there is no "
            "regime to check the model's fidelity in, and the query is an observational one."
        )
    missing = [name for name in (*intervention, outcome) if name not in data]
    if missing:
        raise KeyError(f"data is missing column(s) {sorted(missing)}")

    answered = model.do(dict(intervention)).see(n_samples, seed=seed)
    predicted = float(np.asarray(answered[outcome], dtype=np.float64).mean())

    rows = _matching_rows(data, intervention, atol)
    parents = sorted(model.graph.parents(outcome))
    observed_column = np.asarray(data[outcome], dtype=np.float64).ravel()

    if rows.size == 0:
        report = FidelityReport(
            outcome=outcome, regime=dict(intervention), support=0, tolerance=tolerance
        )
    else:
        rng = np.random.default_rng(seed)
        probe = rows if rows.size <= _MAX_PROBE_ROWS else rng.choice(rows, _MAX_PROBE_ROWS, False)
        errors: list[float] = []
        for row in probe:
            assignment = {p: float(np.asarray(data[p]).ravel()[row]) for p in parents}
            drawn = model.do(assignment).see(n_samples, seed=seed) if assignment else answered
            predicted_row = float(np.asarray(drawn[outcome], dtype=np.float64).mean())
            errors.append(abs(predicted_row - float(observed_column[row])))
        observed = observed_column[rows]
        spread = float(observed.std())
        mean_error = float(np.mean(errors))
        report = FidelityReport(
            outcome=outcome,
            regime=dict(intervention),
            support=int(rows.size),
            probed=len(errors),
            predicted_mean=predicted,
            observed_mean=float(observed.mean()),
            standardised_error=mean_error / spread if spread > 0.0 else mean_error,
            tolerance=tolerance,
        )

    hedge: Hedge | None = None
    if report.support == 0:
        hedge = Hedge(
            reason=(
                "no-factual-support: the logs contain no row in the queried regime, so this answer "
                "is pure extrapolation and the model's fit there is untested"
            ),
            detail={"regime": dict(intervention)},
            downgraded_from="fitted",
        )
    elif not report.trustworthy:
        hedge = Hedge(
            reason=(
                "model-fidelity: the outcome's fitted mechanism mispredicts the factual rows of "
                f"this regime by {report.standardised_error:.2f} observed standard deviations "
                f"(tolerance {tolerance:.2f}), so its interventional answer here is not to be "
                "read as a point estimate"
            ),
            detail={
                "predicted_mean": report.predicted_mean,
                "observed_mean": report.observed_mean,
                "support": report.support,
            },
            downgraded_from="fitted",
        )

    return Certificate(
        claim=(
            f"E[{outcome} | do({', '.join(f'{k}={v}' for k, v in intervention.items())})] "
            f"= {predicted:.6g} from a fitted SCM; {report.summary()}"
        ),
        estimand=EstimandSpec(query="do", target="mean", domains=(outcome,)),
        kind=Kind.IDENTIFIED if hedge is None else Kind.EMPIRICAL,
        value=None,
        alpha=None,
        assumptions=(
            Assumption(
                name="mechanism-fidelity-in-regime",
                params={"tolerance": tolerance},
                checkable=True,
                diagnostic={
                    "standardised_error": report.standardised_error,
                    "support": report.support,
                    "probed": report.probed,
                },
            ),
        ),
        method="fitted-scm-query-with-fidelity-gate",
        witness=None,
        hedge=hedge,
        provenance=Provenance.create(),
    )

causalrl.scm.fidelity.FidelityReport dataclass

How well the outcome's fitted mechanism predicts the regime being queried.

standardised_error is the mean absolute prediction error over the probed rows, divided by the observed spread of the outcome in that regime. Standardising is what makes one tolerance meaningful across outcomes measured in different units; when the observed spread is zero it falls back to the absolute error, since "off by 3" against a constant is still off by 3.

Source code in src/causalrl/scm/fidelity.py

@dataclass(frozen=True)
class FidelityReport:
    """How well the outcome's fitted mechanism predicts the regime being queried.

    ``standardised_error`` is the mean absolute prediction error over the probed rows, divided by
    the observed spread of the outcome in that regime. Standardising is what makes one tolerance
    meaningful across outcomes measured in different units; when the observed spread is zero it
    falls back to the absolute error, since "off by 3" against a constant is still off by 3.
    """

    outcome: str
    regime: dict[str, float]
    support: int = 0
    probed: int = 0
    predicted_mean: float = float("nan")
    observed_mean: float = float("nan")
    standardised_error: float = float("inf")
    tolerance: float = 2.0

    @property
    def trustworthy(self) -> bool:
        """Whether the mechanism predicts this regime well enough to point-answer a query in it."""
        return self.support > 0 and self.standardised_error <= self.tolerance

    def summary(self) -> str:
        if self.support == 0:
            return (
                f"{self.outcome}: NO FACTUAL SUPPORT for regime {self.regime} -- "
                "the query extrapolates entirely"
            )
        return (
            f"{self.outcome}: {'ok' if self.trustworthy else 'FAILS'} "
            f"(standardised error {self.standardised_error:.2f} vs tolerance {self.tolerance:.2f}, "
            f"predicted {self.predicted_mean:.4g} vs observed {self.observed_mean:.4g}, "
            f"support {self.support} rows)"
        )

trustworthy property

Whether the mechanism predicts this regime well enough to point-answer a query in it.

causalrl.data.builder.TrajectoryLogBuilder

Accumulate rows for a :class:TrajectoryLog across many calls instead of just one.

The counterpart to :meth:TrajectoryLog.from_rows for a producer that does not have all its rows up front: push them in as they happen and call :meth:freeze whenever a :class:TrajectoryLog is needed -- as many times as needed, including more than once per builder. :meth:push is the primitive; :meth:episode is a context manager built on top of it that fixes episode_id for every push inside it.

Source code in src/causalrl/data/builder.py

class TrajectoryLogBuilder:
    """Accumulate rows for a :class:`TrajectoryLog` across many calls instead of just one.

    The counterpart to :meth:`TrajectoryLog.from_rows` for a producer that does not have all its
    rows up front: push them in as they happen and call :meth:`freeze` whenever a
    :class:`TrajectoryLog` is needed -- as many times as needed, including more than once per
    builder. :meth:`push` is the primitive; :meth:`episode` is a context manager built on top of it
    that fixes ``episode_id`` for every push inside it.
    """

    def __init__(self, metadata: Mapping[str, Any] | None = None) -> None:
        """Start an empty builder. ``metadata`` is attached to every log :meth:`freeze` returns."""
        self._cols: dict[str, list[Any]] = {
            "entity_id": [],
            "episode_id": [],
            "t": [],
            "kind": [],
            "name": [],
            "value": [],
            "regime": [],
            "observed": [],
        }
        self._metadata: dict[str, Any] = dict(metadata or {})

    def __len__(self) -> int:
        """Rows pushed so far, across every episode, regardless of whether frozen yet."""
        return len(self._cols["entity_id"])

    @property
    def metadata(self) -> dict[str, Any]:
        """A copy of the log-level metadata :meth:`freeze` will attach to the next snapshot."""
        return dict(self._metadata)

    def push(
        self,
        *,
        episode_id: int,
        kind: str,
        name: str,
        value: Any,
        t: int,
        entity_id: int = 0,
        regime: str = "observed",
        observed: bool = True,
    ) -> None:
        """Append one row.

        Prefer :meth:`episode` when stepping through one episode's rows in order -- it saves
        repeating ``episode_id`` at every call. Call this directly when a caller already tracks
        its own episode id per row (replaying an existing sequence of row dicts into a builder,
        for instance) and a ``with`` block would only add ceremony.

        ``entity_id`` defaults to 0, the single-trajectory case this class exists for (one
        controller, one robot, one running experiment); ``regime``/``observed`` default exactly as
        :meth:`TrajectoryLog.from_rows` defaults them, so the two stay interchangeable. Values are
        stored as given -- push does no coercion or validation of its own -- and only converted to
        the schema's NumPy dtypes inside :meth:`freeze`, the same place ``from_rows`` does its own
        conversion, so a malformed value is reported the same way it would be for any other caller
        of :class:`TrajectoryLog`, not through a second, builder-specific error path.
        """
        self._cols["entity_id"].append(entity_id)
        self._cols["episode_id"].append(episode_id)
        self._cols["t"].append(t)
        self._cols["kind"].append(kind)
        self._cols["name"].append(name)
        self._cols["value"].append(value)
        self._cols["regime"].append(regime)
        self._cols["observed"].append(observed)

    @contextmanager
    def episode(self, episode_id: int) -> Generator[EpisodeWriter]:
        """Open a push scope that stamps every row inside it with ``episode_id``.

        Yields an :class:`EpisodeWriter` bound to ``episode_id``; call :meth:`~EpisodeWriter.push`
        on it instead of repeating the id at every call. The handle stops accepting pushes the
        moment this block exits, including when the body raises: the ``finally`` that closes it
        runs either way, so an exception partway through an episode still leaves the builder
        holding exactly the rows pushed before the failure, ready for more pushes or a
        :meth:`freeze` -- the exception itself propagates to the caller unchanged, nothing here
        swallows it.
        """
        writer = EpisodeWriter(self, episode_id)
        try:
            yield writer
        finally:
            writer.close()

    def freeze(self) -> TrajectoryLog:
        """Materialise everything pushed so far into a genuine :class:`TrajectoryLog`.

        Repeatable, not a one-shot terminal operation: call it as often as you like, including
        after pushing more rows following an earlier call. That is the deliberate choice between
        the two honest options for what "appending after freeze" should mean -- refuse it
        outright, or allow it and let freeze be called again -- because a live producer's entire
        reason for existing is that it keeps running after someone has inspected a snapshot, and
        the streaming certificate this class exists to feed
        (:func:`~causalrl.ope.ipw.stream_policy_value`) is exactly the kind of thing wanted
        mid-run, not only once at the very end. Refusing further pushes after the first freeze
        would force a caller to either stop the producer just to get a certificate, or maintain a
        fresh builder per snapshot and stitch the results back together -- both defeat the point
        of having an incremental builder at all.

        Safe to call repeatedly precisely because it is not free of a copy: :class:`TrajectoryLog`
        converts every column list into a new NumPy array in its constructor, so a log returned by
        an earlier ``freeze()`` shares no storage with this builder and cannot be changed by rows
        pushed after it was returned -- each snapshot is a real, independent log, not a view.
        """
        return TrajectoryLog(self._cols, self._metadata)

metadata property

A copy of the log-level metadata :meth:freeze will attach to the next snapshot.

__init__(metadata=None)

Start an empty builder. metadata is attached to every log :meth:freeze returns.

Source code in src/causalrl/data/builder.py

def __init__(self, metadata: Mapping[str, Any] | None = None) -> None:
    """Start an empty builder. ``metadata`` is attached to every log :meth:`freeze` returns."""
    self._cols: dict[str, list[Any]] = {
        "entity_id": [],
        "episode_id": [],
        "t": [],
        "kind": [],
        "name": [],
        "value": [],
        "regime": [],
        "observed": [],
    }
    self._metadata: dict[str, Any] = dict(metadata or {})

__len__()

Rows pushed so far, across every episode, regardless of whether frozen yet.

Source code in src/causalrl/data/builder.py

def __len__(self) -> int:
    """Rows pushed so far, across every episode, regardless of whether frozen yet."""
    return len(self._cols["entity_id"])

push(*, episode_id, kind, name, value, t, entity_id=0, regime='observed', observed=True)

Append one row.

Prefer :meth:episode when stepping through one episode's rows in order -- it saves repeating episode_id at every call. Call this directly when a caller already tracks its own episode id per row (replaying an existing sequence of row dicts into a builder, for instance) and a with block would only add ceremony.

entity_id defaults to 0, the single-trajectory case this class exists for (one controller, one robot, one running experiment); regime/observed default exactly as :meth:TrajectoryLog.from_rows defaults them, so the two stay interchangeable. Values are stored as given -- push does no coercion or validation of its own -- and only converted to the schema's NumPy dtypes inside :meth:freeze, the same place from_rows does its own conversion, so a malformed value is reported the same way it would be for any other caller of :class:TrajectoryLog, not through a second, builder-specific error path.

Source code in src/causalrl/data/builder.py

def push(
    self,
    *,
    episode_id: int,
    kind: str,
    name: str,
    value: Any,
    t: int,
    entity_id: int = 0,
    regime: str = "observed",
    observed: bool = True,
) -> None:
    """Append one row.

    Prefer :meth:`episode` when stepping through one episode's rows in order -- it saves
    repeating ``episode_id`` at every call. Call this directly when a caller already tracks
    its own episode id per row (replaying an existing sequence of row dicts into a builder,
    for instance) and a ``with`` block would only add ceremony.

    ``entity_id`` defaults to 0, the single-trajectory case this class exists for (one
    controller, one robot, one running experiment); ``regime``/``observed`` default exactly as
    :meth:`TrajectoryLog.from_rows` defaults them, so the two stay interchangeable. Values are
    stored as given -- push does no coercion or validation of its own -- and only converted to
    the schema's NumPy dtypes inside :meth:`freeze`, the same place ``from_rows`` does its own
    conversion, so a malformed value is reported the same way it would be for any other caller
    of :class:`TrajectoryLog`, not through a second, builder-specific error path.
    """
    self._cols["entity_id"].append(entity_id)
    self._cols["episode_id"].append(episode_id)
    self._cols["t"].append(t)
    self._cols["kind"].append(kind)
    self._cols["name"].append(name)
    self._cols["value"].append(value)
    self._cols["regime"].append(regime)
    self._cols["observed"].append(observed)

episode(episode_id)

Open a push scope that stamps every row inside it with episode_id.

Yields an :class:EpisodeWriter bound to episode_id; call :meth:~EpisodeWriter.push on it instead of repeating the id at every call. The handle stops accepting pushes the moment this block exits, including when the body raises: the finally that closes it runs either way, so an exception partway through an episode still leaves the builder holding exactly the rows pushed before the failure, ready for more pushes or a :meth:freeze -- the exception itself propagates to the caller unchanged, nothing here swallows it.

Source code in src/causalrl/data/builder.py

@contextmanager
def episode(self, episode_id: int) -> Generator[EpisodeWriter]:
    """Open a push scope that stamps every row inside it with ``episode_id``.

    Yields an :class:`EpisodeWriter` bound to ``episode_id``; call :meth:`~EpisodeWriter.push`
    on it instead of repeating the id at every call. The handle stops accepting pushes the
    moment this block exits, including when the body raises: the ``finally`` that closes it
    runs either way, so an exception partway through an episode still leaves the builder
    holding exactly the rows pushed before the failure, ready for more pushes or a
    :meth:`freeze` -- the exception itself propagates to the caller unchanged, nothing here
    swallows it.
    """
    writer = EpisodeWriter(self, episode_id)
    try:
        yield writer
    finally:
        writer.close()

freeze()

Materialise everything pushed so far into a genuine :class:TrajectoryLog.

Repeatable, not a one-shot terminal operation: call it as often as you like, including after pushing more rows following an earlier call. That is the deliberate choice between the two honest options for what "appending after freeze" should mean -- refuse it outright, or allow it and let freeze be called again -- because a live producer's entire reason for existing is that it keeps running after someone has inspected a snapshot, and the streaming certificate this class exists to feed (:func:~causalrl.ope.ipw.stream_policy_value) is exactly the kind of thing wanted mid-run, not only once at the very end. Refusing further pushes after the first freeze would force a caller to either stop the producer just to get a certificate, or maintain a fresh builder per snapshot and stitch the results back together -- both defeat the point of having an incremental builder at all.

Safe to call repeatedly precisely because it is not free of a copy: :class:TrajectoryLog converts every column list into a new NumPy array in its constructor, so a log returned by an earlier freeze() shares no storage with this builder and cannot be changed by rows pushed after it was returned -- each snapshot is a real, independent log, not a view.

Source code in src/causalrl/data/builder.py

def freeze(self) -> TrajectoryLog:
    """Materialise everything pushed so far into a genuine :class:`TrajectoryLog`.

    Repeatable, not a one-shot terminal operation: call it as often as you like, including
    after pushing more rows following an earlier call. That is the deliberate choice between
    the two honest options for what "appending after freeze" should mean -- refuse it
    outright, or allow it and let freeze be called again -- because a live producer's entire
    reason for existing is that it keeps running after someone has inspected a snapshot, and
    the streaming certificate this class exists to feed
    (:func:`~causalrl.ope.ipw.stream_policy_value`) is exactly the kind of thing wanted
    mid-run, not only once at the very end. Refusing further pushes after the first freeze
    would force a caller to either stop the producer just to get a certificate, or maintain a
    fresh builder per snapshot and stitch the results back together -- both defeat the point
    of having an incremental builder at all.

    Safe to call repeatedly precisely because it is not free of a copy: :class:`TrajectoryLog`
    converts every column list into a new NumPy array in its constructor, so a log returned by
    an earlier ``freeze()`` shares no storage with this builder and cannot be changed by rows
    pushed after it was returned -- each snapshot is a real, independent log, not a view.
    """
    return TrajectoryLog(self._cols, self._metadata)