# Reinforcement Learning for Load-balanced Parallel Particle Tracing

arXiv: https://arxiv.org/abs/2109.05679 · HTML: https://mlsanity.com/paper/2109.05679/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 9 of its 88 references that are in the map.
- First posted on arXiv 13 Sep 2021.

## Builds on (most cited first, 9 of 9)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [On the Global Convergence Rates of Softmax Policy Gradient Methods](https://mlsanity.com/paper/2005.06392/) (2020)
- [On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift](https://mlsanity.com/paper/1908.00261/) (2019)
- [Parameter estimation in softmax decision-making models with linear objective functions](https://mlsanity.com/paper/1502.04635/) (2015)
- [On Optimal Trees for Irregular Gather and Scatter Collectives](https://mlsanity.com/paper/1711.08731/) (2017)

## Abstract (arXiv)

We explore an online reinforcement learning (RL) paradigm to dynamically optimize parallel particle tracing performance in distributed-memory systems. Our method combines three novel components: (1) a work donation algorithm, (2) a high-order workload estimation model, and (3) a communication cost model. First, we design an RL-based work donation algorithm. Our algorithm monitors workloads of processes and creates RL agents to donate data blocks and particles from high-workload processes to low-workload processes to minimize program execution time. The agents learn the donation strategy on the fly based on reward and cost functions designed to consider processes' workload changes and data transfer costs of donation actions. Second, we propose a workload estimation model, helping RL agents estimate the workload distribution of processes in future computations. Third, we design a communication cost model that considers both block and particle data exchange costs, helping RL agents make effective decisions with minimized communication costs. We demonstrate that our algorithm adapts to different flow behaviors in large-scale fluid dynamics, ocean, and weather simulation data. Our algorithm improves parallel particle tracing performance in terms of parallel efficiency, load balance, and costs of I/O and communication for evaluations with up to 16,384 processors.