daab9b9d14
Train on Apple MPS when available (the free-roll recurrent depth is the cost, so the GPU is what lets a bigger/longer model train); STEPS is now a CLI arg; the served demo walk is longer (16 actions) to showcase the horizon. The integer contract + committed tiny conformance vector are unchanged. This is how the served C48 /world-model model (+1.4 dB) was trained. Co-Authored-By: Claude Opus 4.8 <[email protected]>
298 lines
13 KiB
Python
298 lines
13 KiB
Python
#!/usr/bin/env python3
|
||
"""Train a small action-conditioned INTERACTIVE WORLD MODEL and export it as an
|
||
INTEGER-exact decoder — the conformance-is-a-hash answer to Genie / Oasis / Matrix /
|
||
fal-WMA. Those generate the next frame from an action through a FLOAT network, so the
|
||
generated world is non-reproducible and ephemeral. Here the world model is integer:
|
||
|
||
frame_0 = g_s(z_0); z_{t+1} = clip( T( z_t ‖ action_t ) ); frame_{t+1} = g_s(z_{t+1})
|
||
|
||
with BOTH the transition T and the synthesis g_s exported as int_hyper IntModels
|
||
(WHM1: per-out-channel int12 weights, fixed-point IGDN/leaky, all i64). The action is
|
||
folded in as constant input planes. So given (model, z_0, action trace) every frame is
|
||
byte-identical on every machine — a generative interactive world you can put in a
|
||
receipt. Paired with wai.world.model's signed/replay-equivalent action trace, both the
|
||
trace AND the frames hash-converge.
|
||
|
||
python tools/wai_int_world_export.py <out-dir>
|
||
"""
|
||
import sys, os, json, struct
|
||
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, warnings
|
||
warnings.filterwarnings('ignore')
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
from wai_sota_int_synth_export import (q_weight_pc, q_weight_conv, conv_int, conv_t_raw,
|
||
igdn_int, run_int_net, AFRAC, GF, SHIFT)
|
||
|
||
FRAME, LAT, C, A = 64, 8, 24, 3 # frame px / latent px / channels / action planes
|
||
WORLD, SPEED = 96, 4 # global world image + camera step (px/action)
|
||
NROLL, ROLL = 40, 8 # training rollouts / steps per rollout
|
||
STEPS = 2500 # optimizer steps (CLI-overridable for scale runs)
|
||
ONE = 1 << AFRAC
|
||
ACTIONS = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)] # stay, right, left, down, up
|
||
SLOPE_Q, SLOPE_FRAC = 410, 12 # leaky 0.1 in fixed point (410/4096)
|
||
|
||
|
||
# ---------- procedural world + environment (the "ground truth" dynamics) ----------
|
||
def make_world():
|
||
rng = np.random.RandomState(7)
|
||
lo = torch.tensor(rng.rand(1, 3, WORLD // 8, WORLD // 8).astype(np.float32))
|
||
w = F.interpolate(lo, size=(WORLD, WORLD), mode='bilinear', align_corners=True)[0].numpy()
|
||
yy, xx = np.mgrid[0:WORLD, 0:WORLD]
|
||
w[0] = np.clip(w[0] + 0.25 * np.sin(xx / 6.0), 0, 1) # structure the model must learn
|
||
w[1] = np.clip(w[1] + 0.25 * np.cos(yy / 7.0), 0, 1)
|
||
w[2] = np.clip(w[2] + 0.20 * ((xx // 12 + yy // 12) % 2), 0, 1)
|
||
return w.astype(np.float32)
|
||
|
||
|
||
def crop(world, cx, cy):
|
||
x0 = int(round(cx)) - FRAME // 2
|
||
y0 = int(round(cy)) - FRAME // 2
|
||
x0 = min(max(x0, 0), WORLD - FRAME)
|
||
y0 = min(max(y0, 0), WORLD - FRAME)
|
||
return world[:, y0:y0 + FRAME, x0:x0 + FRAME]
|
||
|
||
|
||
def rollout(world, rng, acts=None):
|
||
cx = cy = WORLD / 2.0
|
||
frames = [crop(world, cx, cy)]
|
||
chosen = []
|
||
for t in range(len(acts) if acts is not None else ROLL):
|
||
a = acts[t] if acts is not None else rng.randint(len(ACTIONS))
|
||
dx, dy = ACTIONS[a]
|
||
cx = min(max(cx + dx * SPEED, FRAME / 2), WORLD - FRAME / 2)
|
||
cy = min(max(cy + dy * SPEED, FRAME / 2), WORLD - FRAME / 2)
|
||
frames.append(crop(world, cx, cy))
|
||
chosen.append(a)
|
||
return np.stack(frames), chosen # frames[ROLL+1,3,64,64], acts[ROLL]
|
||
|
||
|
||
def action_planes(a): # [A,8,8] float at unit scale
|
||
dx, dy = ACTIONS[a]
|
||
p = torch.zeros(A, LAT, LAT)
|
||
p[0] = dx; p[1] = dy; p[2] = 0.0
|
||
return p
|
||
|
||
|
||
# ---------- learned model (float; the encoder/trainer) ----------
|
||
class GDN2d(nn.Module):
|
||
def __init__(self, ch):
|
||
super().__init__()
|
||
self.rg = nn.Parameter(torch.randn(ch, ch) * 0.01)
|
||
self.rb = nn.Parameter(torch.zeros(ch))
|
||
|
||
def gb(self):
|
||
return F.softplus(self.rg), F.softplus(self.rb + 1.0)
|
||
|
||
def forward(self, x): # inverse GDN: x * sqrt(beta + gamma·x²)
|
||
g, b = self.gb()
|
||
norm = b[None, :, None, None] + torch.einsum('ij,njhw->nihw', g, x * x)
|
||
return x * torch.sqrt(norm)
|
||
|
||
|
||
class World(nn.Module):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.E = nn.Sequential(nn.Conv2d(3, C, 4, 2, 1), nn.ReLU(),
|
||
nn.Conv2d(C, C, 4, 2, 1), nn.ReLU(),
|
||
nn.Conv2d(C, C, 4, 2, 1))
|
||
self.Tc = nn.ModuleList([nn.Conv2d(C + A, C, 3, 1, 1), nn.Conv2d(C, C, 3, 1, 1)])
|
||
self.G = nn.ModuleList([nn.ConvTranspose2d(C, C, 4, 2, 1), GDN2d(C),
|
||
nn.ConvTranspose2d(C, C, 4, 2, 1), GDN2d(C),
|
||
nn.ConvTranspose2d(C, 3, 4, 2, 1)])
|
||
|
||
def trans(self, z, a):
|
||
ap = torch.stack([action_planes(ai) for ai in a]).to(z.device)
|
||
x = F.leaky_relu(self.Tc[0](torch.cat([z, ap], 1)), 0.1)
|
||
return torch.clamp(self.Tc[1](x), -1, 1)
|
||
|
||
def synth(self, z):
|
||
x = z
|
||
for m in self.G:
|
||
x = m(x)
|
||
return x
|
||
|
||
def forward(self, frames, acts): # frames[B,T+1,3,64,64], acts[B][T]
|
||
z = self.E(frames[:, 0])
|
||
preds = [self.synth(z)]
|
||
for t in range(ROLL):
|
||
z = self.trans(z, [a[t] for a in acts])
|
||
preds.append(self.synth(z))
|
||
return torch.stack(preds, 1)
|
||
|
||
|
||
# ---------- WHM1 writer (matches int_hyper_synth::from_bytes) ----------
|
||
def whm1(layers):
|
||
out = bytearray(b'WHM1') + struct.pack('<IIII', AFRAC, GF, SHIFT, len(layers))
|
||
for L in layers:
|
||
if L['type'] == 'convT':
|
||
out += struct.pack('<BIIIIII', 0, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'], L['opad'])
|
||
out += np.array(L['weight'], dtype='<i2').tobytes()
|
||
out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
|
||
elif L['type'] == 'conv':
|
||
out += struct.pack('<BIIIII', 1, L['in_ch'], L['out_ch'], L['k'], L['stride'], L['pad'])
|
||
out += np.array(L['weight'], dtype='<i2').tobytes()
|
||
out += np.array(L['mult'], dtype='<i8').tobytes() + np.array(L['bias'], dtype='<i8').tobytes()
|
||
elif L['type'] == 'igdn':
|
||
out += struct.pack('<BI', 2, L['c'])
|
||
out += np.array(L['gamma'], dtype='<i4').tobytes() + np.array(L['beta'], dtype='<i8').tobytes()
|
||
elif L['type'] == 'leaky':
|
||
out += struct.pack('<B', 4) + np.array([L['slope_q']], dtype='<i8').tobytes() + struct.pack('<I', L['slope_frac'])
|
||
return bytes(out)
|
||
|
||
|
||
def ptq_conv(layer): # nn.Conv2d → int dict
|
||
wi, sc = q_weight_conv(layer.weight.detach().numpy())
|
||
return {'type': 'conv', 'in_ch': wi.shape[1], 'out_ch': wi.shape[0], 'k': wi.shape[2],
|
||
'stride': layer.stride[0], 'pad': layer.padding[0], 'weight': wi.reshape(-1).tolist(),
|
||
'mult': np.round(sc * (2 ** SHIFT)).astype(np.int64).tolist(),
|
||
'bias': np.round(layer.bias.detach().numpy() * ONE).astype(np.int64).tolist()}
|
||
|
||
|
||
def ptq_convT(layer):
|
||
wi, sc = q_weight_pc(layer.weight.detach().numpy())
|
||
return {'type': 'convT', 'in_ch': wi.shape[0], 'out_ch': wi.shape[1], 'k': wi.shape[2],
|
||
'stride': layer.stride[0], 'pad': layer.padding[0], 'opad': layer.output_padding[0],
|
||
'weight': wi.reshape(-1).tolist(),
|
||
'mult': np.round(sc * (2 ** SHIFT)).astype(np.int64).tolist(),
|
||
'bias': np.round(layer.bias.detach().numpy() * ONE).astype(np.int64).tolist()}
|
||
|
||
|
||
def ptq_igdn(layer):
|
||
with torch.no_grad():
|
||
g, b = layer.gb()
|
||
return {'type': 'igdn', 'c': b.shape[0],
|
||
'gamma': np.round(g.numpy() * (2 ** GF)).astype(np.int64).reshape(-1).tolist(),
|
||
'beta': np.round(b.numpy() * (2 ** (2 * AFRAC))).astype(np.int64).tolist()}
|
||
|
||
|
||
# ---------- integer reference decode (mirrors the Rust int_world loop) ----------
|
||
def planes_int(a):
|
||
dx, dy = ACTIONS[a]
|
||
p = np.zeros((A, LAT, LAT), np.int64)
|
||
p[0] = dx * ONE; p[1] = dy * ONE
|
||
return p
|
||
|
||
|
||
def gs_decode(gs, z): # z[C,8,8]@2^afrac → frame[64,64,3] u8
|
||
x = run_int_net(gs, z)
|
||
half = 1 << (AFRAC - 1)
|
||
x = np.clip(x, 0, ONE)
|
||
pix = ((x * 255 + half) >> AFRAC).clip(0, 255).astype(np.uint8) # [3,64,64]
|
||
return np.ascontiguousarray(pix.transpose(1, 2, 0)) # HWC
|
||
|
||
|
||
def fnv1a(buf):
|
||
h = 14695981039346656037
|
||
for b in buf:
|
||
h = ((h ^ b) * 1099511628211) & 0xFFFFFFFFFFFFFFFF
|
||
return h
|
||
|
||
|
||
def decode_world(T, gs, z0, acts):
|
||
z = z0.copy()
|
||
frames = [gs_decode(gs, z)]
|
||
for a in acts:
|
||
inp = np.concatenate([z, planes_int(a)], 0) # [C+A,8,8]
|
||
z = run_int_net(T, inp)
|
||
z = np.clip(z, -ONE, ONE)
|
||
frames.append(gs_decode(gs, z))
|
||
return frames
|
||
|
||
|
||
def main(outdir):
|
||
os.makedirs(outdir, exist_ok=True)
|
||
torch.manual_seed(0)
|
||
world = make_world()
|
||
rng = np.random.RandomState(123)
|
||
# Half random walks, half OUT-AND-BACK (go some direction, then reverse) so the
|
||
# model is trained toward long-horizon consistency: reversing a path should
|
||
# return where you came from.
|
||
data = []
|
||
for i in range(NROLL):
|
||
if i % 2 == 0:
|
||
data.append(rollout(world, rng))
|
||
else:
|
||
half = ROLL // 2
|
||
d = rng.randint(1, len(ACTIONS)) # a non-stay direction
|
||
rev = [2, 1, 4, 3][d - 1] # its reverse (x↔, y↔)
|
||
data.append(rollout(world, rng, acts=[d] * half + [rev] * (ROLL - half)))
|
||
frames = torch.tensor(np.stack([d[0] for d in data]))
|
||
acts = [d[1] for d in data]
|
||
|
||
# GPU when available (Apple MPS) — the free-roll recurrent backprop DEPTH is the
|
||
# training cost, so the GPU is what unblocks a longer horizon (bigger ROLL).
|
||
dev = 'mps' if torch.backends.mps.is_available() else ('cuda' if torch.cuda.is_available() else 'cpu')
|
||
print(f' training on {dev} (ROLL={ROLL}, C={C}, steps={STEPS})')
|
||
m = World().to(dev)
|
||
frames = frames.to(dev)
|
||
opt = torch.optim.Adam(m.parameters(), 1.5e-3)
|
||
for step in range(STEPS):
|
||
pred = m(frames, acts)
|
||
loss = F.mse_loss(pred, frames)
|
||
opt.zero_grad(); loss.backward(); opt.step()
|
||
m.eval()
|
||
with torch.no_grad():
|
||
snr_f = 10 * np.log10((frames ** 2).sum().item() / ((m(frames, acts) - frames) ** 2).sum().item() + 1e-12)
|
||
m.cpu(); frames = frames.cpu() # back to CPU for PTQ + the integer reference decode
|
||
|
||
# PTQ both nets → WHM1
|
||
T_layers = [ptq_conv(m.Tc[0]), {'type': 'leaky', 'slope_q': SLOPE_Q, 'slope_frac': SLOPE_FRAC}, ptq_conv(m.Tc[1])]
|
||
gs_layers = [ptq_convT(m.G[0]), ptq_igdn(m.G[1]), ptq_convT(m.G[2]), ptq_igdn(m.G[3]), ptq_convT(m.G[4])]
|
||
T_model = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': T_layers}
|
||
gs_model = {'afrac': AFRAC, 'gf': GF, 'shift': SHIFT, 'layers': gs_layers}
|
||
|
||
# a longer demo walk (showcases the horizon): wander out and partway back
|
||
demo_acts = [1, 1, 1, 3, 3, 1, 4, 2, 2, 2, 3, 4, 1, 0, 2, 4]
|
||
df, _ = rollout(world, None, acts=demo_acts)
|
||
with torch.no_grad():
|
||
z0 = m.E(torch.tensor(df[None, 0]))[0].numpy()
|
||
z0_int = np.round(z0 * ONE).astype(np.int64)
|
||
|
||
frames_int = decode_world(T_model, gs_model, z0_int, demo_acts)
|
||
# quality vs ground-truth crops
|
||
gt = (df.transpose(0, 2, 3, 1) * 255).astype(np.float64)
|
||
pr = np.stack(frames_int).astype(np.float64)
|
||
snr_i = 10 * np.log10((gt ** 2).sum() / ((pr - gt) ** 2).sum() + 1e-12)
|
||
|
||
# LONG-HORIZON CONSISTENCY probe: walk out and back the same distance; does the
|
||
# integer world return to where it started? (reported honestly — consistency is
|
||
# the hard problem for world models; we measure, not assume.)
|
||
n = max(2, ROLL // 2)
|
||
back = [1] * n + [2] * n # +x n, then −x n
|
||
bf = decode_world(T_model, gs_model, z0_int, back)
|
||
ret_match = bool((bf[0] == bf[-1]).all())
|
||
ret_snr = 10 * np.log10((bf[0].astype(np.float64) ** 2).sum()
|
||
/ ((bf[-1].astype(np.float64) - bf[0].astype(np.float64)) ** 2).sum() + 1e-12)
|
||
print(f' walk-and-return ({n}+{n}): start==end byte-exact? {ret_match} (return SNR {ret_snr:.1f} dB)')
|
||
|
||
hashes = [format(fnv1a(f.tobytes()), '016x') for f in frames_int]
|
||
clip_hash = format(fnv1a(b''.join(bytes.fromhex(h) for h in hashes)), '016x')
|
||
|
||
open(f'{outdir}/world_T.bin', 'wb').write(whm1(T_layers))
|
||
open(f'{outdir}/world_gs.bin', 'wb').write(whm1(gs_layers))
|
||
json.dump({'c': C, 'lat': LAT, 'frame': FRAME, 'a': A, 'afrac': AFRAC,
|
||
'z0': z0_int.reshape(-1).tolist(), 'actions': demo_acts,
|
||
'frame_hashes': hashes, 'clip_hash': clip_hash},
|
||
open(f'{outdir}/world.json', 'w'))
|
||
# raw RGB frames for the demo (optional viewing)
|
||
np.stack(frames_int).astype(np.uint8).tofile(f'{outdir}/world_frames.rgb')
|
||
print(f'int world model: {len(frames_int)} frames {FRAME}x{FRAME}, T {len(whm1(T_layers))}B + g_s '
|
||
f'{len(whm1(gs_layers))}B, float SNR {float(snr_f):.1f} dB / integer {snr_i:.1f} dB')
|
||
print(f' clip_hash {clip_hash} frame0 {hashes[0]}')
|
||
|
||
|
||
if __name__ == '__main__':
|
||
out = sys.argv[1] if len(sys.argv) > 1 else '/tmp/int_world'
|
||
# optional SOTA-scale overrides: FRAME LAT C ROLL WORLD (FRAME must = LAT·8)
|
||
if len(sys.argv) > 2:
|
||
FRAME = int(sys.argv[2])
|
||
if len(sys.argv) > 3:
|
||
LAT = int(sys.argv[3])
|
||
if len(sys.argv) > 4:
|
||
C = int(sys.argv[4])
|
||
if len(sys.argv) > 5:
|
||
ROLL = int(sys.argv[5])
|
||
if len(sys.argv) > 6:
|
||
WORLD = int(sys.argv[6])
|
||
if len(sys.argv) > 7:
|
||
STEPS = int(sys.argv[7])
|
||
main(out)
|