SuperHQ Motion experiment: locomotion and action checkpoints, runtime, clips, evaluation
6adbbb1 verified | /** Dense FP32 inference for the small Motion MLP. Feature construction is a separate layer. */ | |
| export class DenseNetwork { | |
| constructor(manifest, buffer) { | |
| if (!['superhq-dense-f32-v1', 'superhq-dense-f32-v2', 'superhq-dense-f32-v3'].includes(manifest.format) || manifest.byte_order !== 'little' || | |
| !(buffer instanceof ArrayBuffer) || buffer.byteLength !== manifest.bytes || buffer.byteLength % 4) { | |
| throw new Error('Invalid dense network manifest or weight buffer.'); | |
| } | |
| const bytes = new Uint8Array(new Uint32Array([1]).buffer); | |
| if (bytes[0] !== 1) throw new Error('This runtime requires a little-endian platform.'); | |
| this.manifest = manifest; | |
| const weights = new Float32Array(buffer); | |
| if (!weights.every(Number.isFinite)) throw new Error('Nonfinite network weights.'); | |
| let offset = 0; | |
| this.layers = manifest.layers.map((layer, i) => { | |
| const { inputs, outputs, weight_offset, bias_offset, activation } = layer; | |
| if (!Number.isInteger(inputs) || !Number.isInteger(outputs) || inputs < 1 || outputs < 1 || | |
| weight_offset !== offset || bias_offset !== offset + inputs * outputs || | |
| (i > 0 && inputs !== manifest.layers[i - 1].outputs) || | |
| !['silu', 'identity'].includes(activation)) throw new Error('Invalid dense layer layout.'); | |
| offset = bias_offset + outputs; | |
| if (offset > weights.length) throw new Error('Layer exceeds weight buffer.'); | |
| return { inputs, outputs, activation, weight: weights.subarray(weight_offset, bias_offset), | |
| bias: weights.subarray(bias_offset, offset), output: new Float32Array(outputs) }; | |
| }); | |
| if (!this.layers.length || offset !== weights.length || offset !== manifest.parameters) { | |
| throw new Error('Weight count differs from manifest.'); | |
| } | |
| this.outputSkip = manifest.output_skip ?? []; | |
| this.angularPrior = manifest.angular_prior; | |
| if ((manifest.format === 'superhq-dense-f32-v3') !== (this.angularPrior === 'articulated_constant_velocity_v1') || | |
| (this.angularPrior && (this.layers[0].inputs !== 850 || this.layers.at(-1).outputs !== 144))) { | |
| throw new Error('Invalid angular prior contract.'); | |
| } | |
| if (!Array.isArray(this.outputSkip) || | |
| (manifest.format === 'superhq-dense-f32-v2' && !this.outputSkip.length) || | |
| (manifest.format === 'superhq-dense-f32-v1' && this.outputSkip.length)) throw new Error('Invalid output skip layout.'); | |
| for (const item of this.outputSkip) { | |
| if (!Number.isInteger(item.input) || item.input < 0 || item.input >= this.layers[0].inputs || | |
| !Number.isInteger(item.output) || item.output < 0 || item.output >= this.layers.at(-1).outputs || | |
| !Number.isFinite(item.scale)) throw new Error('Invalid output skip connection.'); | |
| } | |
| } | |
| /** The returned scratch buffer is reused by the next call. Copy it when retaining a prediction. */ | |
| forward(input) { | |
| if (!(input instanceof Float32Array) || input.length !== this.layers[0].inputs || !input.every(Number.isFinite)) { | |
| throw new Error('Expected a finite Float32Array of the network input width.'); | |
| } | |
| let value = input; | |
| for (const layer of this.layers) { | |
| const { inputs, outputs, weight, bias, output, activation } = layer; | |
| for (let row = 0; row < outputs; row++) { | |
| let sum = bias[row]; | |
| const start = row * inputs; | |
| for (let col = 0; col < inputs; col++) sum += weight[start + col] * value[col]; | |
| output[row] = activation === 'silu' ? sum / (1 + Math.exp(-sum)) : sum; | |
| } | |
| if (!output.every(Number.isFinite)) throw new Error('Nonfinite network output.'); | |
| value = output; | |
| } | |
| for (const item of this.outputSkip) value[item.output] += input[item.input] * item.scale; | |
| if (this.angularPrior) addAngularPrior(input, value); | |
| if (!value.every(Number.isFinite)) throw new Error('Nonfinite network output after skip connections.'); | |
| return value; | |
| } | |
| } | |
| function rotation(input, offset) { | |
| const a = Array.from(input.subarray(offset, offset + 3)); | |
| const b = Array.from(input.subarray(offset + 3, offset + 6)); | |
| const na = Math.hypot(...a); | |
| if (na < 1e-6) throw new Error('Degenerate angular prior rotation.'); | |
| const x = a.map(v => v / na), dot = x.reduce((s, v, i) => s + v * b[i], 0); | |
| const y = b.map((v, i) => v - dot * x[i]), nb = Math.hypot(...y); | |
| if (nb < 1e-6) throw new Error('Degenerate angular prior rotation.'); | |
| for (let i = 0; i < 3; i++) y[i] /= nb; | |
| const z = [x[1]*y[2]-x[2]*y[1], x[2]*y[0]-x[0]*y[2], x[0]*y[1]-x[1]*y[0]]; | |
| return [x, y, z]; // Columns, matching the Python six-dimensional representation. | |
| } | |
| function addAngularPrior(input, output) { | |
| for (let j = 0; j < 22; j++) { | |
| const previous = rotation(input, 276 + 6*j), current = rotation(input, 552 + 6*j); | |
| for (let col = 0; col < 2; col++) { | |
| // current * transpose(previous) * currentColumn | |
| const intermediate = previous.map(p => p.reduce((s, v, i) => s + v*current[col][i], 0)); | |
| for (let row = 0; row < 3; row++) { | |
| const extrapolated = current.reduce((s, c, i) => s + c[row]*intermediate[i], 0); | |
| output[6*j + 3*col + row] += extrapolated-current[col][row]; | |
| } | |
| } | |
| } | |
| } | |