causalrl 0.99.0__py3-none-any.whl

causalrl/__init__.py

@@ -0,0 +1,274 @@
+# pyright: reportUnsupportedDunderAll=false
+"""causalrl: causal intervention-selection and causal-RL research tools.
+
+The stable public API is loaded lazily so graph algorithms and tabular components can be used
+without installing the optional PyTorch-backed SCM and neural functionality.
+"""
+
+from importlib import import_module as _import_module
+from importlib.metadata import version as _pkg_version
+from typing import cast as _cast
+
+__version__ = _pkg_version("causalrl")
+
+_EXPORTS: dict[str, tuple[str, str]] = {
+    "DOVI": ("causalrl.agents.dovi", "DOVI"),
+    "UCDTR": ("causalrl.agents.offline_online", "UCDTR"),
+    "Agent": ("causalrl.agents.base", "Agent"),
+    "BruteForceInterventionTS": ("causalrl.agents.scbandit", "BruteForceInterventionTS"),
+    "BehavioralCloning": ("causalrl.imitation", "BehavioralCloning"),
+    "BenchmarkEstimate": ("causalrl.eval.benchmark", "BenchmarkEstimate"),
+    "CPDAG": ("causalrl.discovery", "CPDAG"),
+    "PAG": ("causalrl.discovery", "PAG"),
+    "CausalEnv": ("causalrl.envs.base", "CausalEnv"),
+    "CausalGame": ("causalrl.games", "CausalGame"),
+    "CausalGraph": ("causalrl.scm.graph", "CausalGraph"),
+    "CausalGraphError": ("causalrl.exceptions", "CausalGraphError"),
+    "CausalImitator": ("causalrl.imitation", "CausalImitator"),
+    "CausalRLError": ("causalrl.exceptions", "CausalRLError"),
+    "CausalThompsonSampling": ("causalrl.agents.bandits", "CausalThompsonSampling"),
+    "ConfoundedGridworld": ("causalrl.envs.suite.gridworld", "ConfoundedGridworld"),
+    "ConfoundedMDP": ("causalrl.envs.base", "ConfoundedMDP"),
+    "ConfoundedTrajectoryDataset": ("causalrl.data.dataset", "ConfoundedTrajectoryDataset"),
+    "CounterfactualBanditEnv": (
+        "causalrl.envs.suite.counterfactual_bandit",
+        "CounterfactualBanditEnv",
+    ),
+    "CounterfactualOptimalPolicy": (
+        "causalrl.agents.counterfactual",
+        "CounterfactualOptimalPolicy",
+    ),
+    "DTREnv": ("causalrl.envs.suite.dtr", "DTREnv"),
+    "DeepDeconfoundedQ": ("causalrl.agents.deep_deconfounded", "DeepDeconfoundedQ"),
+    "Domain": ("causalrl.identification.id_algorithm", "Domain"),
+    "Estimand": ("causalrl.identification.id_algorithm", "Estimand"),
+    "FixedSetThompsonSampling": ("causalrl.agents.scbandit", "FixedSetThompsonSampling"),
+    "FunctionalMechanism": ("causalrl.scm.mechanisms", "FunctionalMechanism"),
+    "LinearGaussianMechanism": ("causalrl.scm.mechanisms", "LinearGaussianMechanism"),
+    "MABUCEnv": ("causalrl.envs.suite.mabuc", "MABUCEnv"),
+    "Mechanism": ("causalrl.scm.mechanisms", "Mechanism"),
+    "NaiveOffline": ("causalrl.agents.baselines", "NaiveOffline"),
+    "NaivePOMISThompsonSampling": ("causalrl.agents.scbandit", "NaivePOMISThompsonSampling"),
+    "NaiveThompsonSampling": ("causalrl.agents.bandits", "NaiveThompsonSampling"),
+    "NeuralMechanism": ("causalrl.scm.mechanisms", "NeuralMechanism"),
+    "NotIdentifiableError": ("causalrl.exceptions", "NotIdentifiableError"),
+    "OnlineOnlyUCB": ("causalrl.agents.baselines", "OnlineOnlyUCB"),
+    "POMISThompsonSampling": ("causalrl.agents.scbandit", "POMISThompsonSampling"),
+    "PrerequisiteLearner": ("causalrl.curriculum", "PrerequisiteLearner"),
+    "RealizabilityError": ("causalrl.exceptions", "RealizabilityError"),
+    "UnverifiedAssumptionError": ("causalrl.exceptions", "UnverifiedAssumptionError"),
+    "SelectionDiagram": ("causalrl.identification.transport", "SelectionDiagram"),
+    "SequentialDTREnv": ("causalrl.envs.suite.seq_dtr", "SequentialDTREnv"),
+    "SequentialMABUCEnv": ("causalrl.envs.suite.seq_mabuc", "SequentialMABUCEnv"),
+    "StructuralCausalBanditEnv": ("causalrl.envs.suite.scbandit", "StructuralCausalBanditEnv"),
+    "StructuralCausalModel": ("causalrl.scm.scm", "StructuralCausalModel"),
+    "TabularMDP": ("causalrl.shaping", "TabularMDP"),
+    "Transition": ("causalrl.data.dataset", "Transition"),
+    "TransportFormula": ("causalrl.identification.transport", "TransportFormula"),
+    "apply_potential_shaping": ("causalrl.shaping", "apply_potential_shaping"),
+    "backdoor_adjustment_set": ("causalrl.identification.criteria", "backdoor_adjustment_set"),
+    "best_response": ("causalrl.games", "best_response"),
+    "causal_curriculum": ("causalrl.curriculum", "causal_curriculum"),
+    "causal_potential": ("causalrl.shaping", "causal_potential"),
+    "causal_q_bounds": ("causalrl.identification.bounds", "causal_q_bounds"),
+    "ipw_sensitivity_bounds": ("causalrl.identification.bounds", "ipw_sensitivity_bounds"),
+    "manski_bounds": ("causalrl.identification.bounds", "manski_bounds"),
+    "conditional_mutual_information": ("causalrl.discovery", "conditional_mutual_information"),
+    "counterfactual_expectation": (
+        "causalrl.identification.counterfactual",
+        "counterfactual_expectation",
+    ),
+    "cumulative_regret": ("causalrl.eval.metrics", "cumulative_regret"),
+    "curriculum_q_learning": ("causalrl.curriculum", "curriculum_q_learning"),
+    "discover": ("causalrl.discovery", "discover"),
+    "discover_interventional": ("causalrl.discovery", "discover_interventional"),
+    "discover_latent": ("causalrl.discovery", "discover_latent"),
+    "effect_of_treatment_on_treated": (
+        "causalrl.identification.counterfactual",
+        "effect_of_treatment_on_treated",
+    ),
+    "estimate_effect": ("causalrl.identification.id_algorithm", "estimate_effect"),
+    "estimate_effect_with_experiments": (
+        "causalrl.identification.id_algorithm",
+        "estimate_effect_with_experiments",
+    ),
+    "estimate_transport_general": (
+        "causalrl.identification.id_algorithm",
+        "estimate_transport_general",
+    ),
+    "estimate_transported_effect": (
+        "causalrl.identification.id_algorithm",
+        "estimate_transported_effect",
+    ),
+    "finite_horizon_regret": ("causalrl.eval.metrics", "finite_horizon_regret"),
+    "generate_logs": ("causalrl.data.dataset", "generate_logs"),
+    "identify_effect": ("causalrl.identification.id_algorithm", "identify_effect"),
+    "identify_effect_with_experiments": (
+        "causalrl.identification.id_algorithm",
+        "identify_effect_with_experiments",
+    ),
+    "identify_transport": ("causalrl.identification.id_algorithm", "identify_transport"),
+    "identify_transport_general": (
+        "causalrl.identification.id_algorithm",
+        "identify_transport_general",
+    ),
+    "imitation_backdoor_set": ("causalrl.imitation", "imitation_backdoor_set"),
+    "ipw_value": ("causalrl.eval.ope", "ipw_value"),
+    "is_backdoor_admissible": ("causalrl.identification.transport", "is_backdoor_admissible"),
+    "is_gid_identifiable": ("causalrl.identification.id_algorithm", "is_gid_identifiable"),
+    "is_identifiable": ("causalrl.identification.criteria", "is_identifiable"),
+    "is_identifiable_effect": ("causalrl.identification.id_algorithm", "is_identifiable_effect"),
+    "is_imitable": ("causalrl.imitation", "is_imitable"),
+    "is_nash_equilibrium": ("causalrl.games", "is_nash_equilibrium"),
+    "is_transportable": ("causalrl.identification.transport", "is_transportable"),
+    "is_transportable_effect": (
+        "causalrl.identification.id_algorithm",
+        "is_transportable_effect",
+    ),
+    "is_transportable_general": (
+        "causalrl.identification.id_algorithm",
+        "is_transportable_general",
+    ),
+    "is_valid_curriculum": ("causalrl.curriculum", "is_valid_curriculum"),
+    "minimal_intervention_sets": (
+        "causalrl.identification.intervention_sets",
+        "minimal_intervention_sets",
+    ),
+    "mixed_nash_equilibria": ("causalrl.games", "mixed_nash_equilibria"),
+    "pomis": ("causalrl.identification.intervention_sets", "pomis"),
+    "pure_nash_equilibria": ("causalrl.games", "pure_nash_equilibria"),
+    "q_learning": ("causalrl.shaping", "q_learning"),
+    "report_to_dict": ("causalrl.eval.benchmark", "report_to_dict"),
+    "requires_experiment": ("causalrl.identification.intervention_sets", "requires_experiment"),
+    "run_episodes": ("causalrl.eval.harness", "run_episodes"),
+    "run_confounded_chain_benchmark": ("causalrl.eval.benchmark", "run_confounded_chain_benchmark"),
+    "run_frontdoor_benchmark": ("causalrl.eval.benchmark", "run_frontdoor_benchmark"),
+    "transport_estimand": ("causalrl.identification.transport", "transport_estimand"),
+    "transport_formula": ("causalrl.identification.transport", "transport_formula"),
+    "transported_effect": ("causalrl.identification.transport", "transported_effect"),
+    "value_iteration": ("causalrl.shaping", "value_iteration"),
+}
+
+__all__ = [
+    "CPDAG",
+    "DOVI",
+    "PAG",
+    "UCDTR",
+    "Agent",
+    "BehavioralCloning",
+    "BenchmarkEstimate",
+    "BruteForceInterventionTS",
+    "CausalEnv",
+    "CausalGame",
+    "CausalGraph",
+    "CausalGraphError",
+    "CausalImitator",
+    "CausalRLError",
+    "CausalThompsonSampling",
+    "ConfoundedGridworld",
+    "ConfoundedMDP",
+    "ConfoundedTrajectoryDataset",
+    "CounterfactualBanditEnv",
+    "CounterfactualOptimalPolicy",
+    "DTREnv",
+    "DeepDeconfoundedQ",
+    "Domain",
+    "Estimand",
+    "FixedSetThompsonSampling",
+    "FunctionalMechanism",
+    "LinearGaussianMechanism",
+    "MABUCEnv",
+    "Mechanism",
+    "NaiveOffline",
+    "NaivePOMISThompsonSampling",
+    "NaiveThompsonSampling",
+    "NeuralMechanism",
+    "NotIdentifiableError",
+    "OnlineOnlyUCB",
+    "POMISThompsonSampling",
+    "PrerequisiteLearner",
+    "RealizabilityError",
+    "SelectionDiagram",
+    "SequentialDTREnv",
+    "SequentialMABUCEnv",
+    "StructuralCausalBanditEnv",
+    "StructuralCausalModel",
+    "TabularMDP",
+    "Transition",
+    "TransportFormula",
+    "UnverifiedAssumptionError",
+    "__version__",
+    "apply_potential_shaping",
+    "backdoor_adjustment_set",
+    "best_response",
+    "causal_curriculum",
+    "causal_potential",
+    "causal_q_bounds",
+    "conditional_mutual_information",
+    "counterfactual_expectation",
+    "cumulative_regret",
+    "curriculum_q_learning",
+    "discover",
+    "discover_interventional",
+    "discover_latent",
+    "effect_of_treatment_on_treated",
+    "estimate_effect",
+    "estimate_effect_with_experiments",
+    "estimate_transport_general",
+    "estimate_transported_effect",
+    "finite_horizon_regret",
+    "generate_logs",
+    "identify_effect",
+    "identify_effect_with_experiments",
+    "identify_transport",
+    "identify_transport_general",
+    "imitation_backdoor_set",
+    "ipw_sensitivity_bounds",
+    "ipw_value",
+    "is_backdoor_admissible",
+    "is_gid_identifiable",
+    "is_identifiable",
+    "is_identifiable_effect",
+    "is_imitable",
+    "is_nash_equilibrium",
+    "is_transportable",
+    "is_transportable_effect",
+    "is_transportable_general",
+    "is_valid_curriculum",
+    "manski_bounds",
+    "minimal_intervention_sets",
+    "mixed_nash_equilibria",
+    "pomis",
+    "pure_nash_equilibria",
+    "q_learning",
+    "report_to_dict",
+    "requires_experiment",
+    "run_confounded_chain_benchmark",
+    "run_episodes",
+    "run_frontdoor_benchmark",
+    "transport_estimand",
+    "transport_formula",
+    "transported_effect",
+    "value_iteration",
+]
+
+
+def __getattr__(name: str) -> object:
+    target = _EXPORTS.get(name)
+    if target is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attribute = target
+    try:
+        value = _cast(object, getattr(_import_module(module_name), attribute))
+    except ModuleNotFoundError as exc:
+        if exc.name == "torch":
+            raise ImportError(
+                f"{name} requires PyTorch support; install the 'causalrl[torch]' extra"
+            ) from exc
+        raise
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(__all__))

causalrl/_backend/__init__.py

@@ -0,0 +1,15 @@
+"""Thin numerics seam. PyTorch today; an alternate backend can re-implement this module."""
+
+import torch
+
+# torch.Generator/Tensor are public API but not re-exported in torch's type stubs.
+Tensor = torch.Tensor
+Generator = torch.Generator  # type: ignore[reportPrivateImportUsage]
+
+
+def default_generator(seed: int | None = None) -> Generator:
+    """Return a torch.Generator, optionally seeded, for reproducible sampling."""
+    gen = Generator()
+    if seed is not None:
+        gen.manual_seed(seed)
+    return gen

causalrl/agents/__init__.py

@@ -0,0 +1 @@
+"""Causal RL agents."""

causalrl/agents/bandits.py

@@ -0,0 +1,61 @@
+from typing import Any
+
+import numpy as np
+
+from causalrl.agents.base import Agent
+
+
+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
+
+
+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/agents/base.py

@@ -0,0 +1,22 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class Agent(ABC):
+    """Minimal agent interface: choose an action, then learn from the reward."""
+
+    @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

causalrl/agents/baselines.py

@@ -0,0 +1,59 @@
+import math
+from typing import Any
+
+import numpy as np
+
+from causalrl.agents.base import Agent
+from causalrl.data.dataset import ConfoundedTrajectoryDataset
+
+
+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
+
+
+class NaiveOffline(Agent):
+    """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)))
+
+    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
+        """Naive-offline does not learn online (fixed policy from the logs)."""

causalrl/agents/counterfactual.py

@@ -0,0 +1,64 @@
+"""Counterfactual decision-making agents (Layer 3).
+
+A model-based policy that decides by querying the counterfactual reward ``E[Y_{do(a)} | intent]``
+on a known SCM — the Regret Decision Criterion of Bareinboim, Forney & Pearl (NeurIPS 2015).
+Conditioning the choice on the agent's own intent (a proxy for the unobserved confounder) recovers
+the per-intent optimum that the best fixed intervention cannot see.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from causalrl.agents.base import Agent
+from causalrl.identification.counterfactual import regret_decision_table
+
+if TYPE_CHECKING:
+    from causalrl.scm.scm import StructuralCausalModel
+
+__all__ = ["CounterfactualOptimalPolicy"]
+
+
+class CounterfactualOptimalPolicy(Agent):
+    """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])]
+
+    def update(self, observation: dict[str, Any], action: int, reward: float) -> None:
+        """No-op: the SCM is known, so there is nothing to learn online."""

causalrl/agents/deep_deconfounded.py

@@ -0,0 +1,80 @@
+from typing import Any
+
+import numpy as np
+import torch
+
+from causalrl.agents.base import Agent
+from causalrl.agents.primitives import bounds_table
+from causalrl.data.dataset import ConfoundedTrajectoryDataset
+
+
+class _QNet(torch.nn.Module):
+    """One-hot state -> Q-values per action."""
+
+    def __init__(self, n_states: int, n_actions: int) -> None:
+        super().__init__()
+        self.net = torch.nn.Sequential(  # type: ignore[reportUnknownMemberType]
+            torch.nn.Linear(n_states, 64),  # type: ignore[reportPrivateImportUsage]
+            torch.nn.ReLU(),  # type: ignore[reportPrivateImportUsage]
+            torch.nn.Linear(64, n_actions),  # type: ignore[reportPrivateImportUsage]
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)  # type: ignore[reportUnknownMemberType,no-any-return]
+
+
+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]