Repository: pytorch/rl

Issue:
[BUG] Transpose bug in `reward2go` when the last dim is not 1
## Describe the bug

When calling `reward2go` with a shape whose last dimension is not 1, the results are incorrect due to reshaping rather than transposing.

## To Reproduce

Minimal example:

```python
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)
```

Output:

```python
tensor([[ 0.7290,  0.8100],
        [ 0.9000,  1.0000],
        [-0.7290, -0.8100],
        [-0.9000, -1.0000]])
```

## Expected behavior

```python
tensor([[ 0.7290, -0.7290],
        [ 0.8100, -0.8100],
        [ 0.9000, -0.9000],
        [ 1.0000, -1.0000]])
```

## System info

```python
import torchrl, numpy, sys
print(torchrl.__version__, numpy.__version__, sys.version, sys.platform)
```

```
0.4.0+8570bd3 1.26.2 3.11.6 (main, Oct  8 2023, 05:06:43) [GCC 13.2.0] linux
```

## Reason and Possible fixes

In the last step of `reward2go` the cumsum is reshaped:

```python
    if cumsum.shape != shape:
        cumsum = cumsum.view(shape)
```

It should be:

```python
    cumsum = cumsum.transpose(-2, -1)
```

to undo an earlier transpose.

I'm happy to contribute a PR with a fix and an extra test.

## Checklist

- [X] I have checked that there is no similar issue in the repo (**required**)
- [X] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [X] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:

Resolving patch:
diff --git a/torchrl/objectives/value/functional.py b/torchrl/objectives/value/functional.py
--- a/torchrl/objectives/value/functional.py
+++ b/torchrl/objectives/value/functional.py
@@ -1376,7 +1376,11 @@ def reward2go(
     cumsum = _custom_conv1d(td0_flat.unsqueeze(1), gammas)
     cumsum = cumsum.squeeze(1)
     cumsum = _inv_pad_sequence(cumsum, num_per_traj)
-    cumsum = cumsum.view_as(reward)
+    cumsum = cumsum.reshape_as(reward)
+    cumsum = cumsum.transpose(-2, -1)
     if cumsum.shape != shape:
-        cumsum = cumsum.view(shape)
+        raise RuntimeError(
+            f"Wrong shape for output reward2go: {cumsum.shape} when {shape} was expected."
+        )
+    #     cumsum = cumsum.view(shape)
     return cumsum


Test patch:
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -137,6 +137,7 @@
 from torchrl.objectives.value.functional import (
     _transpose_time,
     generalized_advantage_estimate,
+    reward2go,
     td0_advantage_estimate,
     td1_advantage_estimate,
     td_lambda_advantage_estimate,
@@ -13246,6 +13247,18 @@ def test_split_pad_reverse_tensordict(self, B, T, device):
         reversed = reversed.reshape(td.shape)
         torch.testing.assert_close(td["observation"], reversed["observation"])
 
+    def test_reward2go(self):
+        reward = torch.zeros(4, 2)
+        reward[3, 0] = 1
+        reward[3, 1] = -1
+        done = torch.zeros(4, 2, dtype=bool)
+        done[3, :] = True
+        r = torch.ones(4)
+        r[1:] = 0.9
+        r = torch.cumprod(r, 0).flip(0)
+        r = torch.stack([r, -r], -1)
+        torch.testing.assert_close(reward2go(reward, done, 0.9), r)
+
     def test_timedimtranspose_single(self):
         @_transpose_time
         def fun(a, b, time_dim=-2):