vanilla ndarray, vanilla rnn evolution

counting-gru.js

/*
  evolving a GRU to count to 10
  in vanilla ndarray
*/
let ndarray = require('ndarray')
let ops = require('ndarray-ops')
let gaussian = require('gaussian')
let distribution = gaussian(0, 1)
let { gemv } = require('ndarray-blas-level2')

let initialParams = Array(20)
  .fill(0)
  .map(() => distribution.ppf(Math.random()))
  .concat(gruInit(10, 2))

function unpack(arr, num) {
  if (num > arr.length) {
    throw new Error('tried to get too many parameters')
  }
  return arr.splice(arr.length - num, num)
}

function jitter(p, sigma) {
  let noise = Array(p.length).fill(0).map(() => distribution.ppf(Math.random()))
  let params = p.map((v, k) => v + sigma * noise[k])

  return { noise, params }
}

function buildModel(initialParams) {
  let params = initialParams.slice()
  let rnn = gru(unpack(params, 420), 10, 2)
  let w1 = ndarray(unpack(params, 10 * 2), [2, 10])

  return { rnn, w1 }
}

function forward(m, x, h0) {
  let h1 = m.rnn(x, h0)
  let y = ff(m.w1, h1)

  return [y, h1]
}

function ff(w, x, b) {
  let y = ndarray(Array(w.shape[0]).fill(0))
  if (b) {
    ops.assign(y, b)
  }
  gemv(1, w, x, 1, y)
  return y
}

function gruInit(hiddenSize, xSize) {
  return Array(6 * hiddenSize)
    .fill(0)
    .concat(
      Array(3 * xSize * hiddenSize + 3 * hiddenSize * hiddenSize)
        .fill(0)
        .map(() => distribution.ppf(Math.random()))
    )
}

function gru(params, hiddenSize, xSize) {
  // params required: 3(x * h) + 3(h * h) + 6(h)
  let paramsRequired =
    3 * xSize * hiddenSize + 3 * hiddenSize * hiddenSize + 6 * hiddenSize

  let bir = ndarray(unpack(params, hiddenSize))
  let bin = ndarray(unpack(params, hiddenSize))
  let bii = ndarray(unpack(params, hiddenSize))
  let bhn = ndarray(unpack(params, hiddenSize))
  let bhr = ndarray(unpack(params, hiddenSize))
  let bhi = ndarray(unpack(params, hiddenSize))

  let Wir = ndarray(unpack(params, xSize * hiddenSize), [hiddenSize, xSize])
  let Whr = ndarray(unpack(params, hiddenSize * hiddenSize), [
    hiddenSize,
    hiddenSize
  ])

  let Wii = ndarray(unpack(params, xSize * hiddenSize), [hiddenSize, xSize])
  let Whi = ndarray(unpack(params, hiddenSize * hiddenSize), [
    hiddenSize,
    hiddenSize
  ])

  let Win = ndarray(unpack(params, xSize * hiddenSize), [hiddenSize, xSize])
  let Whn = ndarray(unpack(params, hiddenSize * hiddenSize), [
    hiddenSize,
    hiddenSize
  ])

  return (x, h0) => {
    // rt=sigmoid(Wirxt+bir+Whrh(t−1)+bhr)
    let rt = ndarray(Array(hiddenSize).fill(0.1))
    ops.add(rt, ff(Wir, x), ff(Whr, h0), bir, bhr)
    rt = sigmoid(rt)

    // it=sigmoid(Wiixt+bii+Whih(t−1)+bhi)
    let it = ndarray(Array(hiddenSize).fill(0.1))
    ops.add(it, ff(Wii, x), ff(Whi, h0), bii, bhi)
    it = sigmoid(it)

    // nt=tanh(Winxt+bin+rt∗(Whnh(t−1)+bhn))
    let nt = ndarray(Array(hiddenSize).fill(0.1))
    ops.mul(rt, rt, ff(Whn, h0, bhn))
    ops.add(nt, ff(Win, x), rt, bin)
    nt = tanh(nt)

    // ht=(1−it)∗nt+it∗h(t−1)
    let h1 = ndarray(Array(hiddenSize).fill(0.1))
    h1.data.forEach((v, k) => {
      h1.data[k] += (1 - it.data[k]) * nt.data[k] + it.data[k] * h0.data[k]
    })

    return h1
  }
}

function tanh(ndarr) {
  return ndarray(ndarr.data.slice().map(v => Math.tanh(v)), ndarr.shape)
}

function sigmoid(ndarr) {
  return ndarray(
    ndarr.data.slice().map(v => 1 / (1 + Math.exp(-v))),
    ndarr.shape
  )
}

function fitness(params) {
  const maxCount = 100
  const countTarget = 10
  let model = buildModel(params)
  let x = ndarray([1, 1])
  let h = ndarray(Array(20).fill(0))
  let y = ndarray([0, 0])
  let count = 0
  do {
    let out = forward(model, x, h)
    y = out[0]
    h = out[1]
    count++
  } while (y.data[0] < y.data[1] && count < maxCount)

  return -1 * Math.abs(countTarget - count)
}

const normalizationWindow = 20
const sigma = 0.2
const alpha = 0.01
let mean = 0
let sd = 1e-8
let headParams = initialParams.slice()
for (let i = 0; i < 10000; i++) {
  let { params, noise } = jitter(headParams, sigma)
  let score = fitness(params)
  const n = 1 / normalizationWindow
  mean = (1 - n) * mean + n * score
  sd = (1 - n) * sd + n * Math.pow(score - mean, 2)

  let advantage = (score - mean) / (sd + 1e-8)
  if (i > normalizationWindow) {
    headParams.forEach((v, k) => {
      headParams[k] += alpha / sigma * advantage * noise[k]
    })
  }
}

// higher fitness is better
// (best possible score is 0)
console.log('fitness is: ' + fitness(headParams))

counting-rnn.js

/*
  evolving a vanilla RNN to count to 10
  in vanilla ndarray
*/
let ndarray = require('ndarray')
let ops = require('ndarray-ops')
let gaussian = require('gaussian')
let distribution = gaussian(0, 1)
let { gemv } = require('ndarray-blas-level2')

let initialParams = Array(480)
  .fill(0)
  .map(() => distribution.ppf(Math.random()))

function unpack(arr, num) {
  if (num > arr.length) {
    throw new Error('tried to get too many parameters')
  }
  return arr.splice(arr.length - num, num)
}

function jitter(p, sigma) {
  let noise = Array(p.length).fill(0).map(() => distribution.ppf(Math.random()))
  let params = p.map((v, k) => v + sigma * noise[k])

  return { noise, params }
}

function buildModel(initialParams) {
  let params = initialParams.slice()
  let why = ndarray(unpack(params, 40), [2, 20])
  let wxh = ndarray(unpack(params, 40), [20, 2])
  let whh = ndarray(unpack(params, 400), [20, 20])

  return { why, wxh, whh }
}

function forward(m, x, h0) {
  let h1 = ndarray(Array(h0.data.length).fill(0), h0.shape)
  ops.add(h1, ff(m.whh, h0), ff(m.wxh, x))
  h1 = tanh(h1)
  let y = ff(m.why, h1)
  return [y, h1]
}

function ff(w, x) {
  let y = ndarray(Array(w.shape[0]).fill(0))
  gemv(1, w, x, 1, y)
  return y
}

function tanh(ndarr) {
  return ndarray(ndarr.data.slice().map(v => Math.tanh(v)), ndarr.shape)
}

function fitness(params) {
  const maxCount = 100
  const countTarget = 10
  let model = buildModel(params)
  let x = ndarray([1, 1])
  let h = ndarray(Array(20).fill(0))
  let y = ndarray([0, 0])
  let count = 0
  do {
    let out = forward(model, x, h)
    y = out[0]
    h = out[1]
    count++
  } while (y.data[0] < y.data[1] && count < maxCount)

  return -1 * Math.abs(countTarget - count)
}

const normalizationWindow = 20
const sigma = 0.2
const alpha = 0.01
let mean = 0
let sd = 1e-8
let headParams = initialParams.slice()
for (let i = 0; i < 10000; i++) {
  let { params, noise } = jitter(headParams, sigma)
  let score = fitness(params)
  const n = 1 / normalizationWindow
  mean = (1 - n) * mean + n * score
  sd = (1 - n) * sd + n * Math.pow(score - mean, 2)

  let advantage = (score - mean) / (sd + 1e-8)
  if (i > normalizationWindow) {
    headParams.forEach((v, k) => {
      headParams[k] += alpha / sigma * advantage * noise[k]
    })
  }
}

// higher fitness is better
// (best possible score is 0)
console.log('fitness is: ' + fitness(headParams))