Metadata-Version: 2.5
Name: coggrid
Version: 0.2.0
Summary: A stationary POMDP for studying compositional generalization in latent space
Project-URL: Homepage, https://github.com/johnschwarcz/coggrid
Project-URL: Repository, https://github.com/johnschwarcz/coggrid
Project-URL: Issues, https://github.com/johnschwarcz/coggrid/issues
Project-URL: Paper, https://arxiv.org/abs/2603.27134
Author-email: John Schwarcz <Johnschwarcz@gmail.com>
License-Expression: MIT
License-File: LICENSE
Keywords: bayesian-inference,compositional-generalization,disentanglement,gymnasium,ideal-observer,pomdp,reinforcement-learning,toy-model
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: gymnasium>=0.29
Requires-Dist: numpy>=1.24
Provides-Extra: dev
Requires-Dist: matplotlib>=3.7; extra == 'dev'
Requires-Dist: mypy>=1.8; extra == 'dev'
Requires-Dist: pytest>=7.4; extra == 'dev'
Requires-Dist: ruff>=0.5; extra == 'dev'
Requires-Dist: seaborn>=0.13; extra == 'dev'
Provides-Extra: viz
Requires-Dist: matplotlib>=3.7; extra == 'viz'
Requires-Dist: seaborn>=0.13; extra == 'viz'
Description-Content-Type: text/markdown

# coggrid

[![CI](https://github.com/johnschwarcz/coggrid/actions/workflows/ci.yml/badge.svg)](https://github.com/johnschwarcz/coggrid/actions/workflows/ci.yml)
[![PyPI](https://img.shields.io/pypi/v/coggrid)](https://pypi.org/project/coggrid/)
[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue)](https://www.python.org/downloads/)
[![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/johnschwarcz/coggrid/blob/main/LICENSE)
[![arXiv](https://img.shields.io/badge/arXiv-2603.27134-b31b1b)](https://arxiv.org/abs/2603.27134)

A stationary POMDP for studying **compositional generalization in latent space** —
the environment and ideal-observer baselines from
[arXiv:2603.27134](https://arxiv.org/abs/2603.27134).

> **Sequential Bayesian inference is represented as navigating a latent space.**

An example episode:
<p align="center">
  <img src="https://raw.githubusercontent.com/johnschwarcz/coggrid/main/docs/images/episode_animation_1.gif" alt="An episode playing back: joint and naive posteriors, their difference, and the evidence stream" width="100%">
</p>

<p align="center"><em>
Top: the optimal and naive observers' posteriors and their difference. Bottom: A stream of observations.
</em></p>


---

## Contrasting Optimal and Naive Bayes

<p align="center">
  <img src="https://raw.githubusercontent.com/johnschwarcz/coggrid/main/docs/images/factorization_cost.png" alt="Factorization regret predicts the naive observer's failures, and its own confidence does not" width="100%">
</p>

The naive observer is not just noisier — Factorization regret decouples confidence and accuracy.

---

## How the environment works

```python
from coggrid import CogGridConfig, World
from coggrid.viz import animate_interaction_phases, plot_interaction_phases

world = World(CogGridConfig(n_vars=500, n_contexts=2, seed=0))
batch = world.sample_episodes(2000)

plot_interaction_phases(world, batch, episode=0, channels=(0,))
animate_interaction_phases(world, batch, episode=0)
```

Each episode:

1. **Contexts.** `n_contexts` latent variables are drawn from a pool of `n_vars`.
   One is designated the **goal**.
2. **Realizations.** Each "active" variable in the context takes one of `n_realizations` discrete
   values.
3. **Likelihood.** Active variables interact through key/query embeddings which map to an `n_realizations`^`n_contexts` likelihood.
5. **Observations.** `n_steps` i.i.d. observations sampled stochastically from the likelihood / prob. of observing 1 aka 1 - prob. of observing 0.

### What's in a batch

`sample_episodes` returns an `EpisodeBatch`:

| Field | Shape | Meaning |
| --- | --- | --- |
| `cfg` | — | The `CogGridConfig` the episodes were drawn from. |
| `split` | — | `"train"` or `"held_out"` — which variable pool the contexts came from. |
| `ctx_inds` | `(batch_size, n_contexts)` | Indices of the active variables, sorted. |
| `goal_ind` | `(batch_size,)` | Which column of `ctx_inds` is the goal. |
| `ctx_vals` | `(batch_size, n_contexts)` | The realization each active variable took. |
| `goal_value` | `(batch_size,)` | The goal variable's realization — the quantity to infer. |
| `observations` | `(batch_size, n_steps, n_observations)` | The boolean observation stream. |
| `rates` | `(batch_size, n_observations, n_realizations, ...)` | The Bernoulli rate of an observation for every joint realization. |
| `marginal_rates` | `(batch_size, n_observations, n_contexts, n_realizations)` | The same for each variable's realizations after marginalizing over all others. |
| `true_rates` | `(batch_size, n_observations)` | `rates` at `ctx_vals` — the rates that actually generated `observations`. |
| `interactions` | `(batch_size, n_observations, k)` | Raw interaction terms, one per ordered pair of active variables. |

A learning agent should see only `ctx_inds`, `goal_ind` and `observations` (and optionally `interactions` if they are not to be learned).

`batch.select(index)` returns a sub-batch, `batch.save(path)` / `EpisodeBatch.load(path)` round-trip it through a compressed `.npz`.

### Interactions shift the joint likelihood in latent space

The inner product of one variable's key with the other's query sets a *phase* that translates an "XOR"-like pattern.

<p align="center">
  <img src="https://raw.githubusercontent.com/johnschwarcz/coggrid/main/docs/images/interaction_phases.png" alt="Variable embeddings, the inner products they produce, and the standard likelihood pattern those phases translate" width="100%">
</p>


Reading left to right:
1. **Embeddings.** Each latent variable carries a **key** and a **query** vector
   per channel. The angle between one variable's key and the other's query gives a score `z`. 
2. **Phase.**  Each score shifts a standard sinusoid by `−2π · likelihood_freq · z`. 
3. **The standard pattern**, Sinasoids are expanded to a repeating pattern through an outer product.
4. **The selected pattern**  A specific pair of scores defines a specific region of the pattern.

### Impact of rotating the embeddings
Turning the keys sweeps the scores, and every panel to the right follows:

<p align="center">
  <img src="https://raw.githubusercontent.com/johnschwarcz/coggrid/main/docs/images/interaction_phases_animated.gif" alt="Turning a variable's key sweeps its interaction score, sliding the selected region across the standard pattern" width="100%">
</p>

Note: This is meant to demonstrate how embeddings impact the likelihood. Embeddings are never actually rotated.

### What a single observation actually says

```python
from coggrid.viz import plot_evidence_likelihood

plot_evidence_likelihood(batch, episode=0)
```

The agent sees a **vector** of
`n_observations`  at once, and its likelihood is the product of the
per-channel rates.

<p align="center">
  <img src="https://raw.githubusercontent.com/johnschwarcz/coggrid/main/docs/images/evidence_likelihood.png" alt="Per-channel rate tables, and the posterior induced by every possible observation vector" width="100%">
</p>


The top row is the per-channel joint likelihoods.
The grid below is every possible observation vector and the belief update it induces.


## Quickstart

```python
from coggrid import CogGridConfig, World, run_observers

world = World(CogGridConfig(n_vars=500, n_contexts=2, seed=0))
batch = world.sample_episodes(2000)
traces = run_observers(batch)

traces["joint"].final()   # {'accuracy': ..., 'p_correct': ..., 'mse': ...}
traces["naive"].final()   # the same three, for the factorized observer
```

## Install

```bash
pip install coggrid          # core
pip install "coggrid[viz]"   # + matplotlib and seaborn for the figures
```

Python 3.10+. No compiled extensions, no GPU, any OS.

From source, for development:

```bash
git clone https://github.com/johnschwarcz/coggrid
cd coggrid
pip install -e ".[dev]"
pytest
```

### Every knob

`CogGridConfig` holds the whole specification of a world. 

| Field | Default | Meaning |
| --- | --- | --- |
| `n_vars` | 500 | Size of the latent-variable pool episodes draw from. |
| `n_contexts` | 2 | Variables active per episode. The joint table is `n_realizations ** n_contexts` wide, so this drives memory. |
| `n_realizations` | 10 | Discrete values each active variable can take. |
| `n_observations` | 5 | Binary observation channels. |
| `n_steps` | 30 | Observation samples per episode — the horizon. |
| `embedding_dim` | 30 | Length of each key/query vector. Must be ≥ `n_observations`, since they are orthogonalized across channels. |
| `likelihood_temp` | 2.0 | Scales the potentials before the sigmoid. Higher pushes rates towards 0/1, making single observations more informative. |
| `likelihood_freq` | 1.0 | Periods in the value profile. Higher partitions the realization axis more finely. |
| `batch_size` | 1000 | Default batch size for `sample_episodes`. |
| `n_held_out_vars` | `None` | Size of the held-out pool. `None` means `n_vars // 10`. Held-out variables are `range(n_held_out_vars)`. |
| `subsample_vars` | `None` | If set, draw contexts from a random subset of this size within each split. |
| `allow_repeated_vars` | `True` | Whether one episode may activate the same variable twice. |
| `seed` | `None` | Seed for the default RNG. `None` means a fresh world each run. |

### The two baselines

| Observer | Uses | Meaning |
| --- | --- | --- |
| `joint` | `batch.rates` — the full interaction table | Bayes-optimal. The ceiling. |
| `naive` | `batch.marginal_rates` — interactions averaged away | Exactly the joint observer with the interactions removed, and nothing else changed. |

Both run in log space with a cumulative sum rather than a per-step
multiply-and-renormalize, so neither underflows at long horizons (tested to
`n_steps=4000`).

```python
from coggrid import CogGridConfig, World, run_observers
from coggrid import factorization_regret, disentanglement

world = World(CogGridConfig(n_vars=500, n_contexts=2, seed=0))
batch = world.sample_episodes(2000)
traces = run_observers(batch)

factorization_regret(traces["joint"], traces["naive"])   # (batch_size, n_steps)
disentanglement(traces["joint"], traces["naive"], batch) # (batch_size, n_steps)
```

**Factorization regret** — `D_KL(B_joint ‖ B_naive)` 
([§3.1](https://arxiv.org/abs/2603.27134)) — measures the cost of factorizing inference of interacting latent variables.

**Dis-entanglement** ([§B.3](https://arxiv.org/abs/2603.27134)) quantifies the 'history-dependence' of an observation, 
by measuring the divergence between the naive marginal belief update `p(o_t | r)` 
and the optimal marginal belief update `p(o_t | r, o_1:t-1)`. 
It is zero when marginal belief dynamics are Markovian.


### Train / held-out splits

```python
world.sample_episodes(1000, split="train")      # <=1 novel variable, goal always familiar
world.sample_episodes(1000, split="held_out")   # every active variable novel
```

`"train"` is the condition under which a subset of the variables are never the
goal and never co-occur. `"held_out"` is the condition where no variable in the
episode has been a goal, nor co-occurred, during training. This is specifically relevant for training and evaluating networks.

---

## Gymnasium interface

```python
from coggrid import CogGridEnv

env = CogGridEnv(seed=0)
obs, info = env.reset()

for _ in range(env.cfg.n_steps):
    obs, reward, terminated, truncated, info = env.step(my_agent(obs))
```

**obs** — a dict:

| Key | Shape | Meaning |
| --- | --- | --- |
| `observation` | `(n_observations,)` int8 | Fresh binary sample each step. |
| `active_vars` | `(n_contexts,)` int64 | Which latent variables are active. Constant within an episode. |
| `goal_context` | scalar int | Which column of `active_vars` you are scored on. |

`active_vars` holds *indices*, not embeddings — learning an embedding per index
from experience is the point of the task. At test time those indices have never
been seen.

**Action** — `Discrete(n_realizations)`: the agent's current guess (MAP) of the goal
variable's value.

**Reward** — 1.0 for a correct guess, every step (`reward_mode="dense"`) or only
on the last step (`reward_mode="terminal"`).

Episodes end with `terminated=True`, not `truncated=True`.

`coggrid.env.register()` adds `CogGrid-v0` to the Gymnasium registry.

`CogGridEnv` takes:

| Argument | Default | Meaning |
| --- | --- | --- |
| `config` | `None` | The world specification, per the table above. `None` builds a default `CogGridConfig()`. |
| `world` | `None` | Draw from an existing `World` instead of building one, so several environments share embeddings and therefore the same split. Overrides `config`. |
| `split` | `"held_out"` | Which variable pool episodes come from — see below. |
| `reward_mode` | `"dense"` | `"dense"` scores every step, `"terminal"` only the last. |
| `buffer_size` | 256 | Episodes generated per internal refill. Sampling in blocks is much cheaper than one at a time; larger uses more memory. |
| `expose_likelihood` | `False` | Put the joint and marginal rate tables in `info` — what an ideal-observer baseline needs, and what a learning agent must not see. |
| `render_mode` | `None` | `"ansi"` makes `render()` return a text summary of the current step. |
| `seed` | `None` | Seeds the world's embeddings and the episode stream. `None` means fresh episodes each run. |

### Batched rollouts

There is no vector environment. Batching is what `World` already does:

```python
rollout = world.sample_episodes(4096)      # every episode at once
rollout.observations                       # (4096, n_steps, n_observations)
```

`sample_episodes` builds the whole batch with one vectorized einsum over the
joint likelihood, which is the cheap path. A gym wrapper around it would only add
a step-by-step interface on top of data that is already complete.

---

## Visualization

Every plotting function takes data, returns a `Figure`, and never calls `show()`.

```python
from coggrid.viz import summary_figure

figures = summary_figure(batch, traces)   # -> 4 figures, narrowing in scope
```

- **`plot_episode`** — one episode's rate surface beside its belief traces (shown
  above).
- **`plot_performance`** — batch-averaged curves, one column per metric
  (accuracy, `p_correct`, MSE, factorization regret), with individual episode
  trajectories behind each mean.
- **`plot_regret_analysis`** — whether regret explains the gap.
- **`plot_belief_shape`** — the posteriors behind the accuracy numbers: the mean
  final belief with every episode aligned on its true value, and the
  distribution of `P(true value)` across episodes at each timestep.

The last two are omitted when `n_contexts == 1`, where the two observers coincide
and there is nothing to factorize away.

Every panel is also available on its own — `plot_likelihood`, `plot_trial`,
`plot_regret`, `plot_relative_accuracy`, `plot_regret_vs_accuracy`,
`plot_map_agreement`, `plot_belief_profile`, `plot_confidence_density` — and each
takes `ax` (or `axes`, or a `subplot_spec`) so you can compose your own layout.

### Watching an episode

`animate_episode` is the `render()` of this package. In a notebook, make it the
last expression in a cell and it plays — no ffmpeg, no `to_jshtml`, no backend to
configure:

```python
from coggrid.viz import animate_episode

animate_episode(batch, traces)                            # plays inline
animate_episode(batch, traces).save("docs/images/episode")   # or write a GIF
animate_episode(batch, traces).to_html()                  # or a scrub player
```

Panels are laid out in rows — realization grids on top, anything that evolves
over time underneath:

| Row | Panels |
| --- | --- |
| `GRID_PANELS` | `joint`, `naive (factorized)`, and `joint − naive` — where factorizing moves probability mass. Each marks the truth in green and the observer's current mode with a white ring, and the axes name which variable is the **goal** (green) and which is **context** (orange). |
| `TRACE_PANELS` | the evidence stream, revealed up to the current step |

`animate_episode` takes `episode` (which episode to play, default 0), `fps`
(playback speed, default 6), and two arguments that decide what is drawn:

- **`extended=True`** — used for the animation at the top of this page — adds the
  goal variable's marginal belief and two panels relating factorization regret to
  dis-entanglement, one accumulating and one per step.
- **`panels`** — a list of rows, each row a list of panel functions, replacing the
  default layout entirely. This is the extension point below.

**Adding your own panel.** A panel draws its furniture into an axis and returns
an updater called with the timestep, so extending the animation never means
editing the library:

```python
from coggrid.viz import GRID_PANELS, TRACE_PANELS, animate_episode

def p_correct_panel(ax, view):
    (line,) = ax.plot([], [])
    p_correct = view.traces["joint"].p_correct[view.episode]
    ax.set(xlim=(0, view.n_steps - 1), ylim=(0, 1), title="P(true value)")

    def update(t):
        line.set_data(range(t + 1), p_correct[: t + 1])

    return update

animate_episode(batch, traces, panels=[GRID_PANELS, [*TRACE_PANELS, p_correct_panel]])
```

`view` is an `EpisodeView` — the batch, traces, episode index and palette, plus
the derived arrays panels keep needing (`joint_grid`, `naive_grid`,
`observations`, `truth`, `goal_regret`, `disentanglement`).

```bash
python examples/04_figures.py --out docs/images/     # the four static figures
python examples/05_animation.py --out docs/images/   # an episode played back, as a GIF
```

---

## Customizing the generative model

Five stages are swappable. Pass a function to `World`; you never edit the
library.

```python
def codebook_embeddings(cfg, rng):
    """An EmbeddingSource: (cfg, rng) -> (keys, queries)."""
    ...

world = World(cfg, embeddings=codebook_embeddings)
```

| Argument | Signature | Replace it to change |
| --- | --- | --- |
| `embeddings` | `EmbeddingSource` | How latent variables relate to each other |
| `contexts` | `ContextSampler` | Which variables are active; the split structure |
| `likelihood` | `LikelihoodModel` | The *form* of the interaction (e.g. add a 3-way term) |
| `realizations` | `RealizationSampler` | The prior over latent values |
| `observations` | `ObservationModel` | Non-stationary, correlated or continuous observations |

Signatures are at the top of [`src/coggrid/world.py`](https://github.com/johnschwarcz/coggrid/blob/main/src/coggrid/world.py).
Because each returns its outputs rather than mutating shared state, a replacement
can be unit-tested on its own.

---

## Memory

The joint likelihood is `batch_size x n_observations x n_realizations ** n_contexts`.
That last term grows fast, so check before you run:

```python
>>> 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
```

A `ResourceWarning` fires automatically above 2 GiB.

---

## Layout

```
src/coggrid/
├── config.py       CogGridConfig — every tunable, validated, immutable
├── generative.py   the generative model as pure functions
├── world.py        World (config + embeddings), EpisodeBatch, and the five
│                   swappable generative-stage signatures
├── observers.py    ideal-observer baselines, factorization regret,
│                   dis-entanglement
├── env.py          CogGridEnv, the Gymnasium environment
└── viz/            plotting and animation
examples/           five runnable scripts
docs/make_assets.py regenerates the images in this README
tests/              behavioral tests, plus the docstring examples
```

Numerical notes worth knowing:

- **Randomness** is an explicit `numpy.random.Generator` throughout — no global
  `np.random` — so any batch is reproducible from a seed.
- **Inference** accumulates log-likelihoods with a cumulative sum rather than
  multiplying and renormalizing per step, so it does not underflow at long
  horizons.
- **Per-step belief updates** are recovered by differencing *log* beliefs.
  Beliefs routinely fall below `1e-12` once an observer is confident, and
  dividing two such numbers loses most of the significant digits.
- **Rates** are clipped away from 0 and 1 before taking logs, and the sigmoid is
  overflow-safe.

---

## Related repositories

| Repository | Contents |
| --- | --- |
| **this one** (`coggrid`) | The environment and ideal-observer baselines. No networks, no torch. |
| [`CognitiveGridworld`](https://github.com/johnschwarcz/CognitiveGridworld) | The reference implementation from the paper: environment *and* trained networks, together, as published. |

Use **this** repo if you want the task. Use the paper repo to reproduce published
results.

---

## Citation

```bibtex
@article{schwarcz2026factorization,
  title  = {Factorization Regret mediates compositional generalization in latent space},
  author = {Schwarcz, John},
  year   = {2026},
  eprint = {2603.27134},
  archivePrefix = {arXiv}
}
```

## License

MIT — see [LICENSE](https://github.com/johnschwarcz/coggrid/blob/main/LICENSE).
