[Performance] Accelerate TD lambda return estimate

[Blonck]

Description

Implement a version of for vec_td_lambda_return_estimate optimized for the case where gamma is a scalar.

Motivation and Context

In the case gamma is a scalar one can do the same optimization as done in #1141. Instead of constructing a big gamma tensor [B, T, T], one can split and pad the consecutive traces:

reward = [r00, r01, r02, r03, r10, r11]
done = [False, False, False, True, False, False]

into

r_transformed = [
    [r00, r01, r02, r03],
    [r10, r11, 0, 0]
]

and then apply a gamma filter to this transformed input.

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• [x ] Bug fix (non-breaking change which fixes an issue)

Closes #1148

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• I have updated the tests accordingly (required for a bug fix or a new feature).

[Blonck]

Performance comparison for shape (16, T, 1) on GPU. Shown is the the ratio runtime(algorithm)/runtime(td_lambda_return_estimate). 1 means the algorithm has identical runtime as the unvectorized version of TD lambda.

[vmoens]

LGTM thanks for this!