Follow the suggestions from MinghongAlexXu Faster experience replay Pre-allocated and memory-mapped experience replay segment tree C++ implementation Write a C++ extension module for Python https://mesonbuild.com/Comparisons.html https://github.com/jaromiru/AI-blog/blob/master/SumTree.py https://pytorch.org/rl/reference/generated/torchrl.data.PrioritizedReplayBuffer.html https://discuss.pytorch.org/t/how-to-make-the-replay-buffer-more-efficient/80986 JAX From PyTorch to JAX: towards neural net frameworks that purify stateful code https://github.com/ikostrikov/jaxrl2 https://github.com/ikostrikov/jaxrl https://github.com/ikostrikov/walk_in_the_park/tree/main/rl Using JAX to accelerate our research IsaacGym JAX integration