feat(playground/nisps): warm-start weight transfer on MLP resize (meml-8jc)

Add WasmIML.extractWeights() and WasmIML.createWithWarmStart() to preserve
learned joystick mappings across output-count changes; resizeMLP() now
transfers hidden-layer weights and shared output nodes instead of cold-starting.
This commit is contained in:
w1n5t0n 2026-04-03 17:27:54 +01:00
parent 2d2cabe159
commit c69b46f8cc
2 changed files with 80 additions and 8 deletions

View file

@ -724,35 +724,45 @@ function outputCountForMode(mode) {
/** /**
* Recreate IML instances with a new output count. * Recreate IML instances with a new output count.
* Resets all weights and training data. * Joystick IML uses warm-start weight transfer to preserve learned mappings.
* Training examples are always cleared (dataset is JS-side and output-count-specific).
*/ */
async function resizeMLP(newOutputCount) { async function resizeMLP(newOutputCount) {
if (newOutputCount === N_OUTPUTS) return; if (newOutputCount === N_OUTPUTS) return;
N_OUTPUTS = newOutputCount; N_OUTPUTS = newOutputCount;
// Extract joystick weights before destroying (for warm-start transfer)
const joySnapshot = imlJoy ? imlJoy.extractWeights() : null;
// Destroy old IML instances (free WASM memory) // Destroy old IML instances (free WASM memory)
if (imlJoy) imlJoy.destroy(); if (imlJoy) imlJoy.destroy();
if (imlHand) imlHand.destroy(); if (imlHand) imlHand.destroy();
imlJoy = await WasmIML.create(N_JOY_INPUTS, N_OUTPUTS, [32, 48, 64], 1000, 1.0, 0.00001); // Joystick IML: warm-start from previous weights when possible
if (joySnapshot) {
imlJoy = await WasmIML.createWithWarmStart(joySnapshot, N_OUTPUTS, 1000, 1.0, 0.00001);
} else {
imlJoy = await WasmIML.create(N_JOY_INPUTS, N_OUTPUTS, [32, 48, 64], 1000, 1.0, 0.00001);
imlJoy.randomiseWeights(spreadLevel);
}
imlJoy.setLogger(msg => console.log('[NISPS:joy]', msg)); imlJoy.setLogger(msg => console.log('[NISPS:joy]', msg));
// Hand IML: fresh init (warm-start for 14-input networks is a future concern)
imlHand = await WasmIML.create(N_HAND_INPUTS, N_OUTPUTS, [48, 48, 64], 1000, 1.0, 0.00001); imlHand = await WasmIML.create(N_HAND_INPUTS, N_OUTPUTS, [48, 48, 64], 1000, 1.0, 0.00001);
imlHand.setLogger(msg => console.log('[NISPS:hand]', msg)); imlHand.setLogger(msg => console.log('[NISPS:hand]', msg));
imlHand.randomiseWeights(spreadLevel);
iml = (inputMode === 'joystick') ? imlJoy : imlHand; iml = (inputMode === 'joystick') ? imlJoy : imlHand;
// Reset dependent state // Reset dependent state
rawParamValues = new Array(N_OUTPUTS).fill(0.5); rawParamValues = new Array(N_OUTPUTS).fill(0.5);
// Randomize with current spread
imlJoy.randomiseWeights(spreadLevel);
imlHand.randomiseWeights(spreadLevel);
// Re-run inference // Re-run inference
iml.setInput(0, joyX); iml.setInput(0, joyX);
iml.setInput(1, joyY); iml.setInput(1, joyY);
iml.process(); iml.process();
console.log(`[NISPS] MLP resized to ${N_OUTPUTS} outputs`); console.log(`[NISPS] MLP resized to ${N_OUTPUTS} outputs (joystick IML warm-started)`);
} }
// ---- Init ---- // ---- Init ----
@ -1842,7 +1852,7 @@ async function setOutputMode(mode, { skipConfirm = false } = {}) {
// Warn about weight reset if resizing (skip during state restore) // Warn about weight reset if resizing (skip during state restore)
if (needsResize && !skipConfirm && iml.dataset.features.length > 0) { if (needsResize && !skipConfirm && iml.dataset.features.length > 0) {
if (!confirm(`Switching to ${mode} mode requires ${targetOutputs} outputs (currently ${N_OUTPUTS}). This will reset the neural network weights and training data. Continue?`)) { if (!confirm(`Switching to ${mode} mode requires ${targetOutputs} outputs (currently ${N_OUTPUTS}). This will reset training examples. Network weights will be partially preserved. Continue?`)) {
syncOutputToggles(outputMode); // revert pill UI syncOutputToggles(outputMode); // revert pill UI
return; return;
} }

View file

@ -405,6 +405,68 @@ export class WasmIML {
this.process(); this.process();
} }
// ---- Public weight snapshot for warm-start transfer ----
/**
* Returns a plain object describing the full network weights for later
* reinjection via WasmIML.createWithWarmStart().
*/
extractWeights() {
return {
layerSizes: [...this.layerSizes], // e.g. [3, 32, 48, 64, 126]
weights: this._getFlatWeights(), // plain Array from fromHeapF32
};
}
/**
* Async static factory: create a new WasmIML with newOutputCount outputs,
* transferring as much of snapshot.weights as possible.
* Hidden-layer weights are copied unchanged; output nodes beyond the old
* count retain their random-initialised values.
*/
static async createWithWarmStart(snapshot, newOutputCount, maxIter, learningRate, convergenceThreshold) {
const oldLayerSizes = snapshot.layerSizes;
const nInputs = oldLayerSizes[0] - 1; // stored with bias (+1), strip it
const hiddenLayers = oldLayerSizes.slice(1, -1); // e.g. [32, 48, 64]
const oldOutputCount = oldLayerSizes[oldLayerSizes.length - 1];
// Create fresh instance with new output count
const newIml = await WasmIML.create(nInputs, newOutputCount, hiddenLayers, maxIter, learningRate, convergenceThreshold);
// Calculate prefix weight count (all layers except the final output layer).
// Each layer l: n_nodes[l] * (n_nodes[l-1] + 1) weights (inputs + bias).
// fullOldSizes: [nInputs+1, ...hiddenLayers, oldOutputCount]
let prefixCount = 0;
for (let l = 1; l < oldLayerSizes.length - 1; l++) {
prefixCount += oldLayerSizes[l] * (oldLayerSizes[l - 1] + 1);
}
const lastHidden = hiddenLayers[hiddenLayers.length - 1]; // e.g. 64
const weightsPerOutputNode = lastHidden + 1; // inputs from last hidden + bias
// Build new weight array, starting from random init
const newWeights = newIml._getFlatWeights();
const oldWeights = snapshot.weights;
// Copy hidden layer weights unchanged
for (let i = 0; i < prefixCount; i++) {
newWeights[i] = oldWeights[i];
}
// Copy output layer weights for nodes that existed in the old network
const sharedOutputNodes = Math.min(oldOutputCount, newOutputCount);
for (let n = 0; n < sharedOutputNodes; n++) {
const oldOff = prefixCount + n * weightsPerOutputNode;
const newOff = prefixCount + n * weightsPerOutputNode;
for (let w = 0; w < weightsPerOutputNode; w++) {
newWeights[newOff + w] = oldWeights[oldOff + w];
}
}
// Nodes beyond old count retain their random init — good for exploration
newIml._setFlatWeights(newWeights);
return newIml;
}
// ---- Flat weight get/set (for storedWeights save/restore) ---- // ---- Flat weight get/set (for storedWeights save/restore) ----
_getFlatWeights() { _getFlatWeights() {
const ptr = this._w.alloc(this._weightCount); const ptr = this._w.alloc(this._weightCount);