Evolution Strategy Optimizer

es.ml

(*
The Evolution Strategy Optimizer is a Hill-Climbing routine based on stochastic gradient
estimation (not to be confused with stochastic gradient descent optimization). The basic idea
is to perturb the current best parameter value using Gaussian noise and explore the function
values of the blackbox objective in the vicinity of the parameters. From these explorations one
can estimate the gradient and perform a gradient update step.

More information can be found in:
- T. Salimans et al. Evolution Strategies as a Scalable Alternative to Reinforcement Learning,
    arXiv:1703.03864, https://arxiv.org/pdf/1703.03864.pdf
- D. Wierstra et al. Natural Evolution Strategies, JMLR 15 (2014) 949-980,
    http://www.jmlr.org/papers/volume15/wierstra14a/wierstra14a.pdf
*)
open Owl
open Core


let rosenbrock ~param:p ~inp:x =
  -1. *. ((p.(0) -. x.(0))**2. +. p.(1) *. (x.(1) -. x.(0)**2.)**2.)


let es_grad ~batch_size:bs ~f:f ~inp:(inp : float array) ~scale:scale =
  let dim = Array.length inp in
  let center = Mat.of_array inp 1 dim in
  let all_perts = Mat.gaussian bs dim in
  let displaced_perts = Mat.(center + (scale $* all_perts)) in
  let all_func_vals = Mat.map_rows (fun m -> f ~inp:(m |> Mat.to_array)) displaced_perts in
  let mu = Stats.mean all_func_vals in
  let sigma = Stats.std all_func_vals in
  let norm_func_vec = Mat.( ((of_array all_func_vals bs 1) -$ mu) /$ sigma ) in
  Mat.(norm_func_vec * all_perts) |> Mat.fold ~axis:0 (+.) 0.


let update_params lr bs scale f params =
  let par_vec = Mat.of_array (Array.copy params) 1 (Array.length params) in
  let grad = es_grad ~batch_size:bs ~f:f ~inp:params ~scale:scale in
  let alpha = lr /. (scale *. (bs |> float_of_int)) in
  Mat.(par_vec  + (alpha $* grad)) |> Mat.to_array


let optimize ~lr:lr ~bs:bs ~scale:scale  ~f:f ~x0:x0 ~epochs:epochs =
  let vals = ref [|f ~inp:x0|] in
  let params = ref [|(Array.copy x0)|] in
  let p = ref (Array.copy x0) in
  let n_step = ref 0 in
  while !n_step <= epochs do
    incr n_step;
    p := (update_params lr bs scale f !p);
    params := Array.append !params [|!p|];
    vals := Array.append !vals [|(f ~inp:!p)|];
  done;
  (!vals, !params)


let br_rosenbrock x y =
  let xm1 = Mat.(x -$ 1.) in
  let add0 = Mat.(xm1 * xm1) in
  let fac1 = Mat.(y - x * x) in
  let add1 = Mat.( 100. $* (fac1 * fac1) ) in
  Mat.(add0 + add1)


let plot_opt ~trace:tr ~loc:loc =
  let x, y = Mat.meshgrid (-2.) 2. (-0.5) 2. 200 200 in
  let z = br_rosenbrock x y  |> Mat.clip_by_value ~amin:0. ~amax:5. in
  let h = Plot.create loc in
  Plot.(contour ~h x y z);
  Plot.(scatter ~h ~spec:[ RGB (255,255,0); Marker "#[0x2666]"; MarkerSize 5. ] (Mat.col tr 0) (Mat.col tr 1));
  Plot.output h

let run_example () =
  let f = rosenbrock ~param:[|1.;100.|] in
  let (vals, params) = optimize ~f:f ~x0:[|-0.23; 1.06|] ~epochs:5000 ~lr:0.00001 ~bs:100 ~scale:0.0001 in
  let loc = "rosenbrock.png" in
  let tr = Mat.of_arrays params in
  plot_opt ~loc:loc ~trace:tr

The output will create a plot similar to the one below

The Rosenbrock function is one of the typical functions used to evaluate optimizers