From c69b46f8cc33679f8d55bb066116555e6e025a62 Mon Sep 17 00:00:00 2001
From: w1n5t0n <w1n5t0n@lnfinitemonkeys.org>
Date: Fri, 3 Apr 2026 17:27:54 +0100
Subject: [PATCH] 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.
---
 playground/js/a-app.js            | 26 +++++++++----
 playground/js/nisps/nisps-wasm.js | 62 +++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 8 deletions(-)

diff --git a/playground/js/a-app.js b/playground/js/a-app.js
index 1cd334d..f9f5213 100644
--- a/playground/js/a-app.js
+++ b/playground/js/a-app.js
@@ -724,35 +724,45 @@ function outputCountForMode(mode) {
 
 /**
  * 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) {
   if (newOutputCount === N_OUTPUTS) return;
   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)
   if (imlJoy) imlJoy.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));
+
+  // 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.setLogger(msg => console.log('[NISPS:hand]', msg));
+  imlHand.randomiseWeights(spreadLevel);
+
   iml = (inputMode === 'joystick') ? imlJoy : imlHand;
 
   // Reset dependent state
   rawParamValues = new Array(N_OUTPUTS).fill(0.5);
 
-  // Randomize with current spread
-  imlJoy.randomiseWeights(spreadLevel);
-  imlHand.randomiseWeights(spreadLevel);
-
   // Re-run inference
   iml.setInput(0, joyX);
   iml.setInput(1, joyY);
   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 ----
@@ -1842,7 +1852,7 @@ async function setOutputMode(mode, { skipConfirm = false } = {}) {
 
   // Warn about weight reset if resizing (skip during state restore)
   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
       return;
     }
diff --git a/playground/js/nisps/nisps-wasm.js b/playground/js/nisps/nisps-wasm.js
index a945f62..f83cfc5 100644
--- a/playground/js/nisps/nisps-wasm.js
+++ b/playground/js/nisps/nisps-wasm.js
@@ -405,6 +405,68 @@ export class WasmIML {
     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) ----
   _getFlatWeights() {
     const ptr = this._w.alloc(this._weightCount);