NeuroEvolution using TensorFlowJS - Part 1

Introduction


This tutorial will walk you through how to implement NeuroEvolution using TensorFlow JS. This NeuralNetwork will learn how to play a simple browser based game that requires a player to jump over blocks and gaps. It is based on the source code found on my GitHub account here https://github.com/dionbeetson/neuroevolution-experiment.

The difference in this solution compared to other solutions out there is that instead of building ML directly into the core game source code (like I saw in a lot of examples); the ML and Game source code are completely decoupled. This simulates a much more realistic real world example, as you would rarely have direct access to the source code when applying ML (you would only have access to specific inputs). I created an an API/SDK for the browser based game to expose some key functionality (the inputs the NeuralNetwork needs), but end state is to use image analysis so this solution is more adaptable to other applications.

References

Although I wrote this codebase myself I took a lot of input/influence from a huge range of great tutorials that already exist. You can see a list of the key tutorials/repos I learned from on the GitHub Readme for this project here.

Let's get started!

Setup

Clone this repo from GitHub:
git clone git@github.com:dionbeetson/neuroevolution-experiment.git
Then run in terminal the below to setup the repository for this tutorial:
npm install

npm run setup-tutorial
This basically moves the js/ai directory to js/ai-source, and leaves everything else in place. If you ever need to refer back to the original source code use js/ai-source as reference.

Now run the below to startup the app:
npm run dev
The browser will load up the application, but you will see 404 errors within the DevTools console. You will though however, have a working browser game. Click 'Play game', and use the spacebar to jump to test things out. Throughout this tutorial you are going to be building a NeuroEvolution implementation to learn how to automatically solve this game.

Quick overview of the UI

  • Left section: Includes base controls for Manual and AI interactions. In order from top to bottom
    • Level: You can select the level (there are currently 3 levels available in the game in order of complexity)
    • Play game: Play an user interacted game
    • Enable drawing: Whether you want to render the drawing on canvas
    • Speed: How fast the game will run
    • Start Evolution: Start the NeuroEvolution
    • Save to localStorage: Save best current NeuralNetwork into localStorage
    • Save to disk: Save best current NeuralNetwork to disk
    • Load from localStorage: Load current NeuralNetwork in local storage into memory and start NeuroEvolution
    • Load from disk: Load a NeuralNetwork from disk into memory and start NeuroEvolution
    • Pause after generation: For debugging, pause after each evolution so you can see / debug it
  • Middle section: The game(s) render within the large middle column
  • Right section: Is used to display debugging information for the NeuroEvolution implementation

This tutorial has 4 key files

For the entire tutorial we will be working within the js/ai/ directory.
  1. ui.js: Controls the UI interactions
  2. NeuralNetwork.js: Is the base neural network
  3. Ai.js: This class instantiates (by default 9) games & corresponding NeuralNetwork instances
  4. Neuroevolution.js: This class will run up to 1500 generations of Ai.js until it successfully passes the selected level

Setup base structure

Let's start to build out enough base logic to be able to predict if we should jump or not using TensorFlowJS.

Create a directory js/ai, and create a file called js/ai/ui.js which will handle basic UI controls. Within this file we want to instantiate the NeuroEvolution class (that we are about to create, and bind a UI button to start the evolution process).

js/ai/ui.js

let neuroEvolution = new NeuroEvolution();

document.querySelector("#btn-ml-start").addEventListener('click', function (event) {
  neuroEvolution.start([]);
}, false);
Now we need to create the skeleton of NeuralNetwork.js, Ai.js & Neuroevolution.js.

For the NeuralNetwork let's just start with a constructor that allows us to send input_node/hidden_node & output_node count arguments.

js/ai/NeuralNetwork.js

class NeuralNetwork {
  constructor(input_nodes, hidden_nodes, output_nodes) {
    // The amount of inputs (eg: player y position, height of next block etc..)
    this.input_nodes = input_nodes;
    // Amount of hidden nodes within the Neural Network)
    this.hidden_nodes = hidden_nodes;
    // The amount of outputs, we will use 2 (will be needed for level 3)
    this.output_nodes = output_nodes;

    // Initialize random weights
    this.input_weights = tf.randomNormal([this.input_nodes, this.hidden_nodes]);
    this.output_weights = tf.randomNormal([this.hidden_nodes, this.output_nodes]);
  }
}
Now we need to create the Ai class. This will include a method that will create 1 game for now and link it with a NeuralNetwork instance. We will start with a start() method to create instances of a NeuralNetwork, a checkGame() method that will be called every 50ms which will call the think() method that we will expand to make the actual prediction to jump or not jump.

js/ai/Ai.js

class Ai {
  #totalGames = 1; // Will increase to 9 games later
  #inputs = 5;
  #neurons = 40;
  #outputs = 2;
  #games = [];
  #gamesRunning = 0;
  #sectionsToSeeAhead = 1;
  #timeTakenDateStart = null;

  start(useImageRecognition, neuralNetworks, completeCallback) {
    this.#timeTakenDateStart = new Date();

    for ( let i = 0; i < this.#totalGames; i++ ) {
      let neuralNetwork;

      if ( undefined !== neuralNetworks && neuralNetworks[i] instanceof NeuralNetwork ) {
        neuralNetwork = neuralNetworks[i];
      } else {
        neuralNetwork = new NeuralNetwork(this.#inputs, this.#neurons, this.#outputs);
      }

      let gameApi;

      if ( useImageRecognition ) {
        gameApi = new GameImageRecognition();
      } else {
        gameApi = new GameApi();
      }

      this.#games[i] = {
        gameApi: gameApi,
        neuralNetwork: neuralNetwork,
        interval: null
      }

      // Debug look ahead
      this.#games[i].gameApi.setHighlightSectionAhead(this.#sectionsToSeeAhead)

      // Start game
      this.#gamesRunning++;
      this.#games[i].gameApi.start();

      this.#games[i].interval = setInterval(this.checkGame.bind(null, this, this.#games, this.#games[i]), 50);
    }
  }

  checkGame(ai, games, game) {
    if( game.gameApi.isSetup() ) {
      ai.think(game);
    }
  }

  think(game) {
  }
}
Now create the NeuroEvolution class that will be responsible for running multiple evolutions of Ai.js until we get an evolution that passes the level.

js/ai/NeuroEvolution.js

class NeuroEvolution {
  #generation = 1;
  #maxGenerations = 1500;
  #useImageRecognition = false;

  start(games, bestPlayerBrainsByFitness) {

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

      for ( let i = 0; i < games.length; i++ ) {
        games[i].gameApi.remove();
      }

      games = undefined;

      this.#generation++;

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

  finishGeneration(games, timeTaken) {
   // TODO
  }

  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 + '%';
  }
}
Now you have your base structure in place.

Reload your browser (http://127.0.0.1:8080/), and click the 'Start evolution' button. You should see 1 game playing, and failing (no jumping), no matter how many times you reload the page. That is because we haven't implemented the prediction logic. Let's implement that using TensorFlowJS in a moment.

Quickly, before that, in the Ai class, let's create 9 games at once instead of 1, as NeuroEvolution will be incredibly slow with only 1 NeuralNetwork at a time. Depending on your computer specs you can increase this to a higher number. I did find that 9 worked quite well though.

js/ai/Ai.js

#totalGames = 9
Reload your browser, and click 'Start evolution' button again. You should see 9 games playing.

Predicting whether to jump or not

To use TensorflowJS to predict to jump, we need to modify a few files. First in Ai.js in the think() method add the below logic in. Essentially we are accessing public information that the JS game exposes in js/game/GameAPI.js (like player position, next block height etc..) and then normalising all inputs to be between 0-1 as TensorFlow JS requires this.

js/ai/Ai.js

think(game) {
  let inputs = [];
  let inputsNormalised = [];

  // Player y
  inputs[0] = (game.gameApi.getPlayerY());
  inputsNormalised[0] = map(inputs[0], 0, game.gameApi.getHeight(), 0, 1);

  // Player x
  inputs[1] = game.gameApi.getPlayerX();
  inputsNormalised[1] = map(inputs[1], inputs[1], game.gameApi.getWidth(), 0, 1);

  let section = game.gameApi.getSectionFromPlayer(this.#sectionsToSeeAhead);

  // 2nd closest section x
  inputs[2] = section.x + section.width;
  inputsNormalised[2] = map(inputs[2], inputs[1], game.gameApi.getWidth(), 0, 1);

  // 2nd closest section y
  inputs[3] = section.y;
  inputsNormalised[3] = map(inputs[3], 0, game.gameApi.getHeight(), 0, 1);

  // 2nd closest section y base
  inputs[4] = section.y + section.height;
  inputsNormalised[4] = map(inputs[4], 0, game.gameApi.getHeight(), 0, 1);

  // Call the predict function in the NeuralNetwork
  let outputs = game.neuralNetwork.predict(inputsNormalised);

  // If input is about 0.5 then jump, if not, stay still
  if ( outputs[0] > 0.5 || outputs[1] > 0.5 ) {
    game.gameApi.jump();
  }
}
We also will need to create the predict() method in the NeuralNetwork class. This method creates an input_layer, hidden_layer and output_layer and returns 2 predictions (which are both numbers between 0-1).

js/ai/NeuralNetwork.js

predict(user_input) {
  let output;
  tf.tidy(() => {
    let input_layer = tf.tensor(user_input, [1, this.input_nodes]);
    let hidden_layer = input_layer.matMul(this.input_weights).sigmoid();
    let output_layer = hidden_layer.matMul(this.output_weights).sigmoid();
    output = output_layer.dataSync();
  });
  return output;
}
Reload your browser, and click the 'start evolution' button.

You will now see 9 games playing, and some should be jumping, and some not. This is completely random. Reload a few times if you need to to ensure this is the behaviour you are seeing.

Once all of your players die you will notice the game stops. As of now, we have used TensorFlow JS to make a prediction about whether to jump or not. This is based on a set of inputs (player X, player Y & section ahead height/Y).

Now we need to add in the NeuroEvolution aspect so that we pick the best progressed games and start a new evolution with those NeuralNetworks (but we will evolve/mutate them slightly before hand).

In NeuroEvolution add a method called calculateFitness(). This method iterates through all games once finished, creates a fitness rating which is just the % progress of the game. It then has logic to increase the fitness of better performing games, so that later on in our logic, there is a higher chance the better performing NeuralNetworks are chosen for the next evolution. Hope that all makes sense.

js/ai/NeuroEvolution.js

calculateFitness(games) {
  for ( let i = 0; i < games.length; i++ ) {
    let game = games[i];
    games[i].fitness = game.gameApi.getProgress() / 100;
    games[i].score = game.gameApi.getScore();
    games[i].progress = game.gameApi.getProgress();
  }

  // The below code makes the better progressed games have a higher fitness so they have a higher chance of being selected for next generation
  games.sort(this.sortByFitness);
  games.reverse();

  let prev = 0;
  for ( let i = 0; i < games.length; i++ ) {
    games[i].fitness = this.#discountRate * prev + games[i].fitness;
    prev = games[i].fitness;
  }

  games.sort(this.sortByFitness);

  return games;
}
In NeuroEvolution.js add a method that will check if at least 1 game has passed the level in the current evolution.

js/ai/NeuroEvolution.js

didAtLeastOneGameCompleteLevel(games) {
  for ( let i = 0; i < games.length; i++ ) {
    if (games[i].gameApi.isLevelPassed() ) {
      return games[i];
    }
  }

  return false;
}
In NeuroEvolution add a method that will pick the exact best player with the highest fitness. We use this to display and keep a track of the best performing NeuralNetwork.

js/ai/NeuroEvolution.js

pickBestGameByActualFitness(games){
  let game;
  let prevFitness = 0;
  for ( let i = 0; i < games.length; i++ ) {
    if (games[i].fitness > prevFitness) {
      game = games[i];
      prevFitness = game.fitness;
    }
  }

  return game;
}
In NeuroEvolution add a method that will pick one (not always the best like the above method) of the best players with the highest fitness

js/ai/NeuroEvolution.js

pickBestGameFromFitnessPool(games) {
  let index = 0;
  let r = random(1);

  while (r > 0 ) {
    if( undefined !== games[index] ) {
      r = r - games[index].fitness;
      index++;
    } else {
      r = 0;
    }
  }
  index--;

  let game = games[index];

  return game;
}
In NeuroEvolution we can now improve the finishGeneration() method with the below logic. This code will:
  1. Calculate fitness for all 9 finished games
  2. Store the best player in a local variable
  3. Add the best player to the bestGames[] array
  4. Ensure we only keep the best 5 games (for memory/performance reasons)
  5. Iterate over all games and mutate the best performing NeuralNetworks into new NeuralNetworks for the next evolution
  6. Start the next NeuroEvolution

js/ai/NeuroEvolution.js

finishGeneration(games, timeTaken) {
  games = this.calculateFitness(games);

  // Did one of the games finish?
  let gamePassedLevel = this.didAtLeastOneGameCompleteLevel(games);

  let bestPlayerByFitness = gamePassedLevel;
  let bestPlayerBrainsByFitness = [];

  if( false === bestPlayerByFitness ){
    bestPlayerByFitness = this.pickBestGameByActualFitness(games);
  }

  this.#bestGames.push(bestPlayerByFitness);
  this.#bestGames.sort(this.sortByFitness);

  // Only keep top 5 best scores
  if( this.#bestGames.length > 5 ) {
    this.#bestGames = this.#bestGames.slice(0, 5);
  }

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

    child = this.mutateNeuralNetwork(bestPlayerA.neuralNetwork.clone());

    bestPlayerBrainsByFitness.push(child);
  }

  this.start(games, bestPlayerBrainsByFitness);
}
Lastly, in NeuroEvolution you will need to add the below variable declarations.

js/ai/NeuroEvolution.js

  #discountRate = 0.95;
  #bestGames = [];
Reload your browser, and click the 'start evolution' button. You should again see 9 games playing, some will also be jumping like last time. However, it still wont progress past the first evolution. This is because we need to wire up Ai.js to call the callback method in the checkGame method and add some logic to check if the game is over. Let's update the checkGame() method to look like the below.

js/ai/Ai.js

checkGame(ai, games, game) {
  if( game.gameApi.isOver() ) {
    clearInterval(game.interval);

    ai.#gamesRunning--;
    document.querySelector("#round-progress").style.width = ((games.length-ai.#gamesRunning)/games.length)*100 + '%';

    if( ai.areAllGamesOver(games) && 0 == ai.#gamesRunning ) {
      let timeTakenDateComplete = new Date();
      let timeTaken = (timeTakenDateComplete - ai.#timeTakenDateStart) / 1000;

      ai.#completeCallback(games, timeTaken);
    }
  } else {
    if( game.gameApi.isSetup() ) {
      ai.think(game);
    }
  }
}
We also need to add the following variable declarations to Ai class.

js/ai/Ai.js

#completeCallback;
Reload your browser, and click the 'start evolution' button.

You should see an error as we didn't declare the method areAllGamesOver(). Let's quickly add the below to Ai

js/ai/Ai.js

areAllGamesOver(games) {
  for ( let i = 0; i < this.#totalGames; i++ ) {
    if( false == games[i].gameApi.isOver() ) {
      return false;
    }
  }

  return true;
}
... Another error, we forgot to set the completeCallback. In Ai, create a constructor up the top after the variable declaration.

js/ai/Ai.js

constructor(completeCallback) {
  this.#completeCallback = completeCallback;
}
Reload your browser, and click the 'start evolution' button.

Ok, we are making progress, but we are now missing the clone() method. We need to be able to clone NeuralNetworks after each generation so we don't interfere with previous NeuralNetwork objects.

In NeuralNetwork class add the clone() method as well as the dispose() method, as that is called from within the clone() method.

js/ai/NeuralNetwork.js

clone() {
  return tf.tidy(() => {
    let clonie = new NeuralNetwork(this.input_nodes, this.hidden_nodes, this.output_nodes);
    clonie.dispose();
    clonie.input_weights = tf.clone(this.input_weights);
    clonie.output_weights = tf.clone(this.output_weights);
    return clonie;
  });
}

dispose() {
  this.input_weights.dispose();
  this.output_weights.dispose();
}
Reload your browser, and click the 'start evolution' button. You will see we are close now, we are just missing the mutateNeuralNetwork() method. This method is responsible for slightly tweaking the cloned NeuralNetwork so that you get a slightly different result in the next evolution. Let's go ahead and add the mutateNeuralNetwork() method into the NeuroEvolution class. I'm not going to go into more detail on this, but you can google it and find this method in a few other examples out there. I also wasn't the original author of this, I just adjusted it to work for my implementation. See GitHub for reference to original source code.

js/ai/NeuralEvolution.js

mutateNeuralNetwork(b) {
  function fn(x) {
    if (random(1) < 0.05) {
      let offset = randomGaussian() * 0.5;
      let newx = x + offset;
      return newx;
    }
    return x;
  }

  let neuralNetwork = b.clone();
  let ih = neuralNetwork.input_weights.dataSync().map(fn);
  let ih_shape = neuralNetwork.input_weights.shape;
  neuralNetwork.input_weights.dispose();
  neuralNetwork.input_weights = tf.tensor(ih, ih_shape);

  let ho = neuralNetwork.output_weights.dataSync().map(fn);
  let ho_shape = neuralNetwork.output_weights.shape;
  neuralNetwork.output_weights.dispose();
  neuralNetwork.output_weights = tf.tensor(ho, ho_shape);
  return neuralNetwork;
}
Reload your browser, and click the 'start evolution' button. Now after each of the 9 games fail the first time, the UI will start a new evolution. If you let this go for a few minutes you should see the game get solved - it normally takes around 30-50 generations with this current implementation.

Congratulations!

This is a very simple implementation of NeuroEvolution. In Part 2, we will go through making improvements to really speed things up including:
  • Implementing crossover within TensorFlowJS
  • Using the best progressing NeuralNetwork to ensure faster NeuroEvolution
  • Saving and loading models
  • Adding debugging information

Source code

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

Part 2

Continue on to NeuroEvolution using TensorFlowJS - Part 2

0 comments:

Post a Comment