smbot/snes.lua

-- Separable Natural Evolution Strategies
-- this particular implementation is based on:
-- http://www.jmlr.org/papers/volume15/wierstra14a/wierstra14a.pdf
-- not to be confused with the Super Nintendo Entertainment System.

local assert = assert
local exp = math.exp
local floor = math.floor
local ipairs = ipairs
local log = math.log
local max = math.max
local min = math.min
local sqrt = math.sqrt
local insert = table.insert
local remove = table.remove

local Base = require "Base"

local nn = require "nn"
local normal = nn.normal
local uniform = nn.uniform
local zeros = nn.zeros

local util = require "util"
local argsort = util.argsort
local cdf = util.cdf
local clamp = util.clamp
local normalize_sums = util.normalize_sums
local pdf = util.pdf
local weighted_mann_whitney = util.weighted_mann_whitney

local xnes = require "xnes"
local make_utility = xnes.make_utility

local Snes = Base:extend()

function Snes:init(dims, popsize, base_rate, sigma, antithetic, adaptive)
    -- heuristic borrowed from CMA-ES:
    self.dims = dims
    self.popsize = popsize or 4 + (3 * floor(log(dims)))
    base_rate = base_rate or 3/5 * (3 + log(dims)) / (dims * sqrt(dims))
    self.param_rate = 1.0
    self.sigma_rate = base_rate
    self.covar_rate = base_rate
    self.sigma = sigma or 1
    self.antithetic = antithetic and true or false
    self.adaptive = adaptive == nil and true or adaptive

    if self.antithetic then self.popsize = self.popsize * 2 end

    self.utility = make_utility(self.popsize)

    self.rate_init = self.sigma_rate

    self.mean = zeros{dims}
    self.std = zeros{dims}
    for i=1, self.dims do self.std[i] = self.sigma end

    self.old_asked = {}
    self.old_noise = {}
    self.old_score = {}
    self.new_asked = {}
    self.new_noise = {}

    self.evals = 0
end

function Snes:params(new_mean)
    if new_mean ~= nil then
        assert(#self.mean == #new_mean, "new parameters have the wrong size")
        for i, v in ipairs(new_mean) do self.mean[i] = v end
    end
    return self.mean
end

function Snes:decay(param_decay, sigma_decay)
    if sigma_decay > 0 then
        for i, v in ipairs(self.std) do
            self.std[i] = v * (1 - self.sigma_rate * sigma_decay)
        end
    end
    if param_decay > 0 then
        for i, v in ipairs(self.mean) do
            self.mean[i] = v * (1 - self.param_rate * param_decay * self.std[i])
        end
    end
end

function Snes:ask_once(asked, noise)
    asked = asked or {}
    noise = noise or {}

    for i=1, self.dims do noise[i] = normal() end
    for i, v in ipairs(noise) do asked[i] = self.mean[i] + self.std[i] * v end

    return asked, noise
end

function Snes:ask_twice(asked0, asked1, noise0, noise1)
    asked0 = asked0 or zeros(self.dims)
    asked1 = asked1 or zeros(self.dims)
    noise0 = noise0 or {}
    noise1 = noise1 or {}

    for i=1, self.dims do noise0[i] = normal() end
    noise0.shape = {#noise0}

    for i, v in ipairs(noise0) do
        asked0[i] = self.mean[i] + self.std[i] * v
        asked1[i] = self.mean[i] - self.std[i] * v
    end
    for i, v in ipairs(noise0) do noise1[i] = -v end

    return asked0, asked1, noise0, noise1
end

function Snes:ask(asked, noise)
    -- return a list of parameters for the user to score,
    -- and later pass to :tell().
    self.mixing = false
    if asked == nil then
        asked = {}
        for i=1, self.popsize do asked[i] = zeros(self.dims) end
    end
    if noise == nil then
        noise = {}
        for i=1, self.popsize do noise[i] = zeros(self.dims) end
    end

    if self.antithetic then
        for i=1, self.popsize do
            self:ask_twice(asked[i+0], asked[i+1], noise[i+0], noise[i+1])
        end
    else
        for i=1, self.popsize do
            self:ask_once(asked[i], noise[i])
        end
    end

    self.asked = asked
    self.noise = noise
    return asked, noise
end

function Snes:ask_mix(start_anew)
    -- TODO: refactor and merge with :ask()?
    self.mixing = true
    if start_anew then
        self.old_asked = {}
        self.old_noise = {}
        self.old_score = {}
    end

    -- perform importance mixing.

    local mean_old = self.mean_old or self.mean
    local mean_new = self.mean
    local std_old = self.std_old or self.std
    local std_new = self.std

    local function compute_probabilities(a)
        local prob_new = 0
        local prob_old = 0
        for i, v in ipairs(a) do
            prob_new = prob_new + pdf(v, mean_new[i], std_new[i])
            prob_old = prob_old + pdf(v, mean_old[i], std_old[i])
        end
        return prob_new, prob_old
    end

    local all_asked, all_noise, all_score = {}, {}, {}

    for p=1, #self.old_asked do
        do
            local pp = floor(uniform() * #self.old_asked) + 1
            local a = self.old_asked[pp]
            local prob_new, prob_old = compute_probabilities(a)
            local accept = min(prob_new / prob_old * (1 - self.min_refresh), 1)
            if uniform() < accept then
                --print(("accepted old sample %i with probability %f"):format(pp, accept))
                insert(all_asked, a)
                insert(all_noise, self.old_noise[pp])
                insert(all_score, self.old_score[pp])
            end
        end

        do
            local a, n = {}, {}
            for i=1, self.dims do n[i] = normal() end
            for i, v in ipairs(n) do a[i] = mean_new[i] + std_new[i] * v end
            local prob_new, prob_old = compute_probabilities(a)
            local accept = max(1 - prob_old / prob_new, self.min_refresh)
            if uniform() < accept then
                --print(("accepted new sample %i with probability %f"):format(#all_asked, accept))
                insert(all_asked, a)
                insert(all_noise, n)
                insert(all_score, false)
            end
        end

        -- TODO: early stopping, making sure it doesn't affect performance.
    end

    while #all_asked > self.popsize do
        local pp = floor(uniform() * #all_asked) + 1
        --print(("removing sample %i to fit popsize"):format(pp))
        remove(all_asked, pp)
        remove(all_noise, pp)
        remove(all_score, pp)
    end

    while #all_asked < self.popsize do
        local a, n = {}, {}
        for i=1, self.dims do n[i] = normal() end
        for i, v in ipairs(n) do a[i] = mean_new[i] + std_new[i] * v end
        --print(("unconditionally added new sample %i"):format(#all_asked))
        insert(all_asked, a)
        insert(all_noise, n)
        insert(all_score, false)
    end

    -- split all_ tables back into old_ and new_.
    self.old_asked, self.old_noise, self.old_score = {}, {}, {}
    self.new_asked, self.new_noise = {}, {}
    for i, score in ipairs(all_score) do
        local a, n = all_asked[i], all_noise[i]
        if score ~= false then
            insert(self.old_asked, a)
            insert(self.old_noise, n)
            insert(self.old_score, score)
        else
            insert(self.new_asked, a)
            insert(self.new_noise, n)
        end
    end

    return self.new_asked, self.new_noise
end

function Snes:tell(scored)
    self.evals = self.evals + #scored

    local asked = self.mixing and self.new_asked or self.asked
    local noise = self.mixing and self.new_noise or self.noise
    if self.mixing then
        -- note: modifies, in-place, externally exposed tables.
        for i, v in ipairs(asked) do insert(self.old_asked, v) end
        for i, v in ipairs(noise) do insert(self.old_noise, v) end
        for i, v in ipairs(scored) do insert(self.old_score, v) end
        asked = self.old_asked
        noise = self.old_noise
        scored = self.old_score
    end
    assert(asked and noise, ":tell() called before :ask()")
    assert(#asked == #noise and #asked == #scored, "length mismatch")
    assert(#scored == self.popsize)

    -- TODO: use a proper ranking function.
    local arg = argsort(scored, function(a, b) return a > b end)

    local g_mean = zeros{self.dims}
    local g_std = zeros{self.dims}

--[[
    local utilize = true
    local utility = self.utility

    if utilize then
        utility = {}
        local const = log(self.popsize * 0.5 + 1)
        for i, v in ipairs(arg) do utility[v] = max(const - log(i), 0) end
        normalize_sums(utility)
    else
        utility = normalize_sums(scored, {})
    end
--]]

    for p=1, self.popsize do
        local noise_p = noise[arg[p]]

        for i, v in ipairs(g_mean) do
            g_mean[i] = v + self.utility[p] * noise_p[i]
        end

        for i, v in ipairs(g_std) do
            local n = noise_p[i]
            g_std[i] = v + self.utility[p] * (n * n - 1)
        end
    end

    local step = {}
    for i, v in ipairs(g_mean) do
        step[i] = self.std[i] * v
    end

    self.mean_old = {}
    for i, v in ipairs(self.mean) do
        self.mean_old[i] = v
        self.mean[i] = v + self.param_rate * step[i]
    end

    local otherwise = {}
    self.std_old = {}
    for i, v in ipairs(self.std) do
        self.std_old[i] = v
        self.std[i] = v * exp(self.sigma_rate * 0.5 * g_std[i])
        otherwise[i] = v * exp(self.sigma_rate * 0.75 * g_std[i])
    end

    if self.adaptive then self:adapt(asked, otherwise, self.utility) end

    return step
end

function Snes:adapt(asked, otherwise, qualities)
    local weights = {}
    for p=1, self.popsize do
        local asked_p = asked[p]
        local prob_now = 0
        local prob_big = 0
        for i, v in ipairs(asked_p) do
            prob_now = prob_now + pdf(v, self.mean[i], self.std[i])
            prob_big = prob_big + pdf(v, self.mean[i], otherwise[i])
        end
        weights[p] = prob_big / prob_now
    end

    local u, p = weighted_mann_whitney(qualities, qualities, nil, weights)
    --print(("u, p: %6.3f, %6.3f"):format(u, p))

    if p < 0.5 - 1 / (3 * (self.dims + 1)) then
        self.sigma_rate = 0.9 * self.sigma_rate + 0.1 * self.rate_init
        --print("learning rate -:", self.sigma_rate)
    else
        self.sigma_rate = min(1.1 * self.sigma_rate, 1)
        --print("learning rate +:", self.sigma_rate)
    end
end

return {
    Snes = Snes,
}