Open xiaomengy opened 2 years ago
Great suggestion So far we have implemented a vectorized version of lambda return that works several orders of magnitude faster than regular for-loops without requiring c++ code. I'd love to see how much we can gain by designing a c++ kernel for that if that's what you have in mind.
Here are the things I would keep in mind in this endeavour:
In RL algorithms, it is every common to compute return-like things from trajectories. It is inefficient to compute such returns with normal python for-loops. In order to improve efficiency, we'd better to abstract some general return ops and implement efficient PyTorch kernels. Some examples of such returns are listed below.