[feature-request] Use torch's jit to reduce walltime

[PartiallyUntyped]

Motivation

RL, unlike DL, is reliant on both CPU and GPU performance, simply because we interact with the environment. DL is not reliant on the CPU because the processing is done in the GPU thus, python is a suitable choice. In contrast, python is horrible for RL simply because function calls, property accesses, element accesses, partial application, currying and so on, take a long time.

ps. related to #56.

Suggestion

See whether frequently used functions can be traced/jitted to reduce time in python's interpreter. This feature request is for >= 1.2.

Alternatives

Numba

See what can be compiled with numba or equivalent and is worth the effort. Issue: adds additional dependencies...

Cython

See what can be compiled with cython and decide whether the effort is worth it. Depends on Cython's compat with torch though and adds additional complexity.

[Miffyli]

Agreed this could be added later on once things have solidified a little (i.e. >= 1.2) as you suggested. Some of the functions are almost completely in pytorch functions (e.g. train function of A2C/PPO), so this should be very doable. The gains from this remain unseen but it would still deserve some testing, at the very least.

I do not know about including numba because it adds the dependency and makes implementing algorithms difficult. Perhaps if/when we have more complicated computations for returns/advantages, this can help writing them in intuitive form but also efficiently (I have had success with this when I was too lazy to figure out how to optimize those well :v).

[araffin]

Hello, i definitely agree with this feature (it was in my internal roadmap). i've tried a bit in the past but that was unsuccessful because jit support was limited and i had issues with saving/loading. Also, on another project, i did not see any gain :/ (but i assume that was my fault ^^`)

[PartiallyUntyped]

Another alternative is Cython, but from my experience, Cython becomes a mess to work with and sometimes workarounds are necessary to actually achieve improvements. Also Cython will require extensive code changes and IDK about its compatibility with Pytorch.

[araffin]

i've tried a bit in the past but that was unsuccessful because jit support was limited and i had issues with saving/loading

I gave it another try recently and support is still too limited compared to what you can do with jax: https://github.com/araffin/sbx

will wait a bit more then (some cools stuff are coming with torch dynamo =))