Offline RewardSum accumulates along the wrong dimension #4538

Offline RewardSum accumulates along the axis before the TensorDict dimension named time. This can sum independent environments or reward components instead of successive timesteps.

import torch
from tensordict import TensorDict
from torchrl.envs import RewardSum

td = TensorDict(
    {"reward": torch.tensor([[1.0], [2.0], [3.0]])},
    batch_size=[3],
    names=["time"],
)
print(RewardSum()(td)["episode_reward"].flatten())
# Actual:   tensor([1., 2., 3.])
# Expected: tensor([1., 3., 6.])

Reproduced on main at 648f50cf07dbb948ffd84b1eefd988dc664ecad1, using Python 3.13.9, PyTorch 2.14.0 and TensorDict 0.14.3 on CPU, with TorchRL imported from that source checkout. The optional C++ extension was not built and is unused by this transform.

RewardSum.forward finds the named time dimension and then subtracts one before calling cumsum. TensorDict batch axes already occupy the leading dimensions of its leaf tensors, so the subtraction selects the wrong axis. In a [batch, time, 1] reward tensor it accumulates across batch entries. The existing one-dimensional offline example hides this because axes 0 and -1 coincide there.

A one-line correction and six failing numeric regressions are prepared. They cover unbatched, batch-first and time-first layouts, nested keys and independent reward components.