A standard approach in reinforcement learning is to train multiple agents in parallel. Neural networks can produce multiple actions with a single forward pass, but stepping the environment in parallel requires multiprocessing. I am wondering if it is possible to design a patch for mujoco that allows multiple simulations to be run in parallel without multiprocessing. In principle, it should be possible to partition each vector in mjData into n chunks and apply simulation updates to each each chunk. I would like to do some work on this as a final project for a parallel architectures class that I am taking. I was hoping you might suggest a feasible approach to this and grant my team access to the subset of the source code that we would need to work on. Thanks!
A few years ago I implemented parts of MuJoCo in OpenCL, using vector types to do what you are suggesting: SIMD instructions operating on multiple mjData in one thread. And I got 4x speedup relative to C code, both on the CPU! GPU performance was disappointing because GPUs have too limited cache-per-core. The problem however is that only the simple parts of the pipeline are amenable to this approach. Control-flow branching is complicated: one has to manually separate the numbers in the batch, do serial processing, and merge again. Ideally the OpenCL compiler would be able to do the merging/auto-vectorization automatically and I would simply port the scalar code, but that did not produce any speedup; in general, compilers are still bad at auto-vectorizing complicated code, even though they make ambitious claims to the contrary. Some day I may return to this project, although it is not at the top of the todo list at the moment. Note that larger simulations leverage SIMD instructions in other ways -- multiplying matrices basically. And once the SIMD machinery of the CPU is occupied, there is nothing to be gained from the proposed approach. As for you doing it on your end, that can only happen if someone acquires the company and open-sources MuJoCo.
I'm glad to hear this is a promising direction from your perspective. It sounds like it's impossible for us to contribute though