Genie — Visual Experiments¶

Goal: Reproduce Genie's core idea on toy data — a Latent Action Model that discovers a small discrete set of controllable actions from unlabelled transitions, using VQ.

What we'll explore¶

#	Experiment	Question answered
1	Discovering actions	Can VQ on frame-to-frame change recover actions with no action labels?
2	Action vocabulary	What does the size of the latent action codebook do to the control interface?
3	Frame-by-frame control	Can we steer a world by choosing a latent action each step?
4	Imitation	Can latent actions inferred from an unseen trajectory be replayed to copy it?

---

Linked theory: research/09-genie.md · builds on research/07-tokenized-world-model.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
import warnings
warnings.filterwarnings('ignore')

np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})

# the TRUE (hidden) actions of the world -- the LAM never sees these labels
TRUE_MOVES = np.array([[1, 0], [-1, 0], [0, 1], [0, -1]])   # right, left, up, down
print('Setup complete.')

import numpy as np import matplotlib.pyplot as plt from sklearn.cluster import KMeans import warnings warnings.filterwarnings('ignore') np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) # the TRUE (hidden) actions of the world -- the LAM never sees these labels TRUE_MOVES = np.array([[1, 0], [-1, 0], [0, 1], [0, -1]]) # right, left, up, down print('Setup complete.')

Setup complete.

---

Experiment 1: Discovering actions¶

Question: Genie has no action labels — only video. Its Latent Action Model infers what happened between frames and quantizes it. Can VQ on the frame-to-frame change recover the underlying actions unsupervised?

We generate trajectories where a hidden policy picks one of four moves each step (plus noise). The LAM sees only consecutive positions, computes the change $x_{t+1}-x_t$ (an inverse-dynamics signal), and VQ-clusters it into latent actions.

In [2]:

np.random.seed(1)

n = 4000
hidden_actions = np.random.randint(0, 4, n)              # unknown to the LAM
displacements = TRUE_MOVES[hidden_actions] + 0.18 * np.random.randn(n, 2)

# Latent Action Model: VQ the frame-to-frame change into |A| latent actions
num_actions = 4
lam = KMeans(n_clusters=num_actions, n_init=5, random_state=0).fit(displacements)
latent_action = lam.predict(displacements)
codebook = lam.cluster_centers_
print('learned latent-action codebook (≈ the hidden moves):')
print(np.round(codebook, 2))

np.random.seed(1) n = 4000 hidden_actions = np.random.randint(0, 4, n) # unknown to the LAM displacements = TRUE_MOVES[hidden_actions] + 0.18 * np.random.randn(n, 2) # Latent Action Model: VQ the frame-to-frame change into |A| latent actions num_actions = 4 lam = KMeans(n_clusters=num_actions, n_init=5, random_state=0).fit(displacements) latent_action = lam.predict(displacements) codebook = lam.cluster_centers_ print('learned latent-action codebook (≈ the hidden moves):') print(np.round(codebook, 2))

learned latent-action codebook (≈ the hidden moves):
[[-0.01  1.  ]
 [ 1.    0.  ]
 [-1.   -0.  ]
 [ 0.01 -1.  ]]

In [3]:

fig, ax = plt.subplots(figsize=(6.6, 6.2))
fig.suptitle('Latent Action Model recovers actions with no labels', fontsize=13, fontweight='bold')

ax.scatter(displacements[:, 0], displacements[:, 1], c=latent_action, cmap='tab10', s=10, alpha=0.5)
ax.scatter(codebook[:, 0], codebook[:, 1], c='black', marker='*', s=320,
           edgecolor='white', linewidth=1.2, zorder=5, label='learned latent actions')
ax.scatter(TRUE_MOVES[:, 0], TRUE_MOVES[:, 1], c='crimson', marker='+', s=200,
           linewidth=2.5, zorder=6, label='true (hidden) moves')
ax.set_xlabel('change in x'); ax.set_ylabel('change in y')
ax.set_title('Frame-to-frame change, coloured by latent action')
ax.legend(loc='upper left', fontsize=9); ax.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(6.6, 6.2)) fig.suptitle('Latent Action Model recovers actions with no labels', fontsize=13, fontweight='bold') ax.scatter(displacements[:, 0], displacements[:, 1], c=latent_action, cmap='tab10', s=10, alpha=0.5) ax.scatter(codebook[:, 0], codebook[:, 1], c='black', marker='*', s=320, edgecolor='white', linewidth=1.2, zorder=5, label='learned latent actions') ax.scatter(TRUE_MOVES[:, 0], TRUE_MOVES[:, 1], c='crimson', marker='+', s=200, linewidth=2.5, zorder=6, label='true (hidden) moves') ax.set_xlabel('change in x'); ax.set_ylabel('change in y') ax.set_title('Frame-to-frame change, coloured by latent action') ax.legend(loc='upper left', fontsize=9); ax.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The frame-to-frame changes form four clusters; the LAM's four learned latent actions (black stars) land almost exactly on the true hidden moves (red plusses), and the colours show each change cleanly assigned to one latent action.
2. Why it looks this way: Distinct actions produce distinct, repeatable changes between frames. VQ clusters those changes, so each codeword captures one mode of change — recovering the action set without ever seeing an action label.
3. Key takeaway: This is Genie's central trick: quantizing the transition between frames turns unlabelled video into a discovered, discrete action space.

---

Experiment 2: Action vocabulary¶

Question: Genie uses a small latent action vocabulary (around 8). What happens to the control interface if the codebook is too small or too large?

We fit the LAM with different vocabulary sizes and inspect the resulting latent actions against the four true modes.

In [4]:

np.random.seed(2)

vocab_sizes = [2, 4, 8]
lams = {}
for k in vocab_sizes:
    km = KMeans(n_clusters=k, n_init=5, random_state=0).fit(displacements)
    # within-action variance = how ambiguous each action is
    lab = km.predict(displacements)
    var = np.mean([displacements[lab == j].var(0).sum() for j in range(k) if (lab == j).any()])
    lams[k] = (km.cluster_centers_, var)
    print(f'vocab={k}: mean within-action variance = {var:.3f}')

np.random.seed(2) vocab_sizes = [2, 4, 8] lams = {} for k in vocab_sizes: km = KMeans(n_clusters=k, n_init=5, random_state=0).fit(displacements) # within-action variance = how ambiguous each action is lab = km.predict(displacements) var = np.mean([displacements[lab == j].var(0).sum() for j in range(k) if (lab == j).any()]) lams[k] = (km.cluster_centers_, var) print(f'vocab={k}: mean within-action variance = {var:.3f}')

vocab=2: mean within-action variance = 0.561
vocab=4: mean within-action variance = 0.065
vocab=8: mean within-action variance = 0.043

In [5]:

fig, axes = plt.subplots(1, 3, figsize=(14, 4.7))
fig.suptitle('Too few actions merge moves; matching the modes gives a clean interface', fontsize=13, fontweight='bold')

for ax, k in zip(axes, vocab_sizes):
    centers, var = lams[k]
    ax.scatter(displacements[:, 0], displacements[:, 1], s=6, alpha=0.2, color='lightgray')
    ax.scatter(TRUE_MOVES[:, 0], TRUE_MOVES[:, 1], c='crimson', marker='+', s=160, linewidth=2, zorder=4)
    ax.scatter(centers[:, 0], centers[:, 1], c='seagreen', marker='*', s=240,
               edgecolor='white', linewidth=1, zorder=5)
    ax.set_title(f'vocab |A| = {k}\nwithin-action var = {var:.2f}', fontsize=10.5)
    ax.set_xlabel('Δx'); ax.set_ylabel('Δy')
    ax.set_aspect('equal'); ax.set_xlim(-1.8, 1.8); ax.set_ylim(-1.8, 1.8)

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 3, figsize=(14, 4.7)) fig.suptitle('Too few actions merge moves; matching the modes gives a clean interface', fontsize=13, fontweight='bold') for ax, k in zip(axes, vocab_sizes): centers, var = lams[k] ax.scatter(displacements[:, 0], displacements[:, 1], s=6, alpha=0.2, color='lightgray') ax.scatter(TRUE_MOVES[:, 0], TRUE_MOVES[:, 1], c='crimson', marker='+', s=160, linewidth=2, zorder=4) ax.scatter(centers[:, 0], centers[:, 1], c='seagreen', marker='*', s=240, edgecolor='white', linewidth=1, zorder=5) ax.set_title(f'vocab |A| = {k}\nwithin-action var = {var:.2f}', fontsize=10.5) ax.set_xlabel('Δx'); ax.set_ylabel('Δy') ax.set_aspect('equal'); ax.set_xlim(-1.8, 1.8); ax.set_ylim(-1.8, 1.8) plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: With |A|=2, two latent actions must each cover two true moves, so they sit between clusters and within-action variance is high — ambiguous control. With |A|=4 each latent action lands on one true mode (low variance). With |A|=8 the modes are split into sub-actions that mean nearly the same thing.
2. Why it looks this way: A codebook smaller than the number of real action modes is forced to merge them, so one latent action triggers different effects — useless as a control knob. A codebook larger than needed fragments single actions redundantly.
3. Key takeaway: A small but sufficient vocabulary makes each latent action a consistent, reusable control — the bottleneck is what forces the latent actions to be meaningful, exactly why Genie keeps |A| small.

---

Experiment 3: Frame-by-frame control¶

Question: Once latent actions are discovered, can a user drive the world by choosing one latent action per step — the interactive part of Genie?

We treat each learned latent action as its displacement (the codeword) and roll out a trajectory by feeding a chosen sequence of latent-action ids.

In [6]:

np.random.seed(3)

# label the 4 learned latent actions by their dominant direction for readability
dir_names = {}
for j, c in enumerate(codebook):
    dir_names[j] = ['right', 'left', 'up', 'down'][np.argmin(np.linalg.norm(TRUE_MOVES - c, axis=1))]
name_to_id = {v: k for k, v in dir_names.items()}


def control_rollout(start, action_names):
    pos = np.array(start, float)
    path = [pos.copy()]
    for nm in action_names:
        pos = pos + codebook[name_to_id[nm]]    # apply the learned latent action
        path.append(pos.copy())
    return np.array(path)


# 'play' a little staircase + return
plan = ['right', 'up', 'right', 'up', 'right', 'down', 'down', 'left', 'left']
path = control_rollout([0, 0], plan)
print('latent-action plan:', plan)

np.random.seed(3) # label the 4 learned latent actions by their dominant direction for readability dir_names = {} for j, c in enumerate(codebook): dir_names[j] = ['right', 'left', 'up', 'down'][np.argmin(np.linalg.norm(TRUE_MOVES - c, axis=1))] name_to_id = {v: k for k, v in dir_names.items()} def control_rollout(start, action_names): pos = np.array(start, float) path = [pos.copy()] for nm in action_names: pos = pos + codebook[name_to_id[nm]] # apply the learned latent action path.append(pos.copy()) return np.array(path) # 'play' a little staircase + return plan = ['right', 'up', 'right', 'up', 'right', 'down', 'down', 'left', 'left'] path = control_rollout([0, 0], plan) print('latent-action plan:', plan)

latent-action plan: ['right', 'up', 'right', 'up', 'right', 'down', 'down', 'left', 'left']

In [7]:

fig, ax = plt.subplots(figsize=(6.6, 6))
fig.suptitle('Steering the world with a sequence of latent actions', fontsize=13, fontweight='bold')

ax.plot(path[:, 0], path[:, 1], '-o', color='seagreen', lw=2, ms=7)
for i, (p, nm) in enumerate(zip(path[1:], plan)):
    ax.annotate(nm, p, textcoords='offset points', xytext=(6, 6), fontsize=8)
ax.scatter(*path[0], c='crimson', s=160, marker='*', zorder=5, label='start')
ax.set_xlabel('x'); ax.set_ylabel('y')
ax.set_title('Chosen latent actions trace a controlled path')
ax.legend(fontsize=9); ax.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(6.6, 6)) fig.suptitle('Steering the world with a sequence of latent actions', fontsize=13, fontweight='bold') ax.plot(path[:, 0], path[:, 1], '-o', color='seagreen', lw=2, ms=7) for i, (p, nm) in enumerate(zip(path[1:], plan)): ax.annotate(nm, p, textcoords='offset points', xytext=(6, 6), fontsize=8) ax.scatter(*path[0], c='crimson', s=160, marker='*', zorder=5, label='start') ax.set_xlabel('x'); ax.set_ylabel('y') ax.set_title('Chosen latent actions trace a controlled path') ax.legend(fontsize=9); ax.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Feeding the hand-picked sequence of latent actions produces exactly the intended staircase-and-return path, each step labelled by the latent action that drove it.
2. Why it looks this way: Because each latent action maps to a consistent change (Experiment 2), chaining them gives predictable, controllable motion. The user is effectively 'playing' the world model one latent action at a time.
3. Key takeaway: A small, consistent latent action space is exactly what makes the generated world interactive — the defining feature of Genie.

---

Experiment 4: Imitation¶

Question: Genie's learned latent action space lets agents imitate behaviour from unseen video. Can we infer the latent actions of an unseen trajectory and replay them to reproduce it?

We generate a fresh demo trajectory (its action labels hidden), use the LAM to infer the latent action at each step from the frame changes, then replay those latent actions from a different start position.

In [8]:

np.random.seed(4)

# an unseen demo: draw a diamond-ish loop with hidden moves
demo_hidden = [0, 2, 0, 2, 1, 3, 1, 3]   # right,up,right,up,left,down,left,down
demo = [np.array([2.0, 0.0])]
for a in demo_hidden:
    demo.append(demo[-1] + TRUE_MOVES[a] + 0.15 * np.random.randn(2))
demo = np.array(demo)

# LAM infers latent actions from the demo's frame-to-frame changes (no labels used)
demo_changes = np.diff(demo, axis=0)
inferred = lam.predict(demo_changes)

# replay the inferred latent actions from a new start
replay = [np.array([-1.0, 1.0])]
for j in inferred:
    replay.append(replay[-1] + codebook[j])
replay = np.array(replay)
print('inferred latent-action ids:', inferred.tolist())
print('as directions:', [dir_names[j] for j in inferred])

np.random.seed(4) # an unseen demo: draw a diamond-ish loop with hidden moves demo_hidden = [0, 2, 0, 2, 1, 3, 1, 3] # right,up,right,up,left,down,left,down demo = [np.array([2.0, 0.0])] for a in demo_hidden: demo.append(demo[-1] + TRUE_MOVES[a] + 0.15 * np.random.randn(2)) demo = np.array(demo) # LAM infers latent actions from the demo's frame-to-frame changes (no labels used) demo_changes = np.diff(demo, axis=0) inferred = lam.predict(demo_changes) # replay the inferred latent actions from a new start replay = [np.array([-1.0, 1.0])] for j in inferred: replay.append(replay[-1] + codebook[j]) replay = np.array(replay) print('inferred latent-action ids:', inferred.tolist()) print('as directions:', [dir_names[j] for j in inferred])

inferred latent-action ids: [1, 0, 1, 0, 2, 3, 2, 3]
as directions: ['right', 'up', 'right', 'up', 'left', 'down', 'left', 'down']

In [9]:

fig, ax = plt.subplots(figsize=(7, 6))
fig.suptitle('Imitation: infer latent actions from unseen video, then replay', fontsize=13, fontweight='bold')

ax.plot(demo[:, 0], demo[:, 1], '-o', color='gray', lw=2, ms=7, label='unseen demo (with noise)')
ax.plot(replay[:, 0], replay[:, 1], '--s', color='seagreen', lw=2, ms=6, label='replay from new start')
ax.scatter(*demo[0], c='dimgray', s=140, marker='*', zorder=5)
ax.scatter(*replay[0], c='crimson', s=140, marker='*', zorder=5, label='replay start')
ax.set_xlabel('x'); ax.set_ylabel('y')
ax.set_title('Replayed shape matches the demo (shifted start)')
ax.legend(loc='upper left', fontsize=9); ax.set_aspect('equal')

plt.tight_layout()
plt.show()

fig, ax = plt.subplots(figsize=(7, 6)) fig.suptitle('Imitation: infer latent actions from unseen video, then replay', fontsize=13, fontweight='bold') ax.plot(demo[:, 0], demo[:, 1], '-o', color='gray', lw=2, ms=7, label='unseen demo (with noise)') ax.plot(replay[:, 0], replay[:, 1], '--s', color='seagreen', lw=2, ms=6, label='replay from new start') ax.scatter(*demo[0], c='dimgray', s=140, marker='*', zorder=5) ax.scatter(*replay[0], c='crimson', s=140, marker='*', zorder=5, label='replay start') ax.set_xlabel('x'); ax.set_ylabel('y') ax.set_title('Replayed shape matches the demo (shifted start)') ax.legend(loc='upper left', fontsize=9); ax.set_aspect('equal') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The gray demo trajectory (noisy) and the green replay trace the same shape, just starting from different points. The inferred latent-action ids printed above read off as the demo's actual move sequence.
2. Why it looks this way: The LAM maps each demo frame-change to its nearest latent action, recovering the action sequence behind the demo. Replaying those clean codewords reproduces the behaviour, denoised, from any start.
3. Key takeaway: A shared latent action space lets behaviour be read off unlabelled video and re-executed — the basis of Genie's claim that it opens a path to training generalist agents by imitation.

---

Summary¶

Experiment	Key insight
1. Discovering actions	VQ on frame-to-frame change recovers the hidden action set with no action labels.
2. Action vocabulary	A small-but-sufficient codebook makes each latent action a consistent control; too small merges moves, too large fragments them.
3. Frame-by-frame control	Chaining chosen latent actions steers the world predictably — the interactive core of Genie.
4. Imitation	Latent actions inferred from unseen video can be replayed to reproduce the behaviour.

Connection to the broader theory¶

Genie completes Tier 9 by applying vector quantization twice — once to tokenize observations into a tokenized world model, and once to discover a discrete latent action space from unlabelled video. Together with GAIA-1, it shows discrete latents now cover both state and action: the tokenize → reason → control pipeline is complete, and the framework is ready for the large-scale world models and VLA systems of the next tier.