alex-petrenko/sample-factory

[Q] Cuda-centric software?

jarlva opened this issue · 2 comments

Hi, had a couple of questions:

. Have you explored cuda-centric software to accelerate work - like cupy instead of numpy? Nvidia claims 50-fold speedup.
. Is SF2.0 ready to "kick the tires"? How much faster (on avg) is it?

Thanks.

Hi @jarlva !

I'd say sf2 is about 1 month from release. Needs a finishing touch, a bunch of work on the documentation, and so on.
It is not going to be much faster for CPU-based envs because sf1 already achieved 100% CPU utilization with the right parameters, and for many CPU envs this would be the bottleneck.

Where sf2 became massively faster is at training vectorized (GPU-based) envs like Isaac Gym or Brax. For example, we can train on Ant in IsaacGym at ~500000 FPS on a single 3090Ti. At this point sf2 should be the fastest solution for all types of environments - CPU or GPU.

Regarding cupy, we barely use numpy for any serious computations. Everything that's possible to do on GPU is already done on GPU and is CUDA-accelerated.
I was tentatively looking at JAX for further acceleration of some parts of the pipeline, but this is something for sf2.1 or sf3.

Thank again @alex-petrenko!