HeatherDB

Search documentation

Loading the documentation index…

Download
← Research

Rehearsing tomorrow: building a world model out of memory

agentsmemory

The problem

Nudge a glass toward the edge of a table and you already know what happens next. You did not calculate anything. You have pushed enough objects around to have a rough sense of what pushing does, and you run that sense forward in your head before your hand commits.

Software that acts in the world — a robot arm, a game character, a script driving someone else's web app — mostly does not have this. It tries something and finds out. That is fine when the worst case is a wasted second and bad when the worst case is a broken glass. What you want is the ability to ask what would happen if I did this and get a usable answer without doing it.

The usual way to get that is to collect a large pile of before-and-after examples and train a network to predict the "after" from the "before". It works, and it costs you a training run, a machine to do it on, and a repeat of both every time the world changes. This walkthrough builds the same thing out of a memory instead. There is no training step at all: you write down what you saw, and the writing is the learning. At the end you will have measured how many moves ahead the result can be trusted, and what breaks if you skip one particular step.

You need a running HeatherDB and numpy.

How a memory can answer "what would happen if"

Everything here is built out of long lists of numbers. Give each square of a room its own list, and each possible move its own list too. Nothing about a list means anything by itself — what matters is that any two of them are different enough to tell apart.

There are two operations. The first folds two lists into a single list that is a mixture of both — the engine calls this bind. The clever part is that folding is reversible: if you know one of the two ingredients you can take it back out and recover the other, a little blurred. That reverse step is unbind.

So here is a note the memory can keep for one square of the floor: fold the move you made together with the square you ended up on, and add that folded list to whatever the square's note already holds. Do that for all four moves and the square's note now holds four folded facts stacked on top of each other.

To ask "if I go south from here, where do I end up?", fetch that square's note and unfold it using the pattern for south. Out comes something close to the list for the square you would land on, plus a smear of interference from the three other moves stacked in the same note.

Close is not good enough, because you want to ask another question starting from the answer. So there is a third step: hand the blurry result back to the memory and ask which square it most resembles. That snaps it onto an exact square, throwing the smear away. This is the step that makes everything downstream work, and later on you will see exactly what it is worth.

Step 1 — give every square and every move a pattern

The room is a 5×5 grid, 25 squares, four moves, and walls that leave you where you were.

rng = np.random.default_rng(0)
V = np.stack([nrm(rng.standard_normal(D)) for _ in range(N * N)])  # one per square
role = {a: nrm(rng.standard_normal(D)) for a in ACTS}              # one per move
print("mean |cos| between two cell codes:",
      round(float(np.abs(V @ V.T)[~np.eye(25, dtype=bool)].mean()), 4))
step 1: mean |cos| between two cell codes: 0.0273

Random lists of 1024 numbers are almost perfectly unalike — the typical overlap between any two squares is under 3%. That is the property the whole thing rests on: 25 things you can stack together and still pull apart.

Step 2 — wander and write down what happens

M = np.zeros((N * N, D))
s = 0
for _ in range(2000):
    a = ACTS[int(rng.integers(4))]
    s2 = world_step(s, a)
    M[s] += bind(role[a], V[s2])   # fold "this move" with "this landing square",
    s = s2                         # add it to the note for the square you left

There is no epoch, no batch and no gradient here. Each line of experience is one addition. The note is correct-as-far-as-it-goes after the first write and better after the next.

Then load the notes into the engine as two collections. dynamics holds one entry per square: the address is the square, the contents are its stack of folded facts. states is just the list of squares — the thing an answer gets snapped onto.

call("/collections/dynamics/bulk_load",
     {"addresses": [v.tolist() for v in V], "counters": [nrm(m).tolist() for m in M]})
call("/collections/states/bulk_load",
     {"addresses": [v.tolist() for v in V], "counters": [v.tolist() for v in V]})
step 2: cells with a map: 25 of 25
       dynamics: {'num_locations': 25, 'total_writes': 25.0, 'current_eta': 0.01,
                  'avg_write_count': 1.0, 'max_write_count': 1.0}

Use bulk_load rather than write here. The ordinary write path merges things that look similar, which is useful when you are storing observations and wrong when you are storing a lookup table you want kept literal.

Step 3 — ask what one move does

def predict(x, a):
    m = lookup(x)                    # fetch this square's stack of folded facts
    q = nrm(unbind(m, role[a]))      # unfold it with the move -> blurry answer
    y = cleanup(q)                   # snap the blur onto the nearest known square
    return y, float(q @ y)           # ...and how strongly it was recognised

Two details that will cost you an hour otherwise. Reading dynamics needs "strategy": "fast", because there the contents of an entry are deliberately not the same as its address, and the patient read keeps feeding the contents back in as a new question and wanders off. Reading states wants the patient "iterative", because there address and contents are the same, so repetition settles rather than drifts.

step 3: one-step accuracy 100/100, mean fidelity 0.443

All hundred square-and-move combinations answered correctly. The second number is how strongly the memory recognised its own answer — the overlap between the blurry result and the square it snapped to. Call it fidelity. At 0.44 it is far from perfect and it does not need to be perfect: it only has to be closer to the right square than to any of the other twenty-four.

Step 4 — imagine a whole walk

Feed each answer back in as the next question.

plan = ["S"] * 4 + ["E"] * 4
x = V[0]
for a in plan:
    x, _ = predict(x, a)
step 4: plan SSSSEEEE
        imagined [5, 10, 15, 20, 21, 22, 23, 24]
        actual   [5, 10, 15, 20, 21, 22, 23, 24]

Eight moves imagined from the top-left corner, eight squares correct. The memory has never seen this route — it only ever saw single moves, in a random order, while wandering.

Step 5 — how far ahead can you trust it

A number like "eight for eight" means nothing without something to compare it to. So run three arms over the same random walks: the full loop, the same loop with the snapping step removed, and a model that predicts nothing at all and just says you are still where you started.

step 5:  H   cleanup   no cleanup   stay put
         1     100%        100%         30%
         2     100%        100%         20%
         4     100%         70%         20%
         8     100%         28%         10%
        16     100%         20%          6%

Read the middle column. Without the snap, each answer's blur becomes part of the next question, and the next unfold amplifies it: fine at two moves, falling apart by four, worthless by eight. With the snap, the error is thrown away every single move, so sixteen moves cost no more accuracy than one.

The right-hand column is the honest floor. Guessing "nothing changed" is right 30% of the time at one move — the grid has walls, and a move into a wall really does change nothing. Any model you build has to beat that, not zero.

The whole thing

import base64, json, urllib.request
import numpy as np

BASE, D = "http://localhost:6398", 1024
AUTH = base64.b64encode(b"admin:blog-demo").decode()

def call(path, body=None, method="POST"):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(BASE + path, data=data, method=method,
        headers={"Content-Type": "application/json", "Authorization": "Basic " + AUTH})
    with urllib.request.urlopen(req, timeout=120) as r:
        return json.loads(r.read() or b"null")

def nrm(v): return v / (np.linalg.norm(v) + 1e-12)
def vec(op, b): return np.array(call("/vec/" + op, b)["result"])

# bind folds two patterns into one. Unnormalised, so that repeated folds add up
# rather than each one rescaling the note it lands in.
def bind(a, b):
    return vec("bind", {"a": list(map(float, a)), "b": list(map(float, b)),
                        "normalize": False})
# unbind takes a known ingredient back out, leaving the other one, blurred.
def unbind(c, k):
    return vec("unbind", {"a": list(map(float, c)), "b": list(map(float, k))})

# The world: 25 squares, four moves, walls leave you where you were.
N, ACTS = 5, ["N", "S", "E", "W"]
DELTA = {"N": (-1, 0), "S": (1, 0), "E": (0, 1), "W": (0, -1)}
def world_step(s, a):
    r, c = divmod(s, N); dr, dc = DELTA[a]; nr, nc = r + dr, c + dc
    return nr * N + nc if 0 <= nr < N and 0 <= nc < N else s

rng = np.random.default_rng(0)
V = np.stack([nrm(rng.standard_normal(D)) for _ in range(N * N)])
role = {a: nrm(rng.standard_normal(D)) for a in ACTS}
print("step 1: mean |cos| between two cell codes:",
      round(float(np.abs(V @ V.T)[~np.eye(25, dtype=bool)].mean()), 4))

M = np.zeros((N * N, D))
s = 0
for _ in range(2000):                     # wander, and write down what happened
    a = ACTS[int(rng.integers(4))]
    s2 = world_step(s, a)
    M[s] += bind(role[a], V[s2])
    s = s2
print("step 2: cells with a map:", int((np.linalg.norm(M, axis=1) > 0).sum()), "of", N * N)

# bulk_load stores exactly what you hand it. The ordinary write path would merge
# entries that resemble each other, which is wrong for a literal lookup table.
call("/collections/dynamics/bulk_load",
     {"addresses": [v.tolist() for v in V], "counters": [nrm(m).tolist() for m in M]})
call("/collections/states/bulk_load",
     {"addresses": [v.tolist() for v in V], "counters": [v.tolist() for v in V]})
print("       dynamics:", call("/collections/dynamics/stats", method="GET"))

# "fast" here: in `dynamics` an entry's contents differ from its address, so the
# patient read would keep re-asking with the contents and drift away.
def lookup(x):
    return np.asarray(call("/collections/dynamics/read",
                           {"query": nrm(x).tolist(), "strategy": "fast"})["result"])
# "iterative" here: in `states` address and contents match, so repeating settles.
def cleanup(q):
    return nrm(np.asarray(call("/collections/states/read",
                               {"query": q.tolist(), "strategy": "iterative"})["result"]))
def predict(x, a):
    q = nrm(unbind(lookup(x), role[a]))
    y = cleanup(q)
    return y, float(q @ y)            # the answer, and how strongly it was recognised

ok, fid = 0, []
for st in range(N * N):
    for a in ACTS:
        y, f = predict(V[st], a); fid.append(f)
        ok += int(np.argmax(V @ y)) == world_step(st, a)
print(f"step 3: one-step accuracy {ok}/100, mean fidelity {np.mean(fid):.3f}")

plan = ["S"] * 4 + ["E"] * 4
x, truth, imagined = V[0], 0, []
for a in plan:
    x, _ = predict(x, a)              # each answer becomes the next question
    truth = world_step(truth, a)
    imagined.append((int(np.argmax(V @ x)), truth))
print("step 4: plan", "".join(plan))
print("        imagined", [i for i, _ in imagined])
print("        actual  ", [t for _, t in imagined])

def score(H, S=50, seed=7):
    r = np.random.default_rng(seed)
    starts = r.integers(0, N * N, size=S)
    seqs = r.integers(0, 4, size=(S, H))
    clean = raw = stay = 0
    for i in range(S):
        t = int(starts[i]); xc = V[t]; xr = V[t]
        for h in range(H):
            a = ACTS[int(seqs[i, h])]
            t = world_step(t, a)
            xc, _ = predict(xc, a)                       # snap between moves
            xr = nrm(unbind(lookup(xr), role[a]))        # snap removed
        clean += int(np.argmax(V @ xc)) == t
        raw += int(np.argmax(V @ xr)) == t
        stay += int(starts[i]) == t                      # the do-nothing baseline
    return 100 * clean / S, 100 * raw / S, 100 * stay / S

print("step 5:  H   cleanup   no cleanup   stay put")
for H in (1, 2, 4, 8, 16):
    c, r_, st = score(H)
    print(f"        {H:2d}     {c:3.0f}%        {r_:3.0f}%        {st:3.0f}%")

What to change for your own world

Your squares are whatever your situations are — the screen a user is on, the position of a robot arm, the layout of a board. If you can list them and tell them apart, you can give each a pattern. If they are continuous, bucket them, or build the pattern from coordinates so that nearby situations get similar patterns.

Your moves are whatever you can do. Four here; there is nothing special about four.

Uncertain worlds work without changes. If a move sometimes sends you one way and sometimes another, both outcomes land in the same note, weighted by how often each actually happened. Swap the second read for analyze and its activated_locations weights give you the odds instead of a single answer.

Watch how many notes one collection holds. Roughly one entry per 32 numbers of length is the comfortable load, so a collection of 1024-long patterns handles tens of entries cleanly, not thousands. Split by situation rather than piling everything into one collection — the snapping step is where crowding hurts, and a crowded snap is what the middle column of step 5 looks like.

When this works

When you can enumerate the situations. A situation the model has never been in has no note, and it will answer anyway, confidently and wrongly. The fidelity number is what tells you this has happened, and the next post uses it as exactly that.

When you snap between moves. The middle column of step 5 is not a tuning problem. Skip the snap and two moves is your horizon, no matter how big you make the patterns.

When order does not secretly matter. Notes stored per square handle any world, including ones where doing A then B differs from B then A. But the shortcut of learning one pattern per move and applying it everywhere only holds where order does not matter — folding is order-blind, so a world that is not cannot be captured that way.

More depth: Build a world model, world model, rollout and the cleanup loop.