smbot/snes.lua

310 lines
8.5 KiB
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 abs = math.abs
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 Snes = Base:extend()
function Snes:init(dims, popsize, base_rate, sigma, antithetic)
-- 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
if self.antithetic then self.popsize = self.popsize * 2 end
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
local mean_new = self.mean
local std_old = self.std_old or self.std
local std_new = self.std
self.new_asked = {}
self.new_noise = {}
local marked = {}
for p=1, min(#self.old_asked, self.popsize) do
local a = self.old_asked[p]
-- TODO: cache probs?
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
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(p, accept))
else
-- insert in reverse as not to screw up
-- the indices when removing later.
insert(marked, 1, p)
end
end
for _, p in ipairs(marked) do
remove(self.old_asked, p)
remove(self.old_noise, p)
remove(self.old_score, p)
end
while #self.old_asked + #self.new_asked < self.popsize do
local a = {}
local 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
-- can't cache here!
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
local accept = max(1 - prob_old / prob_new, self.min_refresh)
if uniform() < accept then
insert(self.new_asked, a)
insert(self.new_noise, n)
--print(("accepted new sample %i with probability %f"):format(0, accept))
end
end
return self.new_asked, self.new_noise
end
function Snes:tell(scored)
self.evals = self.evals + #scored
local asked = self.asked
local noise = self.noise
if self.mixing then
asked = self.old_asked
noise = self.old_noise
-- note that these modify tables referenced externally in-place.
for i, v in ipairs(self.new_asked) do insert(asked, v) end
for i, v in ipairs(self.new_noise) do insert(noise, v) end
for i, v in ipairs(scored) do insert(self.old_score, v) end
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
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[p]
for i, v in ipairs(g_mean) do
g_mean[i] = v + utility[p] * noise_p[i]
end
for i, v in ipairs(g_std) do
local n = noise_p[i]
g_std[i] = v + utility[p] * (n * n - 1)
end
end
local step = {}
for i, v in ipairs(g_mean) do
step[i] = self.std[i] * v
end
for i, v in ipairs(self.mean) do
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
self:adapt(asked, otherwise, utility)
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 p = weighted_mann_whitney(qualities, qualities, nil, weights)
--print("p:", 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,
}