619 lines
28 KiB
TypeScript
619 lines
28 KiB
TypeScript
import assert from 'node:assert/strict';
|
|
import { existsSync } from 'node:fs';
|
|
import { join } from 'node:path';
|
|
import test from 'node:test';
|
|
import { fileURLToPath } from 'node:url';
|
|
import {
|
|
DEFAULT_STIMULATION_MS,
|
|
GAMEBOY_MS_PER_FRAME,
|
|
NeuralAgent,
|
|
type AgentState,
|
|
type RewardEvent,
|
|
} from '../src/agent/agent';
|
|
import {
|
|
AGENT_CHUNK_NAMES,
|
|
agentFromChunks,
|
|
agentToChunks,
|
|
decodeEnvelope,
|
|
encodeEnvelope,
|
|
type AgentManifest,
|
|
} from '../src/agent/envelope';
|
|
import type { BrainDataset } from '../src/dataset/format';
|
|
import { loadBrainDatasetFromDir } from '../src/dataset/load-node';
|
|
import type { DecoderConfig } from '../src/readout/decoder';
|
|
import { gameboyDecoderConfig, toButtonMask } from '../src/readout/presets/gameboy';
|
|
import { toyDataset, xorshift } from './fixtures/toy-dataset';
|
|
import { MotorDecoder } from './legacy/decoder';
|
|
import { FlyBrain as LegacyFlyBrain } from './legacy/lif';
|
|
|
|
const FRAME_WIDTH = 160;
|
|
const FRAME_HEIGHT = 144;
|
|
const MAGIC = 'FLYBRAIN';
|
|
|
|
/**
|
|
* Synthetic connectome big enough for the loop to be worth measuring.
|
|
*
|
|
* The shared toy dataset has four neurons and one command role, so a Game Boy readout on it can
|
|
* never fire a pulse channel and the exclusive group has one real candidate. This fixture keeps
|
|
* the same shape (CSR, roles, retina columns) but sizes the populations so that: the noise budget
|
|
* of 300 kicks spread over 4096 neurons leaves the network sub-threshold and the connectome
|
|
* actually drives it, all eight `command_*` roles exist with distinct neurons, and every
|
|
* kenyon -> mbon edge is positive so plasticity has slots to move.
|
|
*/
|
|
function syntheticConnectome(seed = 20260915): BrainDataset {
|
|
const random = xorshift(seed);
|
|
const neurons = 4096;
|
|
const kenyon = { from: 0, to: 1024 };
|
|
const mbon = { from: 1024, to: 1280 };
|
|
const roles: Record<string, number[]> = {
|
|
kenyon: range(kenyon.from, kenyon.to),
|
|
mbon: range(mbon.from, mbon.to),
|
|
reward_pam: range(1280, 1344),
|
|
};
|
|
for (let command = 0; command < 8; command++) roles[`command_${command}`] = range(1344 + command * 32, 1376 + command * 32);
|
|
const visual = range(1600, 2400);
|
|
|
|
const indptr = new Uint32Array(neurons + 1);
|
|
const targets: number[] = [];
|
|
const weights: number[] = [];
|
|
for (let source = 0; source < neurons; source++) {
|
|
indptr[source] = targets.length;
|
|
const isKenyon = source >= kenyon.from && source < kenyon.to;
|
|
for (let edge = 0; edge < 10; edge++) {
|
|
// Four of every Kenyon cell's ten edges land on an MBON with a positive weight: those are
|
|
// the ones the plasticity rule is allowed to select.
|
|
const forced = isKenyon && edge < 4;
|
|
const target = forced
|
|
? mbon.from + Math.floor(random() * (mbon.to - mbon.from))
|
|
: Math.floor(random() * neurons);
|
|
const magnitude = 1 + Math.floor(random() * 40);
|
|
targets.push(target);
|
|
weights.push(!forced && random() < 0.2 ? -magnitude : magnitude);
|
|
}
|
|
}
|
|
indptr[neurons] = targets.length;
|
|
|
|
const visualXY = new Float32Array(visual.length * 2);
|
|
for (let column = 0; column < visual.length; column++) {
|
|
visualXY[column * 2] = (column % 40) * 7.5;
|
|
visualXY[column * 2 + 1] = Math.floor(column / 40) * 5.25;
|
|
}
|
|
return {
|
|
meta: {
|
|
schemaVersion: 1, dataset: 'synthetic', neurons, edges: targets.length, roles,
|
|
visual: { population: 'synthetic-retina', count: visual.length },
|
|
},
|
|
indptr,
|
|
targets: Uint32Array.from(targets),
|
|
weights: Int16Array.from(weights),
|
|
visualIndices: Uint32Array.from(visual),
|
|
visualHemisphere: Uint8Array.from(visual.map((_, column) => column % 2)),
|
|
visualXY,
|
|
};
|
|
}
|
|
|
|
function range(from: number, to: number): number[] {
|
|
return Array.from({ length: to - from }, (_, index) => from + index);
|
|
}
|
|
|
|
/** Deterministic RGBA frames, cycled by the loops below so both sides see the same images. */
|
|
function framePool(count: number, seed = 4242): Uint8Array[] {
|
|
const random = xorshift(seed);
|
|
return Array.from({ length: count }, () => {
|
|
const rgba = new Uint8Array(FRAME_WIDTH * FRAME_HEIGHT * 4);
|
|
for (let index = 0; index < rgba.length; index++) rgba[index] = Math.floor(random() * 256);
|
|
return rgba;
|
|
});
|
|
}
|
|
|
|
/** Rewards, boot mode and manual-input frames, on prime-ish periods so they interleave. */
|
|
function rewardsFor(frame: number): RewardEvent[] {
|
|
if (frame % 137 === 0) return [{ value: -0.4, stimulationMs: 40 }, { value: 0.25 }];
|
|
if (frame % 50 === 0) return [{ value: 0.6 }];
|
|
return [];
|
|
}
|
|
const bootFor = (frame: number): boolean => Math.floor(frame / 120) % 2 === 0;
|
|
const learnFor = (frame: number): boolean => frame % 90 !== 0;
|
|
|
|
/**
|
|
* The Game Boy preset with the exclusive group's timings pinned to `MotorDecoder`'s own hardcoded
|
|
* constants: a 400 ms hold, a 400 ms decision period and a 0.08 fatigue gain.
|
|
*
|
|
* The oracle below proves that `NeuralAgent` reproduces the original worker loop frame by frame,
|
|
* and the loop it is compared against holds `MotorDecoder`, whose timings are not configurable.
|
|
* `docs/design/room-escape.md` sections 1 and 3 moved the live preset's hold, fatigue gain and
|
|
* hysteresis and switched on its blocked-direction cooldown, any of which would make the oracle fail
|
|
* for a reason that has nothing to do with the agent's glue. The live numbers are pinned in
|
|
* `tests/readout.test.ts` instead, and `tools/golden.ts`'s `agent` scenario carries them into the
|
|
* Rust port.
|
|
*/
|
|
function legacyGameboyConfig(): DecoderConfig {
|
|
const config = gameboyDecoderConfig();
|
|
return {
|
|
...config,
|
|
exclusive: {
|
|
...config.exclusive!,
|
|
decisionMs: 400,
|
|
holdMs: 400,
|
|
fatigueGain: 0.08,
|
|
hysteresis: 1.15,
|
|
blockedFatigue: 0,
|
|
blockedMs: 0,
|
|
},
|
|
};
|
|
}
|
|
|
|
function agentFor(data: BrainDataset): NeuralAgent {
|
|
return new NeuralAgent(data, { decoder: legacyGameboyConfig(), frame: { width: FRAME_WIDTH, height: FRAME_HEIGHT } });
|
|
}
|
|
|
|
/** The original worker's initialize + per-frame sequence, written out by hand. */
|
|
class LegacyLoop {
|
|
readonly brain: LegacyFlyBrain;
|
|
readonly decoder = new MotorDecoder();
|
|
remainder = 0;
|
|
|
|
constructor(data: BrainDataset, firstFrame: Uint8Array) {
|
|
this.brain = new LegacyFlyBrain(data);
|
|
this.brain.plasticity.enabled = false;
|
|
this.brain.step(2_500);
|
|
this.brain.plasticity.enabled = true;
|
|
this.decoder.calibrate(this.brain.rates);
|
|
this.brain.setVisualFrame(firstFrame);
|
|
}
|
|
|
|
/** Returns the button mask and the number of 1-ms steps taken. */
|
|
frame(image: Uint8Array, events: RewardEvent[], boot: boolean, learn: boolean): { mask: number; steps: number } {
|
|
this.remainder += GAMEBOY_MS_PER_FRAME;
|
|
const steps = Math.floor(this.remainder);
|
|
this.remainder -= steps;
|
|
this.brain.step(steps);
|
|
const mask = this.decoder.decode(this.brain.rates, this.brain.ms, boot);
|
|
this.brain.setVisualFrame(image);
|
|
for (const event of events) this.brain.reward(event.stimulationMs ?? DEFAULT_STIMULATION_MS);
|
|
if (learn) this.brain.plasticity.reinforce(events.reduce((sum, event) => sum + event.value, 0), this.brain.ms);
|
|
return { mask, steps };
|
|
}
|
|
}
|
|
|
|
function cloneState(state: AgentState): AgentState {
|
|
const { network } = state;
|
|
return {
|
|
...state,
|
|
network: {
|
|
...network,
|
|
membrane: network.membrane.slice(),
|
|
refractory: network.refractory.slice(),
|
|
lastSpikeMs: network.lastSpikeMs.slice(),
|
|
visualDrive: network.visualDrive.slice(),
|
|
rates: { ...network.rates },
|
|
plasticity: {
|
|
...network.plasticity,
|
|
gains: network.plasticity.gains.slice(),
|
|
traces: network.plasticity.traces.slice(),
|
|
touched: network.plasticity.touched.slice(),
|
|
},
|
|
},
|
|
decoder: {
|
|
...state.decoder,
|
|
baseline: { ...state.decoder.baseline },
|
|
heldUntil: { ...state.decoder.heldUntil },
|
|
nextAllowed: { ...state.decoder.nextAllowed },
|
|
fatigue: { ...state.decoder.fatigue },
|
|
},
|
|
};
|
|
}
|
|
|
|
// --- Oracle: the agent loop against the original worker's glue, frame by frame.
|
|
|
|
test('ORACLE: NeuralAgent reproduces the original worker loop over 600 frames', () => {
|
|
const data = syntheticConnectome();
|
|
const frames = framePool(16);
|
|
const legacy = new LegacyLoop(data, frames[0]!);
|
|
const agent = agentFor(data);
|
|
agent.warmup(frames[0]!);
|
|
|
|
// Warm-up must land on identical state before a single frame is ticked.
|
|
assert.deepEqual(agent.network.exportState(), legacy.brain.exportState());
|
|
assert.equal(agent.network.ms, 2_500);
|
|
|
|
let fired = 0;
|
|
let directions = 0;
|
|
for (let frame = 1; frame <= 600; frame++) {
|
|
const image = frames[frame % frames.length]!;
|
|
const events = rewardsFor(frame);
|
|
const boot = bootFor(frame);
|
|
const learn = learnFor(frame);
|
|
const expected = legacy.frame(image, events, boot, learn);
|
|
const result = agent.tick(image, { rewards: events, boot, learn });
|
|
assert.equal(result.steps, expected.steps, `step count differs at frame ${frame}`);
|
|
assert.equal(toButtonMask(result.active), expected.mask, `button mask differs at frame ${frame}`);
|
|
if (result.active.some(channel => ['a', 'b', 'start', 'select'].includes(channel))) fired++;
|
|
if (result.active.some(channel => ['up', 'down', 'left', 'right'].includes(channel))) directions++;
|
|
}
|
|
|
|
// A run that never fires a pulse or never holds a direction would pass vacuously.
|
|
assert.ok(fired > 0, 'no pulse channel ever fired');
|
|
assert.ok(directions > 0, 'no direction was ever held');
|
|
assert.deepEqual(agent.network.exportState(), legacy.brain.exportState());
|
|
assert.ok(agent.plasticity.statistics().updates > 0, 'plasticity never updated');
|
|
assert.ok(agent.plasticity.statistics().changed > 0, 'no gain ever moved');
|
|
|
|
const before = legacy.decoder.exportState();
|
|
const after = agent.exportState();
|
|
assert.deepEqual(after.decoder.baseline, before.baseline);
|
|
assert.deepEqual(after.decoder.heldUntil, before.heldUntil);
|
|
assert.deepEqual(after.decoder.nextAllowed, before.nextAllowed);
|
|
assert.deepEqual(after.decoder.fatigue, before.fatigue);
|
|
assert.equal(after.decoder.current, before.direction);
|
|
assert.equal(after.decoder.nextDecision, before.nextDirectionDecision);
|
|
assert.equal(after.remainder, legacy.remainder);
|
|
assert.equal(after.warmedUp, true);
|
|
assert.equal(after.version, 1);
|
|
});
|
|
|
|
test('one Game Boy frame is 70224 dot clocks and the loop never drifts', () => {
|
|
assert.equal(GAMEBOY_MS_PER_FRAME, 1000 / (4_194_304 / 70_224));
|
|
const data = toyDataset();
|
|
const agent = new NeuralAgent(data, { decoder: gameboyDecoderConfig(), frame: { width: 8, height: 4 }, warmupMs: 10 });
|
|
assert.equal(agent.msPerFrame, GAMEBOY_MS_PER_FRAME);
|
|
assert.equal(agent.warmupMs, 10);
|
|
assert.equal(agent.ready, false);
|
|
const blank = new Uint8Array(8 * 4 * 4);
|
|
agent.warmup();
|
|
assert.equal(agent.ready, true);
|
|
let steps = 0;
|
|
for (let frame = 0; frame < 1000; frame++) steps += agent.tick(blank).steps;
|
|
// 1000 frames at 59.7275 fps is 16742.7 ms: the fractional remainder must be carried, not lost.
|
|
assert.equal(steps, Math.floor(1000 * GAMEBOY_MS_PER_FRAME));
|
|
assert.equal(agent.network.ms, 10 + steps);
|
|
const remainder = agent.exportState().remainder;
|
|
assert.ok(remainder >= 0 && remainder < 1);
|
|
});
|
|
|
|
test('warmup calibrates once and tick before warmup is refused', () => {
|
|
const data = toyDataset();
|
|
const agent = new NeuralAgent(data, { decoder: gameboyDecoderConfig(), frame: { width: 4, height: 4 }, warmupMs: 50 });
|
|
const blank = new Uint8Array(4 * 4 * 4);
|
|
assert.throws(() => agent.tick(blank), /Warm up the agent/);
|
|
agent.warmup();
|
|
assert.throws(() => agent.warmup(), /already warmed up/);
|
|
// A wrong-sized frame is refused before the network advances, not half way through the tick.
|
|
const before = agent.exportState();
|
|
assert.throws(() => agent.tick(new Uint8Array(3)), /RGBA bytes/);
|
|
assert.throws(() => agent.resetTransients(new Uint8Array(3)), /RGBA bytes/);
|
|
assert.deepEqual(agent.exportState(), before);
|
|
|
|
// Calibration happens on settled rates with plasticity re-enabled. Roles the dataset does not
|
|
// declare calibrate to zero, which is what makes their score exactly 1 and keeps them silent.
|
|
assert.equal(agent.plasticity.enabled, true);
|
|
assert.deepEqual(agent.exportState().decoder.baseline, {
|
|
command_0: agent.network.rates.command_0,
|
|
command_1: 0, command_2: 0, command_3: 0, command_4: 0, command_5: 0, command_6: 0, command_7: 0,
|
|
});
|
|
assert.ok(agent.network.rates.command_0! > 0);
|
|
});
|
|
|
|
test('the frame size defaults to the kernel retina and is validated', () => {
|
|
const data = toyDataset();
|
|
const agent = new NeuralAgent(data, { decoder: gameboyDecoderConfig() });
|
|
assert.deepEqual(agent.frame, { width: 160, height: 144 });
|
|
assert.throws(() => new NeuralAgent(data, { decoder: gameboyDecoderConfig(), frame: { width: 0, height: 4 } }), /frame size/);
|
|
assert.throws(() => new NeuralAgent(data, { decoder: gameboyDecoderConfig(), warmupMs: -1 }), /warmupMs/);
|
|
assert.throws(() => new NeuralAgent(data, { decoder: gameboyDecoderConfig(), msPerFrame: 0 }), /msPerFrame/);
|
|
});
|
|
|
|
// --- The game-recovery hook.
|
|
|
|
test('resetTransients drops holds and traces but keeps the RNG, clock and gains', () => {
|
|
const data = syntheticConnectome();
|
|
const frames = framePool(8, 77);
|
|
const agent = agentFor(data);
|
|
agent.warmup(frames[0]!);
|
|
for (let frame = 1; frame <= 60; frame++) {
|
|
agent.tick(frames[frame % frames.length]!, { rewards: rewardsFor(frame), boot: bootFor(frame) });
|
|
}
|
|
|
|
const before = agent.exportState();
|
|
assert.ok(before.network.plasticity.traces.some(value => value !== 0), 'no eligibility to clear');
|
|
assert.ok(Object.values(before.decoder.heldUntil).some(value => value > 0), 'no hold to clear');
|
|
|
|
agent.resetTransients(frames[3]!);
|
|
const after = agent.exportState();
|
|
|
|
// Kept: everything the rollback did not invalidate.
|
|
assert.equal(after.network.ms, before.network.ms);
|
|
assert.equal(after.network.rng, before.network.rng);
|
|
assert.equal(after.network.populationRate, before.network.populationRate);
|
|
assert.deepEqual(after.network.rates, before.network.rates);
|
|
assert.deepEqual(after.network.membrane, before.network.membrane);
|
|
assert.deepEqual(after.network.plasticity.gains, before.network.plasticity.gains);
|
|
assert.equal(after.network.plasticity.updates, before.network.plasticity.updates);
|
|
assert.deepEqual(after.decoder.baseline, before.decoder.baseline);
|
|
assert.equal(after.remainder, before.remainder);
|
|
|
|
// Dropped: the timeline that no longer exists.
|
|
assert.ok(after.network.plasticity.traces.every(value => value === 0));
|
|
assert.ok(after.network.plasticity.touched.every(value => value === before.network.ms));
|
|
assert.equal(after.network.plasticity.signal, 0);
|
|
assert.ok(Object.values(after.decoder.heldUntil).every(value => value === 0));
|
|
assert.ok(Object.values(after.decoder.nextAllowed).every(value => value === before.network.ms + 480));
|
|
assert.equal(after.decoder.current, null);
|
|
assert.ok(Object.values(after.decoder.fatigue).every(value => value === 0));
|
|
assert.equal(after.decoder.nextDecision, before.network.ms);
|
|
|
|
// The replacement image became the visual drive.
|
|
assert.notDeepEqual(after.network.visualDrive, before.network.visualDrive);
|
|
});
|
|
|
|
// --- Checkpoints.
|
|
|
|
test('an exported agent resumes bit-exactly in a fresh agent', () => {
|
|
const data = syntheticConnectome();
|
|
const frames = framePool(8, 909);
|
|
const source = agentFor(data);
|
|
source.warmup(frames[0]!);
|
|
for (let frame = 1; frame <= 200; frame++) {
|
|
source.tick(frames[frame % frames.length]!, { rewards: rewardsFor(frame), boot: bootFor(frame), learn: learnFor(frame) });
|
|
}
|
|
|
|
const restored = agentFor(data);
|
|
assert.equal(restored.ready, false);
|
|
restored.importState(source.exportState());
|
|
assert.equal(restored.ready, true);
|
|
assert.deepEqual(restored.exportState(), source.exportState());
|
|
assert.equal(restored.compatibility(), source.compatibility());
|
|
|
|
for (let frame = 201; frame <= 320; frame++) {
|
|
const image = frames[frame % frames.length]!;
|
|
const options = { rewards: rewardsFor(frame), boot: bootFor(frame), learn: learnFor(frame) };
|
|
const expected = source.tick(image, options);
|
|
const actual = restored.tick(image, options);
|
|
assert.deepEqual(actual, expected, `diverged at frame ${frame}`);
|
|
}
|
|
assert.deepEqual(restored.exportState(), source.exportState());
|
|
assert.deepEqual(restored.snapshot(), source.snapshot());
|
|
});
|
|
|
|
test('snapshot reports the network without exposing its arrays', () => {
|
|
const data = syntheticConnectome();
|
|
const frames = framePool(4, 31);
|
|
const agent = agentFor(data);
|
|
agent.warmup(frames[0]!);
|
|
agent.tick(frames[1]!, { rewards: [{ value: 1 }] });
|
|
const snapshot = agent.snapshot();
|
|
assert.equal(snapshot.ms, agent.network.ms);
|
|
assert.equal(snapshot.populationRate, agent.network.populationRate);
|
|
assert.deepEqual(snapshot.rates, agent.network.rates);
|
|
assert.notEqual(snapshot.rates, agent.network.rates);
|
|
assert.equal(snapshot.spikeTimes.length, agent.network.lastSpikeMs.length);
|
|
assert.deepEqual(snapshot.spikeTimes, Float32Array.from(agent.network.lastSpikeMs));
|
|
assert.equal(snapshot.learning.version, 'fly-kc-mbon-rstdp-v2');
|
|
assert.ok(snapshot.learning.synapses > 0);
|
|
});
|
|
|
|
test('a rejected checkpoint leaves the agent exactly as it was', () => {
|
|
const data = syntheticConnectome();
|
|
const frames = framePool(8, 1234);
|
|
const agent = agentFor(data);
|
|
agent.warmup(frames[0]!);
|
|
for (let frame = 1; frame <= 100; frame++) agent.tick(frames[frame % frames.length]!, { rewards: rewardsFor(frame) });
|
|
const early = agent.exportState();
|
|
for (let frame = 101; frame <= 200; frame++) agent.tick(frames[frame % frames.length]!, { rewards: rewardsFor(frame) });
|
|
const current = agent.exportState();
|
|
|
|
const reject = (state: AgentState, pattern: RegExp) => {
|
|
assert.throws(() => agent.importState(state), pattern);
|
|
assert.deepEqual(agent.exportState(), current, 'a failed import must not change the agent');
|
|
};
|
|
|
|
// Version and warm-up flag.
|
|
reject({ ...cloneState(early), version: 2 as unknown as 1 }, /checkpoint version/);
|
|
reject({ ...cloneState(early), warmedUp: 'yes' as unknown as boolean }, /warm-up flag/);
|
|
|
|
// A remainder outside [0, 1) would desynchronize the network from the environment clock.
|
|
reject({ ...cloneState(early), remainder: 1.25 }, /frame remainder/);
|
|
reject({ ...cloneState(early), remainder: -0.5 }, /frame remainder/);
|
|
reject({ ...cloneState(early), remainder: Number.NaN }, /frame remainder/);
|
|
|
|
// Corrupt plasticity: rejected by the network before anything is written.
|
|
const badGains = cloneState(early);
|
|
badGains.network.plasticity.gains[0] = 5;
|
|
reject(badGains, /plasticity values/);
|
|
const badTopology = cloneState(early);
|
|
badTopology.network.plasticity.topology += 1;
|
|
reject(badTopology, /plasticity topology/);
|
|
|
|
// Corrupt readout only: the network half imports cleanly first, so this is the case that needs
|
|
// the rollback. Without it the agent would keep a frame-100 network under a frame-200 readout.
|
|
const badDecoder = cloneState(early);
|
|
badDecoder.decoder.heldUntil.up = Number.POSITIVE_INFINITY;
|
|
reject(badDecoder, /decoder checkpoint/);
|
|
const badFatigue = cloneState(early);
|
|
badFatigue.decoder.fatigue.left = 4;
|
|
reject(badFatigue, /decoder fatigue/);
|
|
|
|
// The valid state it rolled back from still loads.
|
|
agent.importState(early);
|
|
assert.deepEqual(agent.exportState(), early);
|
|
});
|
|
|
|
// --- Envelope.
|
|
|
|
/** Encodes a manifest verbatim, so decode's own structural checks can be exercised. */
|
|
function encodeRaw(magic: string, manifest: object, chunks: Uint8Array[]): ArrayBuffer {
|
|
const magicBytes = new TextEncoder().encode(magic);
|
|
const manifestBytes = new TextEncoder().encode(JSON.stringify(manifest));
|
|
const total = magicBytes.length + 8 + manifestBytes.length + chunks.reduce((sum, chunk) => sum + 4 + chunk.byteLength, 0);
|
|
const output = new Uint8Array(total);
|
|
const view = new DataView(output.buffer);
|
|
let offset = 0;
|
|
output.set(magicBytes, offset); offset += magicBytes.length;
|
|
view.setUint32(offset, manifestBytes.length, true); offset += 4;
|
|
output.set(manifestBytes, offset); offset += manifestBytes.length;
|
|
for (const chunk of chunks) {
|
|
view.setUint32(offset, chunk.byteLength, true); offset += 4;
|
|
output.set(chunk, offset); offset += chunk.byteLength;
|
|
}
|
|
view.setUint32(offset, crc32(output.subarray(0, offset)), true);
|
|
return output.buffer;
|
|
}
|
|
|
|
/** Independent CRC32, so the test does not trust the implementation it is checking. */
|
|
function crc32(bytes: Uint8Array): number {
|
|
let crc = 0xffffffff;
|
|
for (const byte of bytes) {
|
|
crc ^= byte;
|
|
for (let bit = 0; bit < 8; bit++) crc = crc & 1 ? (crc >>> 1) ^ 0xedb88320 : crc >>> 1;
|
|
}
|
|
return (crc ^ 0xffffffff) >>> 0;
|
|
}
|
|
|
|
test('an envelope round-trips its manifest and chunks', () => {
|
|
const chunks = { alpha: Uint8Array.of(1, 2, 3), beta: new Uint8Array(0), gamma: Uint8Array.from({ length: 300 }, (_, i) => i & 0xff) };
|
|
const buffer = encodeEnvelope(MAGIC, { note: 'hello', count: 7 }, chunks);
|
|
const decoded = decodeEnvelope(buffer, MAGIC);
|
|
assert.equal(decoded.manifest.schemaVersion, 2);
|
|
assert.equal(decoded.manifest.note, 'hello');
|
|
assert.equal(decoded.manifest.count, 7);
|
|
assert.deepEqual(decoded.manifest.chunks, ['alpha', 'beta', 'gamma']);
|
|
assert.deepEqual({ ...decoded.chunks }, chunks);
|
|
// Chunk names cannot reach a prototype key.
|
|
assert.equal(Object.getPrototypeOf(decoded.chunks), null);
|
|
// The footer is a CRC32 over everything before it.
|
|
const bytes = new Uint8Array(buffer);
|
|
assert.equal(new DataView(buffer).getUint32(bytes.length - 4, true), crc32(bytes.subarray(0, -4)));
|
|
assert.equal(new TextDecoder().decode(bytes.subarray(0, 8)), MAGIC);
|
|
});
|
|
|
|
test('an envelope rejects a foreign magic, a bad schema and a bad chunk name', () => {
|
|
const buffer = encodeEnvelope(MAGIC, {}, { alpha: Uint8Array.of(9) });
|
|
assert.throws(() => decodeEnvelope(buffer, 'OTHERMAG'), /Not a OTHERMAG envelope/);
|
|
assert.throws(() => decodeEnvelope(new Uint8Array(3).buffer, MAGIC), /Not a FLYBRAIN envelope/);
|
|
assert.throws(() => encodeEnvelope('', {}, {}), /magic/);
|
|
assert.throws(() => encodeEnvelope(MAGIC, {}, { 'bad-name': Uint8Array.of(1) }), /chunk name/);
|
|
|
|
assert.throws(() => decodeEnvelope(encodeRaw(MAGIC, { schemaVersion: 1, chunks: [] }, []), MAGIC), /envelope schema: 1/);
|
|
assert.throws(() => decodeEnvelope(encodeRaw(MAGIC, { schemaVersion: 2, chunks: ['bad-name'] }, [Uint8Array.of(1)]), MAGIC), /Invalid envelope chunks/);
|
|
assert.throws(() => decodeEnvelope(encodeRaw(MAGIC, { schemaVersion: 2, chunks: ['a', 'a'] }, [Uint8Array.of(1), Uint8Array.of(2)]), MAGIC), /Invalid envelope chunks/);
|
|
assert.throws(() => decodeEnvelope(encodeRaw(MAGIC, { schemaVersion: 2, chunks: [7] }, [Uint8Array.of(1)]), MAGIC), /Invalid envelope chunks/);
|
|
assert.throws(() => decodeEnvelope(encodeRaw(MAGIC, { schemaVersion: 2, chunks: 'alpha' }, []), MAGIC), /Invalid envelope chunks/);
|
|
});
|
|
|
|
test('an envelope rejects corruption, truncation and trailing bytes', () => {
|
|
const chunks = { alpha: Uint8Array.from({ length: 64 }, (_, i) => i), beta: Uint8Array.of(5, 6) };
|
|
const original = new Uint8Array(encodeEnvelope(MAGIC, { note: 'x' }, chunks));
|
|
|
|
const flipped = original.slice();
|
|
flipped[flipped.length - 10] ^= 0x01;
|
|
assert.throws(() => decodeEnvelope(flipped.buffer as ArrayBuffer, MAGIC), /checksum mismatch/);
|
|
|
|
const flippedManifest = original.slice();
|
|
flippedManifest[20] ^= 0x20;
|
|
assert.throws(() => decodeEnvelope(flippedManifest.buffer as ArrayBuffer, MAGIC), /checksum mismatch|JSON|Unsupported|Invalid/);
|
|
|
|
const trailing = new Uint8Array(original.length + 3);
|
|
trailing.set(original);
|
|
assert.throws(() => decodeEnvelope(trailing.buffer, MAGIC), /checksum mismatch/);
|
|
|
|
// Trailing data with a valid footer: only the length bookkeeping can catch it.
|
|
const padded = new Uint8Array(original.length + 4);
|
|
padded.set(original.subarray(0, original.length - 4));
|
|
new DataView(padded.buffer).setUint32(padded.length - 4, crc32(padded.subarray(0, -4)), true);
|
|
assert.throws(() => decodeEnvelope(padded.buffer, MAGIC), /trailing data/);
|
|
|
|
const shortChunk = original.slice(0, original.length - 8);
|
|
assert.throws(() => decodeEnvelope(shortChunk.buffer as ArrayBuffer, MAGIC), /truncated|checksum mismatch/);
|
|
|
|
const manifestBytes = new TextEncoder().encode(JSON.stringify({ schemaVersion: 2, chunks: ['alpha'] }));
|
|
const truncatedManifest = new Uint8Array(MAGIC.length + 4 + manifestBytes.length - 5);
|
|
truncatedManifest.set(new TextEncoder().encode(MAGIC));
|
|
new DataView(truncatedManifest.buffer).setUint32(MAGIC.length, manifestBytes.length, true);
|
|
assert.throws(() => decodeEnvelope(truncatedManifest.buffer, MAGIC), /manifest is truncated/);
|
|
|
|
// A chunk header that claims more bytes than the file holds.
|
|
const overlong = encodeRaw(MAGIC, { schemaVersion: 2, chunks: ['alpha', 'beta'] }, [Uint8Array.of(1, 2, 3)]);
|
|
assert.throws(() => decodeEnvelope(overlong, MAGIC), /truncated/);
|
|
});
|
|
|
|
test('agentToChunks and agentFromChunks survive a full envelope round trip', () => {
|
|
const data = syntheticConnectome();
|
|
const frames = framePool(8, 5150);
|
|
const source = agentFor(data);
|
|
source.warmup(frames[0]!);
|
|
for (let frame = 1; frame <= 150; frame++) {
|
|
source.tick(frames[frame % frames.length]!, { rewards: rewardsFor(frame), boot: bootFor(frame), learn: learnFor(frame) });
|
|
}
|
|
|
|
const { manifest, chunks } = agentToChunks(source.exportState());
|
|
assert.deepEqual(Object.keys(chunks), [...AGENT_CHUNK_NAMES]);
|
|
assert.equal(chunks.membrane!.byteLength, data.meta.neurons * 4);
|
|
assert.equal(chunks.lastSpikeMs!.byteLength, data.meta.neurons * 8);
|
|
assert.equal(chunks.visualDrive!.byteLength, data.visualIndices.length * 4);
|
|
|
|
const buffer = encodeEnvelope(MAGIC, { ...manifest, compatibility: source.compatibility() }, chunks);
|
|
const decoded = decodeEnvelope(buffer, MAGIC);
|
|
assert.equal(decoded.manifest.compatibility, source.compatibility());
|
|
const state = agentFromChunks(decoded.manifest as unknown as AgentManifest, decoded.chunks);
|
|
|
|
const restored = agentFor(data);
|
|
restored.importState(state);
|
|
assert.deepEqual(restored.exportState(), source.exportState());
|
|
for (let frame = 151; frame <= 200; frame++) {
|
|
const image = frames[frame % frames.length]!;
|
|
const options = { rewards: rewardsFor(frame), boot: bootFor(frame), learn: learnFor(frame) };
|
|
assert.deepEqual(restored.tick(image, options), source.tick(image, options), `diverged at frame ${frame}`);
|
|
}
|
|
assert.deepEqual(restored.exportState(), source.exportState());
|
|
});
|
|
|
|
test('agentFromChunks rejects an incomplete or misaligned checkpoint', () => {
|
|
const data = syntheticConnectome();
|
|
const agent = agentFor(data);
|
|
agent.warmup();
|
|
const { manifest, chunks } = agentToChunks(agent.exportState());
|
|
|
|
assert.throws(() => agentFromChunks({ ...manifest, agentVersion: 2 as unknown as 1 }, chunks), /checkpoint version/);
|
|
assert.throws(() => agentFromChunks({ ...manifest, decoder: undefined as unknown as AgentManifest['decoder'] }, chunks), /manifest is incomplete/);
|
|
for (const name of AGENT_CHUNK_NAMES) {
|
|
const missing = { ...chunks };
|
|
delete missing[name];
|
|
assert.throws(() => agentFromChunks(manifest, missing), new RegExp(`missing ${name}`));
|
|
}
|
|
const partial = { ...chunks, membrane: chunks.membrane!.subarray(0, chunks.membrane!.byteLength - 1) };
|
|
assert.throws(() => agentFromChunks(manifest, partial), /partial element/);
|
|
});
|
|
|
|
// --- Real dataset, when the artifacts are present in this worktree.
|
|
|
|
const dataDir = join(fileURLToPath(new URL('../../../', import.meta.url)), 'data', 'fafb-v783');
|
|
|
|
test('NeuralAgent runs on the real FAFB v783 dataset', {
|
|
skip: existsSync(join(dataDir, 'meta.json')) ? false : 'data/fafb-v783/meta.json is absent in this worktree',
|
|
}, async () => {
|
|
const data = await loadBrainDatasetFromDir(dataDir);
|
|
const agent = agentFor(data);
|
|
const frames = framePool(4, 783);
|
|
agent.warmup(frames[0]!);
|
|
assert.equal(agent.network.ms, 2_500);
|
|
assert.ok(agent.network.populationRate > 0);
|
|
|
|
let held = 0;
|
|
for (let frame = 1; frame <= 30; frame++) {
|
|
const result = agent.tick(frames[frame % frames.length]!, { rewards: frame % 10 === 0 ? [{ value: 1 }] : [] });
|
|
assert.ok(Number.isFinite(result.spikes) && result.spikes >= 0);
|
|
held += result.active.length;
|
|
}
|
|
const snapshot = agent.snapshot();
|
|
assert.ok(Object.values(snapshot.rates).every(Number.isFinite));
|
|
assert.ok(Object.keys(snapshot.rates).length >= 8);
|
|
assert.ok(Number.isFinite(snapshot.populationRate) && snapshot.populationRate > 0);
|
|
assert.ok(held > 0, 'the readout never held a channel');
|
|
assert.equal(snapshot.learning.synapses, 16_384);
|
|
|
|
const compatibility = agent.compatibility();
|
|
assert.ok(compatibility.includes('lif-1ms-f64-v2'), compatibility);
|
|
assert.ok(compatibility.includes('fly-kc-mbon-rstdp-v2'), compatibility);
|
|
assert.ok(compatibility.includes(data.fingerprint!), compatibility);
|
|
assert.equal(compatibility.split('/').length, 3);
|
|
});
|