[Question] [Improvement] CUDA vectorized environment

[AlessandroZavoli]

I recently come across a paper about Isaac Gym by NVIDIA https://arxiv.org/pdf/2108.10470.pdf
They suggest that massive parallelization of the environment computation on the GPU can improve RL training, with respect to mixed CPU-GPU architectures (due to bottleneck in data transfer from host to device and vice versa). NVIDIA declares that researchers can achieve the same level of success as OpenAI’s supercomputer on a single GPU in about 10 hours.

The idea seems sound. A tentative implementation within Stable Baselines using Numba Kernels has been reported in a short undergraduate project report
https://www.sihao.dev/2021/05/21/increasing_sample_throughput_for_rl_environments_using_cuda/
Yet, I suspect that this implementation may have some pitfalls, as the training results are strange/questionable

Is someone interested in exploring this idea and/or suggesting how to implement a "CudaVecEnv" correctly?

[Miffyli]

Related issue #314

With high batchsize, you can get some of the boost of Isaac Gym by just doing the CudaVecEnv (VecEnv that, underneath, just calls the Isaac Gym's already-parallel envs).

To get the full boost, the problem is modifying all the buffers to be CUDA (or PyTorch) tensors and interconnecting it with Isaac Gym, which would not be a very backwards compatible change or at least would include bunch of code for switching between buffers, and at the very least they need to be done in pytorch instead of numpy arrays.

[AlessandroZavoli]

Thank you for your answer.
I see the point related to those buffers and the fact that a full boost cannot be obtained.

Yet, I was just wondering how to create that CudaVecEnv(VecEnv) class, that can be applied to a generic environment where the system dynamics (i.e., the step function) is just some algebraic function (e.g., cart-pole or rocket landing)
Maybe there is some example of reimplementing VecEnv that I can follow to better understand who to do it?
Or do you have some further advice?

[Miffyli]

Hmm unfortunately there are no docs on creating VecEnvs, but you can copy the DummyVecEnv code and start modifying from there. The major difference to normal Gym Envs is that step and reset etc return multiple observations (and take in multiple actions) at the same time :)

[araffin]

Yet, I was just wondering how to create that CudaVecEnv(VecEnv) class, that can be applied to a generic environment where the system dynamics (i.e., the step function) is just some algebraic function (e.g., cart-pole or rocket landing)

for isaac gym, you can take a look at this wrapper (adapter is here) and their vec env implementation

You should also take a look a env pool (only on cpu, but quite fast though), we have a wrapper to use it with SB3: https://github.com/sail-sg/envpool/blob/master/examples/sb3_examples/ppo.py

[araffin]

closing as some VenvEnvWrapper exist to handle GPU vectorized envs.

[LyuJZ]

Hi, thanks a lot for the well-documented stable baselines3. Now I am using Isaac Gym Preview4. May I ask if it is possible to give some examples to wrap IsaacGymEnvs into VecEnv? I noticed this issue was mentioned before. And some tips have been given in the issue #772. However, it seems it is for Isaac Gym Preview3. Could you give one example for Isaac Gym Preview 4?