Accelerating AI Training: A C++ Thread Pool For Lunar Lander

Accelerating AI training: A C++ thread pool for lunar lander

Training AI, especially with evolutionary methods like Genetic Algorithms (GAs), can be computationally intensive. Each generation involves evaluating the fitness of many individuals, often by running them through simulations. To speed up this process for our Lunar Lander AI, I have implemented a C++ thread pool to parallelize the evaluation phase.

In my NNEngineTrainer class, the train_generation method is responsible for assessing the performance of each neural network (individual) in the current GA population. This involves:

1. Iterating through each individual in the population.
2. For each individual, running multiple game simulations (one for each predefined LayoutInfo which specifies starting and landing pad positions).
3. Calculating a fitness score for the individual based on its performance across all layouts.

Without parallelization, these simulations would run sequentially, making each generation take a considerable amount of time, especially with large populations or many layouts.

To address this, I utilize a thread pool, specifically the tp::thread_pool class using a library that provides a straightforward way to manage a pool of worker threads and distribute tasks among them.

// In nn_engine_trainer.cpp
#include "thread/thread_pool.hpp" // Include the thread pool header

template <typename T>
void Training::NNEngineTrainer<T>::train_generation(
    const std::vector<LayoutInfo>& layouts, size_t population_size,
    std::vector<double>& fitness_scores,   // Output
    std::vector<double>& all_member_steps) // Output
{
    // ... (initialization) ...

    tp::thread_pool pool; // Create a thread pool instance

    // Outer loop: Iterate over each member of the population
    for (size_t member_id = 0; member_id < population_size; ++member_id)
    {
        // Define a simulation task for the current member as a lambda
        auto sim_task = [this, member_id, &layouts, &fitness_scores, &all_member_steps, /*...captures for config...*/ ]() {
            GameLogicCpp<T> game_sim(true); // Create a game instance per task
            // ... (configure game_sim) ...

            double total_fitness_for_member = 0.0;
            double total_steps_for_member   = 0.0;

            // Inner loop: Iterate over each layout for the current member
            for (const auto& layout_info : layouts)
            {
                game_sim.reset(static_cast<T>(layout_info.spad_x1), static_cast<T>(layout_info.lpad_x1));
                // ... (run simulation for this layout) ...
                // ... (calculate fitness_for_layout) ...
                total_fitness_for_member += fitness_for_layout;
                total_steps_for_member += static_cast<double>(steps_this_layout_local);
            }

            // Store results (thread-safe due to unique member_id index)
            fitness_scores[member_id]   = total_fitness_for_member;
            all_member_steps[member_id] = total_steps_for_member;
        }; // End of sim_task lambda

        if (this->multithread_) {
            pool.push_task(sim_task); // Submit task to the pool
        } else {
            sim_task(); // Execute synchronously if multithreading is disabled
        }
    } // End of population loop

    if (this->multithread_) {
        pool.wait_for_tasks(); // Wait for all submitted tasks to complete
    }
}

1. Configuration: A multithread_ boolean flag (loaded from config.txt) controls whether the thread pool is used.
2. Pool Creation: Inside train_generation, a tp::thread_pool pool; object is created. By default, this pool initializes with a number of threads equal to the hardware concurrency (e.g., number of CPU cores).
3. Task Definition: For each member_id (each neural network individual), a lambda function sim_task is defined. This lambda encapsulates the entire simulation process for one individual across all specified layouts.
   • Each task creates its own GameLogicCpp<T> game_sim instance. This is important for thread safety, as each thread needs its own independent game state to work with.
4. Task Submission:
   • If multithread_ is true, pool.push_task(sim_task); submits the lambda to the thread pool’s queue. The pool then assigns this task to an available worker thread.
   • If multithread_ is false, sim_task(); is called directly, executing the simulation synchronously in the main thread.
5. Synchronization: After all tasks (one for each population member) have been submitted (if multithreading), pool.wait_for_tasks(); is called. This blocks the main thread until all tasks in the pool have completed their execution. This ensures that all fitness scores are calculated before proceeding to the GA’s selection and evolution steps.
6. Results: The fitness_scores and all_member_steps vectors are populated by each task. Since each task writes to a unique index (member_id), this access is inherently thread-safe without needing explicit locks around vector writes.

• Speed: The most significant benefit is a dramatic reduction in the time taken per generation. If you have N cores, you can (ideally) evaluate N individuals simultaneously.
• Responsiveness: While training is CPU-bound, offloading work to other threads can keep the main thread.
• Scalability: The approach scales well with the number of available CPU cores.

Using multithreading through a d thread pool, the NNEngineTrainer can efficiently train complex neural networks for the Lunar Lander, making the exploration of different GA parameters and network architectures much more feasible.

The code for this implementation is available on Github here.