reinfors: an RL search/sampling engine in rust with caller-owned networks and training in python
I'd like to share reinfors, an open-source RL library that runs the engine in rust while the network and training loop remain ordinary python. Here is a quick overview, with full documentation available in the repo for anyone that's interested.
Motivation
During my own research, I found that existing open-source libraries did not offer the balance between modularity/composability and performance that I wanted. Python-first stacks are highly flexible but make simulation/search the bottleneck. Native frameworks (e.g. all-C++) keep the hot loop fast, but typically pull training into native code with it. Fully fused pipelines are generally the fastest, but their specialisation sacrifices flexibility.
Approach
reinfors is designed so that everyday use requires only python. Composing, training and evaluating never touch lower-level code. You compose an engine from the built-in games and algorithms, pass it an inference callback, and keep the network, optimizer and training loop as ordinary python code (pytorch, JAX, or anything else). Throughput-sensitive parts of the pipeline (simulation, search, episode orchestration, batch assembly) run in rust underneath. The engine pools inference requests across games and search leaves into numpy batches, so the python boundary is crossed rarely.
Does the python boundary cost performance?
This was my primary concern, so I ran some benchmarking experiments. I compared reinfors against an all-C++ libtorch implementation (OpenSpiel) on chess alphazero training, with the network, search budget and gradient intensity held equal. reinfors came out slightly ahead on throughput, and its trained agents performed slightly better in head-to-head games. Note that this was run on a single AWS instance (g5.2xlarge — one A10G GPU, four physical cores), so it is not a general claim, but it is evidence that the boundary need not cost meaningful performance. See link in repo README for details if interested.
What's in it today:
- Games — chess, backgammon, connect4, snake, gridworld, three poker variants
- Algorithms — alphazero, MCTS/treestrap, expectimax, minimax, DQN, PPO, CFR, Deep CFR, MCCFR
Current limitations:
- Fixed discrete action spaces only
- Sequential and simultaneous games only (no mixed-phase games that switch between sequential and simultaneous)
- The game/algorithm catalogue is still small
- New games and algorithms are written against rust traits (python is only for composing and training, not for defining new components)
I hope this may be useful for researchers wanting to experiment with the existing game/algorithm combinations (the examples should get you running quickly). For those comfortable with rust, the codebase is also designed to be extended. New games, policies and learners are written against small composable traits, and the extension guides should provide sufficient details.
Feedback and contributions are very welcome!