NeuroEvolution using TensorFlow JS - Part 2

Introduction


This tutorial is part 2, if you have not completed NeuroEvolution using TensorFlowJS - Part 1, I highly recommend you do that first. As it will:
  • Explain how to setup the codebase
  • How to run the local dev environment
  • How to setup the basic NeruoEvolution implementation

In this tutorial we will be:
  • Adding visual debugging information
  • Adding more inputs for TensorFlow to use to predict jumping
  • Implement crossover within TensorFlowJS
  • Use the best progressing NeuralNetwork to ensure faster NeuroEvolution
  • Add saving and loading of NeuralNetwork models

This tutorial is also based on the source code found on my GitHub account here https://github.com/dionbeetson/neuroevolution-experiment.

Let's get started!

Add in visual/debugging

We will be adding in useful debugging information to the right hand column and to the game itself.

Chart the progress of each evolution

We are going to use ChartJS to track the progress (as a percentage) and the score for the best game of each evolution. Create a file called js/ai/NeuroEvolutionChart.js and copy the below logic into it. It is a basic ChartJS implementation that I won't go into detail in this tutorial.

js/ai/NeuroEvolutionChart.js

class NeuroEvolutionChart {
  constructor() {
    this.chart = null;
    this.config = {
    type: 'line',
    data: {
     labels: [],
     datasets: [{
      label: '',
      backgroundColor: '#999999',
      borderColor: '#EEEEEE',
          borderWidth: 3,
      data: [],
      fill: false,
     },{
      label: '',
      backgroundColor: '#91C7B1',
      borderColor: '#A1CFBC',
          borderWidth: 2,
      data: [],
      fill: false,
     }]
    },
    options: {
     responsive: true,
     title: {
      display: false
     },
     tooltips: {
      mode: 'index',
      intersect: false,
     },
        legend: {
          display: false
        },
     hover: {
      mode: 'nearest',
      intersect: true
     },
     scales: {
      xAxes: [{
       display: false,
       scaleLabel: {
        display: true,
        labelString: 'Round'
       }
      }],
      yAxes: [{
       display: true,
            ticks: {
              beginAtZero: true,
              min: 0,
              max: 100
            },
       scaleLabel: {
        display: false,
        labelString: 'Progress'
       }
      },{
       display: true,
            position: 'right',
            ticks: {
              beginAtZero: true,
              min: 0,
              max: 200
            },
       scaleLabel: {
        display: false,
        labelString: 'Score'
       }
      }]
     }
    }
   };

    this.setup();
  }

  setup(){
    var ctx = document.getElementById('trend').getContext('2d');
    this.chart = new Chart(ctx, this.config);
  }
}
Now, add an update() method that will be responsible for pushing the data points onto the chart. It will also only keep 12 data points (the very first evolution, and the latest 11 evolutions). This is to avoid clutter on the page.

js/ai/NeuroEvolutionChart.js

update(progress, score) {
  // Purge middle data (to avoid clutter)
  if( this.config.data.labels.length > 12 ) {
    this.config.data.labels.splice(1, 1);
    this.config.data.datasets[0].data.splice(1, 1);
    this.config.data.datasets[1].data.splice(1, 1);
  }

  this.config.data.labels.push((this.config.data.labels.length+1))
  this.config.data.datasets[0].data.push(progress);
  this.config.data.datasets[1].data.push(score);

  this.chart.update();
}
Back in the NeuroEvolution class add the following variable declaration to the top

js/ai/NeuroEvolution.js

#neuroEvolutionChart = new NeuroEvolutionChart();
Then in the same file in the finishGeneration(games, timeTaken) method, right after the line this.#bestGames.sort(this.sortByFitness); add in.

js/ai/NeuroEvolution.js

// Update UI - Chart
this.#neuroEvolutionChart.update(bestPlayerByFitness.progress, bestPlayerByFitness.score);
Reload your browser, and click the 'start evolution' button. After each evolution you will start to see the chart in the top right update.

Additional debugging visuals

Let's quickly add in some more useful debug information which includes:
  • The best game progress to date
  • Track the best score from each evolution
  • A progress bar in the header to track round and generation progress

In the NeuroEvolution class add the below methods

js/ai/NeuroEvolution.js

updateUIaddBestGenerationToBestScore( pickBestPlayerByFitness, timeTaken ) {
  let bestScore = document.createElement("li");
  bestScore.innerHTML = pickBestPlayerByFitness.score + ' (' + pickBestPlayerByFitness.progress.toFixed(1) + '%) (' + pickBestPlayerByFitness.fitness.toFixed(3) + ') (' + timeTaken + 's)';
  this.#bestScores.insertBefore(bestScore, document.querySelector("li:first-child"));
}

updateUIRoundInformation() {
  document.querySelector("#round-current").innerHTML = this.#generation;
  document.querySelector("#round-total").innerHTML = this.#maxGenerations;
  document.querySelector("#round-progress").style.width = '0%';
  document.querySelector("#generation-progress").style.width = (this.#generation/this.#maxGenerations)*100 + '%';
}

updateUIBestPlayerScore( bestGame ) {
  document.querySelector("#best-player-score").innerHTML = bestGame.score + " points (" + bestGame.progress.toFixed(1) + "%)";
}
In the first line of the start() method add:

js/ai/NeuroEvolution.js

processGeneration(games, bestPlayerBrainsByFitness) {
  this.updateUIRoundInformation();
  ...
Then right after where you added in this.#neuroEvolutionChart.update... add in:

js/ai/NeuroEvolution.js

// Update UI
this.updateUIaddBestGenerationToBestScore(bestPlayerByFitness, timeTaken);
this.updateUIBestPlayerScore(this.#bestGames[0]);
this.updateUIRoundInformation();
Lastly, add the following variable declaration to the top

js/ai/NeuroEvolution.js

#bestScores = null;
and the following constructor

js/ai/NeuroEvolution.js

constructor() {
  this.#bestScores = document.querySelector("#best-scores");
}
Reload your browser, and click the 'start evolution' button. You should start seeing debug information in your right hand column.

Another useful debug is to visualise the inputs that are being used to predict the jump() action.

Back in the Ai class within the think() method, just before the game.brain.predict() line, add in this code block to push debug points into the game canvas.

js/ai/Ai.js

game.gameApi.setDebugPoints([
  {
    x: 0,
    y: inputs[0]
  },
  {
    x: inputs[2],
    y: 0
  },
  {
    x: 0,
    y: inputs[3]
  },
  {
   x: 0,
   y: inputs[4]
  }
]);
Reload your browser, and click the 'start evolution' button. You should start seeing crosshairs for the inputs being used by TensorFlow JS on every prediction call.

Sometimes it is useful to be able to 'pause the game before the next evolution begins. You will see a checkbox in the left column labelled 'Pause after generation'. Let's wire it up, as it currently is not.

In ui.js add in the below change handler to capture user interaction.

js/ai/ui.js

document.querySelector("#ml-pause-before-next-generation").addEventListener("change", function() {
  neuroEvolution.pauseBeforeNextGeneration = this.checked;

  if( false == neuroEvolution.pauseBeforeNextGeneration && neuroEvolution.pausedGames.length > 0 && neuroEvolution.pausedBestNeuralNetworksByFitness.length > 0 ) {
    neuroEvolution.start(neuroEvolution.pausedGames, neuroEvolution.pausedBestNeuralNetworksByFitness);
  }
});

neuroEvolution.pauseBeforeNextGeneration = document.querySelector("#ml-pause-before-next-generation").checked;
This isn't enough though, we need to build in the logic to actually pause the game. We are basically going to wrap the ai.start(bestPlayerBrainsByFitness); call within the start() method in NeuroEvolution.js in an IF statement based on if the pause checkbox is checked. If it is checked, we will then store the current games within a local class variable so we can retrieve it once the checkbox is unchecked.

The completed method looks like this now.

js/ai/NeuroEvolution.js

start(games, bestPlayerBrainsByFitness) {
  this.updateUIRoundInformation();

  if ( this.#generation < this.#maxGenerations ) {

    if( false == this.#pauseBeforeNextGeneration ){
      for ( let i = 0; i < games.length; i++ ) {
        games[i].gameApi.remove();
      }

      games = undefined;

      this.#pausedGames = [];
      this.#pausedBestNeuralNetworksByFitness = [];

      this.#generation++;

      const ai = new Ai(this.finishGeneration.bind(this));
      ai.start(this.#useImageRecognition, bestPlayerBrainsByFitness);

    } else {
      this.#pausedGames = games;
      this.#pausedBestNeuralNetworksByFitness = bestPlayerBrainsByFitness;

      for ( let i = 0; i < games.length; i++ ) {
        games[i].gameApi.show();
      }
    }
  } else {
    this.enableSpeedInput();
  }
}
Then add in the following variable declarations in the NeuroEvolution class.

js/ai/NeuroEvolution.js

#pausedGames = [];
#pausedBestNeuralNetworksByFitness = [];
#pauseBeforeNextGeneration = false;
And then add in the required getters/setters in the NeuroEvolution class.

js/ai/NeuroEvolution.js

get pauseBeforeNextGeneration() {
  return this.#pauseBeforeNextGeneration;
}

set pauseBeforeNextGeneration( pauseBeforeNextGeneration ) {
  this.#pauseBeforeNextGeneration = pauseBeforeNextGeneration;
}

get pausedGames() {
  return this.#pausedGames;
}

get pausedBestNeuralNetworksByFitness() {
  return this.#pausedBestNeuralNetworksByFitness;
}

get bestGames() {
  return this.#bestGames;
}

set pauseBeforeNextGeneration( pauseBeforeNextGeneration ) {
  this.#pauseBeforeNextGeneration = pauseBeforeNextGeneration;
}

Speeding up NeuroEvolution

Ok that is debugging done for now. Let's move onto improving our NeuroEvolution to solve this game a little faster. This involves 3 things:
  1. Adding more inputs for TensorFlow to use to predict jumping
  2. After each evolution randomly choosing to use the 'best game to date', as this can keep things on track if we have a single bad evolution randomly (trust me it happens)
  3. Implementing a crossover method so instead of mutating only 1 of the best games, we take 2 of the best games and breed a child NeuralNetwork from their 2 results (which is called a crossover)
  4. If a game passes the level, don't mutate or crossover, as we already have a successful NeuralNetwork

Adding more inputs for TensorFlow JS

In the Ai class we are going to add in 3 more inputs for (isPlayerJumping, playerVelocity & canPlayerJump). Within the think() method after the existing 5 inputs, add in these additional inputs.

js/ai/Ai.js

// Is player jumping
inputs[5] = (game.gameApi.isPlayerJumping() ? 1 : 0);
inputsNormalised[5] = map(inputs[5], 0, 1, 0, 1);

// Player velocity
inputs[6] = (game.gameApi.getPlayerVelocity() ? 1 : 0);
inputsNormalised[6] = map(inputs[6], -1.1, 1.1, 0, 1);

// Can play jump?
inputs[7] = (game.gameApi.canPlayerJump() ? 1 : 0);
inputsNormalised[7] = map(inputs[7], 0, 1, 0, 1);
Then change the variable declaration for inputs in Ai class to 8

js/ai/Ai.js

#inputs = 8; 

Implement crossover within TensorFlowJS

In the NeuroEvolution class within the completeGeneration() method, we are going to modify the breeding section to make it a little smarter. Make the for loop look like this.

js/ai/NeuroEvolution.js

// Breeding
for (let i = 0; i < games.length; i++) {
  let bestPlayerA = this.pickBestGameFromFitnessPool(games);
  let bestPlayerB = this.pickBestGameFromFitnessPool(games);
  let child;

  child = this.mutateNeuralNetwork(this.crossoverNeuralNetwork(bestPlayerA.neuralNetwork.clone(), bestPlayerB.neuralNetwork.clone()));
 
  bestPlayerBrainsByFitness.push(child);
}
Now we need to add in the crossover method to the same file.

js/ai/NeuroEvolution.js

crossoverNeuralNetwork(neuralNetworkOne, neuralNetworkTwo) {
  let parentA_in_dna = neuralNetworkOne.input_weights.dataSync();
  let parentA_out_dna = neuralNetworkOne.output_weights.dataSync();
  let parentB_in_dna = neuralNetworkTwo.input_weights.dataSync();
  let parentB_out_dna = neuralNetworkTwo.output_weights.dataSync();

  let mid = Math.floor(Math.random() * parentA_in_dna.length);
  let child_in_dna = [...parentA_in_dna.slice(0, mid), ...parentB_in_dna.slice(mid, parentB_in_dna.length)];
  let child_out_dna = [...parentA_out_dna.slice(0, mid), ...parentB_out_dna.slice(mid, parentB_out_dna.length)];

  let child = neuralNetworkOne.clone();
  let input_shape = neuralNetworkOne.input_weights.shape;
  let output_shape = neuralNetworkOne.output_weights.shape;

  child.dispose();

  child.input_weights = tf.tensor(child_in_dna, input_shape);
  child.output_weights = tf.tensor(child_out_dna, output_shape);

  return child;
}
Reload your browser, and click the 'start evolution' button. You should start seeing games solve themselves a bit faster in fewer generations (normally around 20).

10% chance to crossover best game

Back in the NeuroEvolution class, within breeding section, let's add one final update so that on random occasions (10% chance), we do a crossover with the best game to date. Make the breeding section look like this

js/ai/NeuroEvolution.js

// Breeding
for (let i = 0; i < games.length; i++) {
  let bestPlayerA = this.pickBestGameFromFitnessPool(games);
  let bestPlayerB = this.pickBestGameFromFitnessPool(games);
  let bestPlayerC = this.#bestGames[0];
  let child;

  if ( random(1) < 0.1) {
    const ai = new Ai();
    let bestPlayerD = new NeuralNetwork(ai.inputs, ai.neurons, ai.outputs);
    child = this.mutateNeuralNetwork(this.crossoverNeuralNetwork(bestPlayerC.neuralNetwork.clone(), bestPlayerD));
  } else {
    child = this.mutateNeuralNetwork(this.crossoverNeuralNetwork(bestPlayerA.neuralNetwork.clone(), bestPlayerB.neuralNetwork.clone()));
  }

  bestPlayerBrainsByFitness.push(child);
}
We also need to expose getters in the Ai class for the input, neuron & output variables. Add the following methods:

js/ai/Ai.js

get inputs(){
  return this.#inputs;
}

get neurons(){
  return this.#neurons;
}

get outputs(){
  return this.#outputs;
}
Reload your browser, and click the 'start evolution' button. You should start seeing games solve themselves a bit faster in fewer generations (normally around 10-15).

Don't mutate on success

Now wrap the Breed functionality in an IF statement to ensure that if a game passes the level, we don't mutate or crossover, we just continue using the successful NeuralNetwork.

js/ai/NeuroEvolution.js

if ( false != gamePassedLevel ) {
  for ( let i = 0; i < games.length; i++ ) {
    if (games[i].gameApi.isLevelPassed() ) {
      games[i].neuralNetwork.save('neuralNetwork');
      for (let ii = 0; ii < games.length; ii++) {
          bestPlayerBrainsByFitness.push(games[i].neuralNetwork.clone());
        }
      }
    }

    console.log('Level Passed:', this.#bestGames[0], this.#bestGames.length, this.#bestGames);
    this.start(games, bestPlayerBrainsByFitness);
  } else {
    // Breeding
    for (let i = 0; i < games.length; i++) {
      let bestPlayerA = this.pickBestGameFromFitnessPool(games);
      let bestPlayerB = this.pickBestGameFromFitnessPool(games);
      let bestPlayerC = this.#bestGames[0];
      let child;

      if ( random(1) < 0.1) {
        const ai = new Ai();
        let bestPlayerD = new NeuralNetwork(ai.inputs, ai.neurons, ai.outputs);
        child = this.mutateNeuralNetwork(this.crossoverNeuralNetwork(bestPlayerC.neuralNetwork.clone(), bestPlayerD));
      } else {
        child = this.mutateNeuralNetwork(this.crossoverNeuralNetwork(bestPlayerA.neuralNetwork.clone(), bestPlayerB.neuralNetwork.clone()));
      }

      bestPlayerBrainsByFitness.push(child);
    }
 
   this.start(games, bestPlayerBrainsByFitness);
}
Reload your browser, and click the 'start evolution' button. Once enough evolutions process to pass the level, you will notice that all future games in evolutions will start passing consistently.

This is a far as we will take this tutorial in regards to evolving and speeding up our NeuroEvolution logic. Next we will move onto saving/loading our NeuralNetwork models.

Saving & Loading our NeuralNetwork models

The last piece to implement is the saving and loading functionality of our NeuralNetwork. This is so once you have a NeuralNetwork that works you can save it and load it back when needed. We will be implementing the:
  1. Ability to automatically save a NeuralNetwork that passes to localStorage
  2. Ability to manually save best NeuralNetwork to localStorage
  3. Ability to manually save best NeuralNetwork to disk
  4. Ability to load from local storage
  5. Ability to load from disk

First we need to add the ability to save within the NeuralNetwork.js file. Add these 2 methods

js/ai/NeuralNetwork.js

stringify() {
  let neuralNetworkToSave = {
    input_weights: this.input_weights.arraySync(),
    output_weights: this.output_weights.arraySync()
  };

  return JSON.stringify(neuralNetworkToSave);
}

save( key ) {
  localStorage.setItem(key, this.stringify());
}

Automatically save a NeuralNetwork that passes to localStorage

In the NeuroEvolution. class within the method finishGeneration(). Add in the below to the IF statement when a game level is passed.

js/ai/NeuroEvolution.js

games[i].neuralNetwork.save('neuralNetwork');
So it should look like this:

js/ai/NeuroEvolution.js

for ( let i = 0; i < games.length; i++ ) {
  if (games[i].gameApi.isLevelPassed() ) {
    games[i].neuralNetwork.save('neuralNetwork');
    for (let ii = 0; ii < games.length; ii++) {
      bestPlayerBrainsByFitness.push(games[i].neuralNetwork.clone());
    }
  }
}

console.log('Level Passed:', this.#bestGames[0], this.#bestGames.length, this.#bestGames);
this.start(games, bestPlayerBrainsByFitness);
Your game will now automatically save the NeuralNetwork model to localStorage key neuralNetwork if it passes the level.

Manually save best NeuralNetwork to localStorage

To manually save to local storage add the below code to ui.js

js/ai/ui.js

document.querySelector("#btn-ml-save-neuralnetwork-localstorage").addEventListener('click', function (event) {
  if( neuroEvolution.bestGames.length > 0 ) {
    neuroEvolution.bestGames[0].neuralNetwork.save('neuralNetwork')
  }
});

Manually save best NeuralNetwork to disk

To manually save to disk you can add the below code to ui.js

js/ai/ui.js

document.querySelector("#btn-ml-save-neuralnetwork-disk").addEventListener('click', function (event) {
  if( neuroEvolution.bestGames.length > 0 ) {
    let neuralNetworkJSON = neuroEvolution.bestGames[0].neuralNetwork.stringify();
    let a = document.createElement("a");
    let file = new Blob([neuralNetworkJSON], {type: 'application/json'});
    a.href = URL.createObjectURL(file);
    a.download = 'neuralNetwork.json';
    a.click();
  }
});

Loading NeuralNetwork Models

To load from local storage or disk we need to create a method to handle loading a NeuralNetwork into memory. Create the below method in ui.js (there is probably a better file location for this method now that I think about it).

js/ai/ui.js

const loadNeuralNetwork = (neuralNetworkData) => {
  const ai = new Ai();
  let neuralNetwork = new NeuralNetwork(ai.inputs, ai.neurons, ai.outputs);
  neuralNetwork.dispose();

  neuralNetwork.input_weights = tf.tensor(neuralNetworkData.input_weights);
  neuralNetwork.output_weights = tf.tensor(neuralNetworkData.output_weights);

  let neuralNetworks = [];
  for ( let i = 0; i < ai.totalGames; i++ ) {
    neuralNetworks.push(neuralNetwork);
  }

  let games = neuroEvolution.pausedGames;
  neuroEvolution.pausedGames;
  neuroEvolution.reset();

  neuroEvolution.start(games, neuralNetworks);
};

We also need to add the below method into the Ai class.

js/ai/Ai.js

get totalGames(){
  return this.#totalGames;
}

Load from local storage

Now to load from localStorage into memory, add the below into ui.js

js/ai/ui.js

document.querySelector("#btn-ml-load-neuralnetwork-localstorage").addEventListener('click', function (event) {
  cachedLevelSections = [];
  let games = neuroEvolution.pausedGames;
  neuroEvolution.pausedGames;
  neuroEvolution.reset();

  let neuralNetworkData = localStorage.getItem('neuralNetwork');

  loadNeuralNetwork(JSON.parse(neuralNetworkData));
});
We also need to add the reset method into the NeuroEvolution class.

js/ai/NeuroEvolution.js

reset() {
  this.#bestGames = [];
  this.#generation = 1;
  this.#pausedGames = [];
  this.#pausedBestNeuralNetworksByFitness = [];
  this.#pauseBeforeNextGeneration = false;
}

Load from disk

Now to load from disk into memory, add the below into ui.js

js/ai/ui.js

document.querySelector("#btn-ml-load-neuralnetwork-disk").addEventListener('change', function (event) {
  let file = this.files[0];
  let reader = new FileReader()
  let textFile = /application\/json/;
  let fileText = '';

  if (file.type.match(textFile)) {
    reader.onload = function (event) {
       let importedNeuralNetwork = JSON.parse(event.target.result);

       loadNeuralNetwork(importedNeuralNetwork);
    }
  }

  reader.readAsText(file);
});
Reload your browser, and click the 'start evolution' button. You will notice if you inspect Chrome DevTools you will see a localStorage key called 'brain' once a level is passed.


You can also now also use the buttons in the left hand menu to save and load the best NeuralNetwork to disk.

Within the GitHub codebase within the models/ directory you will see 3 existing models that can be loaded up to pass levels 1, 2 & 3.

Try it yourself and get your NeuroEvolution implementation to pass Level 1, then Level 2, then finally Level 3.

I get consistent results along the lines of:

  • Level 1: Takes 5-20 generations
  • Level 2: Takes 5-20 generations
  • Level 3: Takes 40-700 generations (as it has to learn to jump blocks and gaps).

Congratulations!

I hope this tutorial was useful for you to learn how to use TensorFlowJS to build a NeuroEvolution implementation. Any questions leave them in the comments below, or tweet me on twitter at @dionbeetson

Source code

All source code for this tutorial can be found here https://github.com/dionbeetson/neuroevolution-experiment.

Part 3 - coming soon...

Coming soon and will be using image recognition to solve the same browser based game.

0 comments:

Post a Comment