coggrid 0.2.0__py3-none-any.whl

coggrid/__init__.py

@@ -0,0 +1,81 @@
+"""coggrid — a stationary POMDP for studying compositional generalization
+in latent space.
+
+The "grid" is the table of observation rates over the joint realizations of the
+active latent variables.
+
+Quick start
+-----------
+>>> from coggrid import CogGridConfig, World, run_observers
+>>> world = World(CogGridConfig(n_vars=200, n_contexts=2, seed=0))
+>>> batch = world.sample_episodes(256)
+>>> traces = run_observers(batch)
+>>> traces["joint"].final()["accuracy"] > traces["naive"].final()["accuracy"]
+True
+
+Gymnasium interface
+-------------------
+>>> from coggrid import CogGridEnv
+>>> env = CogGridEnv(seed=0)
+>>> obs, info = env.reset()
+>>> obs, reward, terminated, truncated, info = env.step(0)
+
+Layout
+------
+``config``       :class:`CogGridConfig` — every tunable, validated, immutable.
+``generative``   The generative model as pure functions.
+``world``        :class:`World`, :class:`EpisodeBatch`, and the five swappable
+                 generative-stage signatures.
+``observers``    Ideal-observer baselines and metrics.
+``env``          :class:`CogGridEnv`, the Gymnasium environment.
+``viz``          Plotting. Every function returns a figure; none call ``show()``.
+"""
+
+from __future__ import annotations
+
+from .config import CogGridConfig
+from .env import CogGridEnv
+from .observers import (
+    BeliefTrace,
+    disentanglement,
+    factorization_regret,
+    joint_observer,
+    naive_observer,
+    run_observers,
+    score_belief,
+)
+from .world import (
+    ContextSampler,
+    EmbeddingSource,
+    EpisodeBatch,
+    LikelihoodModel,
+    ObservationModel,
+    RealizationSampler,
+    World,
+)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "__version__",
+    # configuration
+    "CogGridConfig",
+    # environment
+    "World",
+    "EpisodeBatch",
+    "CogGridEnv",
+    # observers
+    "BeliefTrace",
+    "joint_observer",
+    "naive_observer",
+    "run_observers",
+    "score_belief",
+    "factorization_regret",
+    "disentanglement",
+    # extension points
+    "EmbeddingSource",
+    "ContextSampler",
+    "LikelihoodModel",
+    "RealizationSampler",
+    "ObservationModel",
+]

coggrid/config.py

@@ -0,0 +1,240 @@
+"""Every tunable of a CogGrid world, in one validated, immutable dataclass.
+
+Nothing in this module touches global state, allocates large arrays, or runs a
+simulation.
+"""
+
+from __future__ import annotations
+
+import warnings
+from dataclasses import dataclass, replace
+from typing import Any
+
+import numpy as np
+
+__all__ = ["CogGridConfig"]
+
+
+@dataclass(frozen=True, slots=True)
+class CogGridConfig:
+    """Immutable description of a CogGrid world.
+
+    The world is a *stationary* POMDP. On each episode:
+
+    1. ``n_contexts`` latent variables are drawn from a pool of ``n_vars``;
+    2. each active variable takes one of ``n_realizations`` discrete values;
+    3. those values jointly determine a Bernoulli rate for each of
+       ``n_observations`` binary observation channels;
+    4. the agent sees ``n_steps`` i.i.d. samples of those channels and must
+       infer the value of the single *goal* variable.
+
+    Because the Bernoulli rates come from **pairwise** interactions between
+    active variables, the joint observation distribution does not factorise over
+    contexts. The gap between an observer that models the interactions and one
+    that assumes independence is the quantity of interest.
+
+    Attributes
+    ----------
+    n_vars:
+        Size of the latent variable pool that contexts are drawn from.
+    n_contexts:
+        Number of simultaneously active latent variables per episode. Note that
+        the joint likelihood table is ``n_realizations ** n_contexts`` wide, so
+        this is the parameter that drives memory use.
+    n_realizations:
+        Number of discrete values each active variable can take.
+    n_observations:
+        Number of binary observation channels.
+    n_steps:
+        Observation samples per episode (the episode horizon).
+    embedding_dim:
+        Dimensionality of the per-variable key/query embeddings. Must be at least
+        ``n_observations``, since embeddings are orthogonalized across channels.
+    likelihood_temp:
+        Scales the interaction potentials before the sigmoid. Larger values push
+        Bernoulli rates towards 0/1, making single observations more informative.
+    likelihood_freq:
+        Number of periods in the sinusoidal value profile. Higher values make
+        the mapping from interaction strength to realization multimodal.
+    batch_size:
+        Default batch size for :meth:`~coggrid.World.sample_episodes`.
+    n_held_out_vars:
+        Size of the held-out ("test") slice of the variable pool. ``None`` means
+        ``n_vars // 10``, resolved to an int during validation. Held-out
+        variables are ``range(n_held_out_vars)``; training variables are the rest.
+    subsample_vars:
+        If set, draw contexts from a random subset of this size within each
+        split. Useful for probing generalization with a fixed small support.
+    allow_repeated_vars:
+        Whether a single episode may activate the same latent variable twice. Set
+        to ``False`` for distinct-variable episodes.
+    seed:
+        Seed for the default RNG. ``None`` means non-reproducible.
+    """
+
+    n_vars: int = 500
+    n_contexts: int = 2
+    n_realizations: int = 10
+    n_observations: int = 5
+    n_steps: int = 30
+    embedding_dim: int = 30
+    likelihood_temp: float = 2.0
+    likelihood_freq: float = 1.0
+    batch_size: int = 1000
+    n_held_out_vars: int | None = None
+    subsample_vars: int | None = None
+    allow_repeated_vars: bool = True
+    seed: int | None = None
+
+    # ------------------------------------------------------------------ setup
+    def __post_init__(self) -> None:
+        positive = {
+            "n_vars": self.n_vars,
+            "n_contexts": self.n_contexts,
+            "n_realizations": self.n_realizations,
+            "n_observations": self.n_observations,
+            "n_steps": self.n_steps,
+            "embedding_dim": self.embedding_dim,
+            "batch_size": self.batch_size,
+        }
+        for name, value in positive.items():
+            if not isinstance(value, (int, np.integer)) or value < 1:
+                raise ValueError(f"{name} must be a positive int, got {value!r}")
+
+        if self.embedding_dim < self.n_observations:
+            raise ValueError(
+                f"embedding_dim ({self.embedding_dim}) must be >= n_observations "
+                f"({self.n_observations}): embeddings are orthogonalized across "
+                "observation channels, which is impossible in a lower-dimensional "
+                "space."
+            )
+        if self.likelihood_temp <= 0:
+            raise ValueError(f"likelihood_temp must be > 0, got {self.likelihood_temp}")
+
+        # Resolve the "auto" default once, so every reader downstream sees an int.
+        held_out = self.n_held_out_vars
+        if held_out is None:
+            held_out = self.n_vars // 10
+        object.__setattr__(self, "n_held_out_vars", held_out)
+        if not 0 <= held_out <= self.n_vars:
+            raise ValueError(
+                f"n_held_out_vars ({held_out}) must be in [0, n_vars={self.n_vars}]"
+            )
+        if held_out == 0:
+            raise ValueError(
+                "n_held_out_vars resolved to 0 — there would be no held-out "
+                "variables to evaluate generalization on. Increase n_vars "
+                "(>= 10) or set n_held_out_vars explicitly."
+            )
+        if held_out == self.n_vars:
+            raise ValueError(
+                "n_held_out_vars equals n_vars — there would be no training "
+                "variables left."
+            )
+
+        if self.subsample_vars is not None:
+            smallest = min(held_out, self.n_vars - held_out)
+            if not 1 <= self.subsample_vars <= smallest:
+                raise ValueError(
+                    f"subsample_vars ({self.subsample_vars}) must be in "
+                    f"[1, {smallest}] (the smaller of the two splits)"
+                )
+
+        if not self.allow_repeated_vars and self.n_contexts > self._split_floor():
+            raise ValueError(
+                f"n_contexts ({self.n_contexts}) exceeds the number of distinct "
+                f"variables available in the smaller split ({self._split_floor()}), "
+                "so distinct-variable episodes are impossible. Either increase "
+                "n_vars / n_held_out_vars or set allow_repeated_vars=True."
+            )
+
+    def _split_floor(self) -> int:
+        """Distinct variables available in the smaller of the two splits."""
+        if self.subsample_vars is not None:
+            return self.subsample_vars
+        return min(self.n_held_out_vars, self.n_vars - self.n_held_out_vars)
+
+    # ------------------------------------------------------- derived quantities
+    @property
+    def n_train_vars(self) -> int:
+        """Number of latent variables reserved for training episodes."""
+        return self.n_vars - self.n_held_out_vars
+
+    @property
+    def n_roll(self) -> int:
+        """Length of the circular value profile (``1 + 2 * n_realizations``).
+
+        The profile is longer than ``n_realizations`` so that interaction
+        strengths can push probability mass "off the end" of the realization
+        axis and wrap around, rather than piling up at the boundary.
+        """
+        return 1 + 2 * self.n_realizations
+
+    @property
+    def realization_shape(self) -> tuple[int, ...]:
+        """Shape of the joint realization axes: ``(n_realizations,) * n_contexts``."""
+        return (self.n_realizations,) * self.n_contexts
+
+    def joint_likelihood_shape(self, batch_size: int | None = None) -> tuple[int, ...]:
+        """Shape of the joint likelihood table for a batch."""
+        n = self.batch_size if batch_size is None else batch_size
+        return (n, self.n_observations, *self.realization_shape)
+
+    def memory_report(self, batch_size: int | None = None) -> str:
+        """Human-readable estimate of the dominant allocation.
+
+        ``n_realizations ** n_contexts`` grows fast; this is the number people
+        need in front of them *before* they wait on an OOM.
+
+        Examples
+        --------
+        >>> print(CogGridConfig(n_contexts=4, n_realizations=20).memory_report(1000))
+        1000 episodes x 5 channels x 20^4 realizations
+          joint likelihood : 6.0 GiB
+          joint belief     : 35.8 GiB
+          peak (approx)    : 41.7 GiB
+        """
+        n = self.batch_size if batch_size is None else batch_size
+        table = int(np.prod(self.joint_likelihood_shape(n), dtype=np.int64)) * 8
+        belief = table * self.n_steps / self.n_observations
+
+        def human(x: float) -> str:
+            for unit in ("B", "KiB", "MiB", "GiB"):
+                if x < 1024:
+                    return f"{x:.1f} {unit}"
+                x /= 1024
+            return f"{x:.1f} TiB"
+
+        return (
+            f"{n} episodes x {self.n_observations} channels x "
+            f"{self.n_realizations}^{self.n_contexts} realizations\n"
+            f"  joint likelihood : {human(table)}\n"
+            f"  joint belief     : {human(belief)}\n"
+            f"  peak (approx)    : {human(table + belief)}"
+        )
+
+    def warn_if_large(
+        self, batch_size: int | None = None, limit_gib: float = 2.0
+    ) -> None:
+        """Emit a ``ResourceWarning`` when the joint table is likely to hurt."""
+        shape = self.joint_likelihood_shape(batch_size)
+        if int(np.prod(shape, dtype=np.int64)) * 8 > limit_gib * 1024**3:
+            warnings.warn(
+                "CogGrid joint likelihood is large:\n"
+                + self.memory_report(batch_size)
+                + "\nReduce batch_size, n_contexts or n_realizations, or sample "
+                "in chunks.",
+                ResourceWarning,
+                stacklevel=3,
+            )
+
+    # ---------------------------------------------------------------- helpers
+    def rng(self, seed: int | np.random.Generator | None = None) -> np.random.Generator:
+        """Build a ``Generator``, preferring an explicit ``seed`` over ``self.seed``."""
+        if isinstance(seed, np.random.Generator):
+            return seed
+        return np.random.default_rng(self.seed if seed is None else seed)
+
+    def replace(self, **changes: Any) -> CogGridConfig:
+        """Return a copy with ``changes`` applied (validation re-runs)."""
+        return replace(self, **changes)

coggrid/env.py

@@ -0,0 +1,287 @@
+"""Gymnasium-style environments.
+
+The environment presents one episode at a time:
+
+* **Observation** — a dict with three keys:
+
+  ``observation``
+      ``(n_observations,)`` int8 binary vector, one fresh sample per step.
+  ``active_vars``
+      ``(n_contexts,)`` int indices of the latent variables active this episode.
+      Constant within an episode. An agent is expected to *learn* an embedding
+      per index, which is where "generalization from experience" comes in — at
+      test time these indices have never been seen.
+  ``goal_context``
+      Which column of ``active_vars`` the agent is scored on. Constant within
+      an episode.
+
+* **Action** — ``Discrete(n_realizations)``: the agent's current guess at the
+  goal variable's value. Actions do **not** influence the observation stream; the
+  world is stationary and the agent is a decoder, not a controller. Reporting a
+  guess every step is what makes the accuracy-versus-time curves in the paper
+  well defined.
+
+* **Reward** — 1.0 for a correct guess. ``reward_mode="dense"`` (default) pays
+  out every step; ``"terminal"`` pays out only on the last step.
+
+Episodes end with ``terminated=True`` rather than ``truncated=True``. The
+horizon is not an externally imposed time limit — it *is* the task, namely "how
+well can you decode from exactly ``n_steps`` samples of evidence" — so there is
+no value left to bootstrap from.
+
+Because the generative model is cheap in batch and expensive one episode at a
+time, :class:`CogGridEnv` samples a buffer of episodes and serves them one by
+one. If you want the batch directly, call
+:meth:`~coggrid.World.sample_episodes` — that is the vectorized path, and no
+environment wrapper is needed for it.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import numpy as np
+from gymnasium import Env
+from gymnasium.spaces import Box, Dict, Discrete, MultiBinary
+
+from .config import CogGridConfig
+from .generative import Split
+from .world import EpisodeBatch, World
+
+__all__ = ["CogGridEnv", "register"]
+
+RewardMode = Literal["dense", "terminal"]
+
+
+def _observation_space(cfg: CogGridConfig) -> Dict:
+    return Dict(
+        {
+            "observation": MultiBinary(cfg.n_observations),
+            "active_vars": Box(
+                low=0,
+                high=cfg.n_vars - 1,
+                shape=(cfg.n_contexts,),
+                dtype=np.int64,
+            ),
+            "goal_context": Discrete(cfg.n_contexts),
+        }
+    )
+
+
+class CogGridEnv(Env):
+    """Single-episode Gymnasium environment.
+
+    Parameters
+    ----------
+    config:
+        A :class:`CogGridConfig`, or ``None`` for the defaults.
+    world:
+        An existing :class:`World` to draw from. Pass one to share embeddings
+        (and therefore the train/held-out split) across several environments.
+        If given, ``config`` is ignored.
+    split:
+        Which variable pool episodes are drawn from — ``"held_out"`` or ``"train"``.
+    reward_mode:
+        ``"dense"`` or ``"terminal"``.
+    buffer_size:
+        How many episodes to generate per internal refill. Larger is faster but
+        uses more memory; see :meth:`CogGridConfig.memory_report`.
+    expose_likelihood:
+        If ``True``, ``info`` carries the full joint and marginal rate tables for
+        the current episode. This is what an ideal-observer baseline needs. Off
+        by default because the arrays are large and an agent should not see them.
+    seed:
+        Seed for the world's embeddings and the episode stream.
+
+    Examples
+    --------
+    >>> env = CogGridEnv(seed=0)
+    >>> obs, info = env.reset()
+    >>> sorted(obs)
+    ['active_vars', 'goal_context', 'observation']
+    >>> total = 0.0
+    >>> for _ in range(env.cfg.n_steps):
+    ...     obs, reward, terminated, truncated, info = env.step(env.action_space.sample())
+    ...     total += reward
+    >>> terminated
+    True
+    """
+
+    metadata = {"render_modes": ["ansi"]}
+
+    def __init__(
+        self,
+        config: CogGridConfig | None = None,
+        *,
+        world: World | None = None,
+        split: Split = "held_out",
+        reward_mode: RewardMode = "dense",
+        buffer_size: int = 256,
+        expose_likelihood: bool = False,
+        seed: int | None = None,
+        render_mode: str | None = None,
+    ) -> None:
+        if reward_mode not in ("dense", "terminal"):
+            raise ValueError(
+                f"reward_mode must be 'dense' or 'terminal', got {reward_mode!r}"
+            )
+        if buffer_size < 1:
+            raise ValueError(f"buffer_size must be >= 1, got {buffer_size}")
+
+        if world is None:
+            cfg = config if config is not None else CogGridConfig()
+            if seed is not None:
+                cfg = cfg.replace(seed=seed)
+            world = World(cfg)
+        self.world = world
+        self.cfg = world.cfg
+
+        self.split: Split = split
+        self.reward_mode: RewardMode = reward_mode
+        self.buffer_size = int(buffer_size)
+        self.expose_likelihood = bool(expose_likelihood)
+        self.render_mode = render_mode
+
+        self.observation_space = _observation_space(self.cfg)
+        self.action_space = Discrete(self.cfg.n_realizations)
+
+        self._stream = np.random.default_rng(
+            self.cfg.seed if seed is None else seed
+        )
+        self._buffer: EpisodeBatch | None = None
+        self._cursor = 0
+        self._episode: EpisodeBatch | None = None
+        self._t = 0
+        self._last_action: int | None = None
+
+    # ------------------------------------------------------------------ buffer
+    def _next_episode(self) -> EpisodeBatch:
+        if self._buffer is None or self._cursor >= len(self._buffer):
+            self._buffer = self.world.sample_episodes(
+                self.buffer_size, split=self.split, rng=self._stream
+            )
+            self._cursor = 0
+        episode = self._buffer.select(self._cursor)
+        self._cursor += 1
+        return episode
+
+    # ---------------------------------------------------------------- gym api
+    def reset(
+        self,
+        *,
+        seed: int | None = None,
+        options: dict[str, Any] | None = None,
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Start a new episode. Returns ``(observation, info)``."""
+        if seed is not None:
+            self._stream = np.random.default_rng(seed)
+            self._buffer = None
+        if options:
+            split = options.get("split")
+            if split is not None:
+                self.split = split
+                self._buffer = None
+
+        self._episode = self._next_episode()
+        self._t = 0
+        self._last_action = None
+        return self._observe(), self._info()
+
+    def step(
+        self, action: int | np.integer
+    ) -> tuple[dict[str, Any], float, bool, bool, dict[str, Any]]:
+        """Submit a guess at the goal variable's value and advance one step."""
+        if self._episode is None:
+            raise RuntimeError("call reset() before step()")
+
+        action = int(action)
+        if not 0 <= action < self.cfg.n_realizations:
+            raise ValueError(
+                f"action {action} outside Discrete({self.cfg.n_realizations})"
+            )
+        self._last_action = action
+
+        correct = action == int(self._episode.goal_value[0])
+        self._t += 1
+        terminated = self._t >= self.cfg.n_steps
+
+        if self.reward_mode == "dense":
+            reward = float(correct)
+        else:
+            reward = float(correct) if terminated else 0.0
+
+        # On the final step there is no further observation to reveal; repeat the
+        # last one so the returned dict always matches observation_space.
+        observation = self._observe(clamp=True)
+        info = self._info()
+        info["correct"] = correct
+        return observation, reward, terminated, False, info
+
+    def render(self) -> str | None:
+        """One-line text summary of the current step."""
+        if self.render_mode != "ansi" or self._episode is None:
+            return None
+        e = self._episode
+        bits = "".join("1" if b else "0" for b in self._current_bits())
+        guess = "-" if self._last_action is None else str(self._last_action)
+        return (
+            f"t={self._t:>3}/{self.cfg.n_steps}  obs={bits}  "
+            f"vars={e.ctx_inds[0].tolist()}  goal_ctx={int(e.goal_ind[0])}  "
+            f"truth={int(e.goal_value[0])}  guess={guess}"
+        )
+
+    # ---------------------------------------------------------------- internals
+    def _current_bits(self, clamp: bool = True) -> np.ndarray:
+        assert self._episode is not None
+        t = min(self._t, self.cfg.n_steps - 1) if clamp else self._t
+        return self._episode.observations[0, t]
+
+    def _observe(self, clamp: bool = True) -> dict[str, Any]:
+        assert self._episode is not None
+        return {
+            "observation": self._current_bits(clamp).astype(np.int8),
+            "active_vars": self._episode.ctx_inds[0].astype(np.int64),
+            "goal_context": int(self._episode.goal_ind[0]),
+        }
+
+    def _info(self) -> dict[str, Any]:
+        assert self._episode is not None
+        e = self._episode
+        info: dict[str, Any] = {
+            "step": self._t,
+            "split": self.split,
+            # Ground truth: for analysis and for scoring, never as agent input.
+            "goal_value": int(e.goal_value[0]),
+            "context_values": e.ctx_vals[0].copy(),
+            "true_rates": e.true_rates[0].copy(),
+        }
+        if self.expose_likelihood:
+            info["rates"] = e.rates[0]
+            info["marginal_rates"] = e.marginal_rates[0]
+        return info
+
+    # ------------------------------------------------------------------ extras
+    @property
+    def episode(self) -> EpisodeBatch | None:
+        """The current episode as a one-element :class:`EpisodeBatch`.
+
+        Handy for running an ideal observer on exactly the episode the agent saw.
+        """
+        return self._episode
+
+
+def register() -> None:
+    """Add ``CogGrid-v0`` to the Gymnasium registry.
+
+    >>> import gymnasium as gym
+    >>> from coggrid.env import register
+    >>> register()
+    >>> env = gym.make("CogGrid-v0")
+    """
+    from gymnasium.envs.registration import register as gym_register
+
+    gym_register(
+        id="CogGrid-v0",
+        entry_point="coggrid.env:CogGridEnv",
+        max_episode_steps=None,  # the env terminates on its own
+    )