giuse/machine_learning_workbench

# frozen_string_literal: true

# Run as: `bundle exec ruby examples/neuroevolution.rb`

# Make sure the gem is installed first with `gem install machine_learning_workbench`
# Alternatively, add `gem 'machine_learning_workbench'` to your Gemfile if using Bundle,
# followed by a `bundle install`
require 'machine_learning_workbench'
# Workbench shorthands
XNES = WB::Optimizer::NaturalEvolutionStrategies::XNES
FFNN = WB::NeuralNetwork::FeedForward

# Let's address the XOR problem, as it requires nonlinear fitting
XOR = {[0,0] => 0, [1,0] => 1, [0,1] => 1, [1,1] => 0}
# A classic [2,2,1] (2 inputs, 2 hidden neurons, 1 output neurons) feed-forward
# network with nonlinear activations can solve this problem.
# To approximate more complex functions, keep the number of inputs and outputs
# fixed (they depend on the problem) and increase the number and/or size of
# hidden neurons. For example: [2, 10, 7, 4, 1].
# NOTE: If your network grows above few thousands of weights, XNES may be too slow.
# Try using SNES for large shallow networks or BDNES for deep networks.
NET = FFNN.new [2,2,1], act_fn: :logistic
# Note: the process is exactly the same, from instantiation to training, for recurrent
# networks using the class `WB::NeuralNetwork::Recursive`.
# Of course RNNs should be applied to sequential tasks, while XOR is static

# We will search for the network's weights with a black-box optimization algorithm
# This means we will search for arrays of numbers, which need to be scored.
# The scoring process will work as follows: use the numbers as weights for the neural
# network, test the network on classifying the 4 cases of XOR, use that count as the
# score for the weights (original array of numbers).

# Hence the fitness looks as follows:
def fitness weights
  # Each list of weights uniquely defines a neural network
  NET.load_weights weights
  # Activate the network on each of the XOR instances
  # - prediction: the output of the network
  # - observation: correct value, our target
  pred_obs = XOR.map do |input, obs|
    # The network can have an arbitrary number of output neurons
    # Since here we have only one, we extract the value as the output
    output = NET.activate(input)[0]
    # Here we interpret the output as classification
    pred = output > 0.5 ? 1 : 0
    # Finally accumulate prediction-observation pairs
    [pred, obs]
  end
  # To build a score out of this, we count the number of correct classifications
  score = Float(pred_obs.count { |pr, ob| pr == ob })
  # That's it, this will score the weights based on their network's performance
end

# Next comes initializing the black-box stochastic optimization algorithm
# We are searching for the network's weights, this gives us the search space dimensionality
# We'll use XNES as we are working with less than 100 dimensions (weights)
nes = XNES.new NET.nweights, method(:fitness), :max, rseed: 0
# Note BDNES requires `NET.nweights_per_layer` rather than `NET.nweights` in initialization:
# nes = WB::Optimizer::NaturalEvolutionStrategies::BDNES.new NET.nweights_per_layer,
#   method(:fitness), :max, rseed: 10
# The random seed is fixed here to ensure a reproducible behavior
# In a real task, best using an oversized network, more iterations, and try several seeds

# NOTE: In practical applications it is best to delegate parallelization to the fitness
# function instead of computing the fitness of one individual at a time. This can be
# achieved by passing  an objective function defined on a _list_ of weight-lists, and
# setting the `parallel_fit` switch to `true`:
# nes = XNES.new NET.nweights,
#   -> (genotypes) { Parallel.map genotypes, &method(:fitness) },
#   :max, rseed: 0, parallel_fit: true


# Nothing left but to run the optimization algorithm
# Depending on the random seed (read: luck)few epochs here will suffice
50.times { nes.train }
# OK! now remember, `NET` currently holds the weights of the last evaluation
# Let's fetch the best individual found so far
best_fit, best_weights = nes.best
# Let's run them again to check they work
result = fitness best_weights
# Note if you defined a parallel fitness above you'll need instead
# result = fitness([best_weights])[0]
puts "The found network achieves a score of #{result} out of #{XOR.size} in the XOR task"
puts "Weights: #{best_weights.to_a}"
puts "Done!"
# That's it! 18 lines and you got a working neuroevolution algorithm, congrats :)