TransZero: Parallel Tree Expansion in MuZero using Transformer Networks
This work addresses the problem of slow planning times in model-based reinforcement learning for researchers and practitioners, enabling faster real-time decision-making in complex environments, though it is incremental as it builds on MuZero.
The paper tackles the sequential bottleneck in Monte Carlo Tree Search for model-based reinforcement learning by introducing TransZero, which uses a transformer network to generate multiple latent future states simultaneously, achieving up to an eleven-fold speedup in wall-clock time compared to MuZero while maintaining sample efficiency.
We present TransZero, a model-based reinforcement learning algorithm that removes the sequential bottleneck in Monte Carlo Tree Search (MCTS). Unlike MuZero, which constructs its search tree step by step using a recurrent dynamics model, TransZero employs a transformer-based network to generate multiple latent future states simultaneously. Combined with the Mean-Variance Constrained (MVC) evaluator that eliminates dependence on inherently sequential visitation counts, our approach enables the parallel expansion of entire subtrees during planning. Experiments in MiniGrid and LunarLander show that TransZero achieves up to an eleven-fold speedup in wall-clock time compared to MuZero while maintaining sample efficiency. These results demonstrate that parallel tree construction can substantially accelerate model-based reinforcement learning, bringing real-time decision-making in complex environments closer to practice. The code is publicly available on GitHub.