NEAT: Simple neuroevolution framework, in Rust

A month or so ago I started working on a neural network implementation in Rust, from scratch. I wasn’t interested in achieving the best performance, or having all the bells and whistles. I had a simple goal of understanding NEAT. For those of you that are not familiar with it, here’s a snippet from the Wikipedia 👇.

a form of artificial intelligence that uses evolutionary algorithms to generate artificial neural networks, parameters, topology and rules

What fascinated me is that the system starts from a totally random set of simple neural networks. Very often, the evolution finds the best architecture and weights in a relatively short amount of time. The first problem I wanted to solve was training the network to be a XOR gate. It managed to do that after a few days of me writing code and figuring out where the bugs were. I was very excited when I found out it evolved the output activation function to be a step function. It figured out the outputs are always round numbers. All by itself. Wow.

Now, to cut the story short, I don’t think I’ll spend more time on it as I’ve achieved my goal. It might be useful for somebody else so I’ve made it public. Here’s a short example.

Problem

If you do a Google search, or look on Youtube, you’ll quickly find out what exactly is a cart pole balancing problem. There is a cart that can only go left and right in a 2D world. A pole is attached to the top of the cart. The goal is to balance the pole for as long as you can. Similar to what children do, balancing a broom in the hand. 🧹

I couldn’t find any existing simulators for this environment so I wrote one myself. You can find its code in the repo that I will link at the end of the post.

Training

I’ve tried to make it super easy to set up a problem, run the evolution, and get a champion neural network out of it. Here’s how to that for the cart pole problem.

There are only 3 arguments when creating a NEAT system:

• number of input neurons
• number of output neurons
• the fitness function that returns a f64

let mut system = NEAT::new(4, 1, |network| {
    let num_simulations = 10;
    let max_steps = 1000;
    let mut env = CartPole::new();

    let mut steps_done = 0;
    let mut fitness = 0.;

    for _ in 0..num_simulations {
        env.reset();

        for _ in 0..max_steps {
            if env.done() {
                break;
            }

            let state = env.state();
            let network_output = network.forward_pass(state.to_vec());
            let env_input = f64::max(-1., f64::min(1., *network_output.first().unwrap()));

            env.step(env_input).unwrap();
            steps_done += 1;
        }

        fitness += env.fitness();
    }

    fitness / num_simulations as f64
});

system.set_configuration(Configuration {
    population_size: 100,
    max_generations: 500,
    stagnation_after: 50,
    node_cost: 1.,
    connection_cost: 1.,
    compatibility_threshold: 2.,
    ..Default::default()
});

system.add_hook(10, |generation, system| {
    println!(
        "Generation {}, best fitness is {}, {} species alive",
        generation,
        system.get_best().2,
        system.species_set.species().len()
    );
});

let (network, fitness) = system.start();

After the system is created the configuration is tweaked for this specific problem. That should help the process to find the best neural networks faster.

In order to see what is happening I’ve added a simple “hook” that runs every 10 generations and gives us some info about the current state.

I’ll omit the code that exports the neural network to a file, but you can find it in the repo. To run the training process on your machine navigate to the examples/cart-pole/ dir and run:

cargo run --release -- train

Balancing the pole

After the training is complete you’ll see a network.bin file is created. We’ll use that to instantiate the network in the simulation. To open the simulation run:

cargo run --release -- visualize

Now drag the network.bin file into the simulation window and watch what happens. Or if you are lazy, watch the video below.

Sometimes the network won’t be able to balance the pole. That happens because the starting parameters are generated randomly, and for some of them it just isn’t possible to put the pole back up.

After a couple of seconds you should see the cart and the pole perfectly in the middle. The network can balance them for eternity. On the other hand, watching a stable system is boring, and for that I’ve added the ability to apply some “wind” by pressing arrow keys.

Future

I suppose I won’t spend more time on this. There are other things I’d like to try. There’s a tiny possibility someone will find this useful. In that case I’ll be happy to chat about it, and potentially find some time if somebody wants to fund it. You can send me a message on Twitter.

This is what I had in mind for the future, even though it probably won’t happen.

• Two poles balancing task example (started it in a different branch)
• Recurrent connections
• Extend the system so it works with both f32 and f64 (might improve performance)
• HyperNEAT
• FS NEAT (feature selection)

All the code is now public and you can find it here.

On to the next project 👋