salesforce / warp-drive

Extremely Fast End-to-End Deep Multi-Agent Reinforcement Learning Framework on a GPU (JMLR 2022)
BSD 3-Clause "New" or "Revised" License
465 stars 78 forks source link

Numba #Enhancement #14

Closed nubonics closed 1 year ago

nubonics commented 2 years ago

Perhaps, I misunderstand [very possible] what numba is for, but maybe it can used to replace learning how to write cuda c code, and instead just write python code that numba can translate [which runs directly on the gpu] into cuda c code?

Emerald01 commented 2 years ago

We always consider Numba as another option to implement the step() function. Yet I think we need to pay attention to some points.

For Numba, if you are saying some built-in implementation for arrays like Numpy did, those functional programming based coding is easy and looks similar to Numpy. However, for a multi-agent RL, it is almost impossible that you can finish your logic by just calling these functions. You will need many custom functionalities, like for loop, complicated logic branching and hashing and so on, and if that is the case, the syntax of Numba is pretty close to CUDA C except that it hides the C pointers. Numba will not translate your Python code to CUDA C code either.

We tested some Numba execution, it is significantly slower than CUDA C with more call-back overhead, but we will try to see if we can release some Numba environment in the near future as well since it is indeed easier to write.

Thank you.

nubonics commented 2 years ago

Ah okay, I was hoping for an easier way to write cuda c without actually have to write it in anything besides python lols.

Yes, I believe that the .step function is the most expensive, therefore it is definetly worth writing the cuda c step function in order to use warp-drive, because as the name suggests, its at the speed of light!!!

Emerald01 commented 1 year ago

@nubonics the numba is fully supported after version 2.0, and you can refer to example_envs about how those Numba envs are implemented. Thank you.