Parallel self-play using Tokio

[jonbinney]

Summary of changes

• multi-threaded self-play in rust, using the tokio asynchronous library
• evaluation cache in rust self-play. This is shared across all worker threads and all games.
• "virtual loss" so that MCTS search can be run multiple times on a single tree without waiting for each evaluation
• After choosing a move, the rust alphazero agent now reuses that part of the MCTS subtree to save work on the next move.
• Renamed self-play config variables that control parallelism so that they make sense for both the rust and the python multithread/multiprocess designs
• Updated the rust package to use the 2024 edition of rust instead of the 2021 edition.
• included the design and plan docs that the claude superpowers plugin created. @adamantivm will like this :-)

Instructions for trying this out

• Download this config file
• Create a python virtual environment, activate it, and run pip install -r deep_quoridor/requirements.txt
• In the deep_quoridor/rust directory, run cargo build --release --all-features --bin selfplay
• From the root directory of the repo, run python -O deep_quoridor/src/train_v2.py <path to downloaded config.yaml> --overrides run_id=<unique name for the run>

Multi-threaded rust self-play with Tokio

There are now a fixed number of worker threads created, and any of them can work on the MCTS tree of whichever game is ready. This lets you choose a number of threads based on the number of cores in the CPU, and choose a number of parallel games to take advantage of the efficiency of batch evaluations. Interestingly, Claude implemented this in a way that immediately picks up a new model version, even mid-game.

Virtual loss

We want to be able to run MCTS searches on a tree without waiting for the evaluator each time. The problem is that every search will choose the same path through the tree, since the total value and visit counts will not have changed. Virtual loss fixes this by having each search temporarily back-propagate a fixed number of losses (defaults to 3, which alphazero used). This makes that path through the tree less attractive to the next thread. Once the evaluations are ready, the virtual losses are removed, and the actual evaluations are back propagated.

If the number of searches done before evaluations are ready (call this K) is too large, then the overall statistics of MCTS will be thrown off. Anecdotally, using K=16 for mcts_n=400 seems to cause no noticeable decrease in performance.

If we want to we could try techniques that are newer/better than virtual loss, like "virtual mean"

Self-play config variables renamed/added

• num_workers is now num_processes to be more explicit. I've been always using 1 rust selfplay process, to take advantage of the shared evaluator cache.
• parallel_games is now games_per_process to make it clearer. In python, each self-play process has exactly one thread that handles all of those games. In rust there can be more than one thread.
• mcts_worker_threads is only used in rust at the moment. This should be set as high as the CPU can handle. Each of these threads can work on any of the games.
• leaf_parallelism is the number of searches that can be done on a single MCTS tree before waiting for evaluations from the GPU. (virtual loss values are used until the evaluations are ready)
• virtual_loss is the number of losses temporarily backpropagated for each node waiting on an evaluation.
• enable_tree_reuse: if true, the mcts subtree of the chosen move is reused for the next move.
• eval_cache_max_size is used to limit the evaluator cache size. one million is a pretty safe value; set it as high as you can depending on the amount of memory the machine has.
• eval_batch_size: the evaluator batches up evaluation requests so that it can send them all to the GPU together. this is the maximum number of evaluations it can send to the GPU at once. setting this larger increases evaluation bandwidth, but also increases latency. if the latency is too high, mcts search threads running on the CPU will be blocked while waiting for the evaluation.
• eval_max_wait_ms: we don't want to wait forever for a full size batch. After this many milliseconds we just run whatever batch is ready to go.

[jonbinney]

Hmmm... not sure why that test is hanging in CI. Debugging now.