from torchrl.objectives.value.functional import reward2go
import torch
reward = torch.zeros(4, 2)
reward[3, 0] = 1
reward[3, 1] = -1
done = torch.zeros(4, 2, dtype=bool)
done[3, :] = True
reward2go(reward, done, 0.9)

tensor([[ 0.7290,  0.8100],
        [ 0.9000,  1.0000],
        [-0.7290, -0.8100],
        [-0.9000, -1.0000]])

tensor([[ 0.7290, -0.7290],
        [ 0.8100, -0.8100],
        [ 0.9000, -0.9000],
        [ 1.0000, -1.0000]])

import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)

    if cumsum.shape != shape:
        cumsum = cumsum.view(shape)

    cumsum = cumsum.transpose(-2, -1)