Repository: pytorch/rl

Issue:
[Feature Request] Support custom, tensor-based decay in TD(lambda)
## Motivation

Currently, the vectorised TD(lambda) function only supports scalar gamma values.
However, in some cases we'll want to provide a gamma vector: for instance, the [`MultiStep`](https://github.com/facebookresearch/rl/blob/d295fae9f532862d60e28ca1f509ad1183cf6a56/torchrl/data/postprocs/postprocs.py#L109) class computes a gamma tensor of the same size as the reward that is to be used to compute decayed reward sums.

## Solution

The proposed solution is to allow [`_custom_conv1d`](https://github.com/facebookresearch/rl/blob/266ddff05ce0c31ba526c0124e0f025ff5273286/torchrl/objectives/returns/functional.py#L167) to receive filters (i.e. stack of gamma values) with a certain batch size. The conv1d operation does not support batching but we can do something like this:

```python
if filter.ndimension() > 2:  # then filter has a batch size
    filter = filter.reshape(batch_size, 1, 1, T)
    functorch.vmap(torch.conv1d, (None, 0))(val_pad, filter)  # (None, 0) tells vmap to expand val_pad and not filter
else:
   do_as_before()
```

In this implementation, `functorch` will do the vectorised map of filter over val_pad (i.e. expand val_pad and execute conv1d over it).

Naturally, the creation of the filter and the reshaping will need to be adapted accordingly.

Tests should be added to account for this feature and placed in `test_costs.py`.
For instance:

```python
batch = 10
time = 11
next_value = torch.randn(batch, time, 1)
done = torch.zeros(batch, time, 1, dtype=torch.bool)
reward = torch.randn(batch, time, 1)

gamma = 0.99
lmbda = 0.95

gamma_tensor = torch.full(gamma, (batch, time, 1))

v1 = vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done)
v2 = vec_td_lambda_return_estimate(gamma_tensor, lmbda, next_state_value, reward, done)

assert_isclose(v1, v2)

# same with last done being true
done[:, -1] = True  # terminating trajectory
gamma_tensor[:, -1] = 0.0

v1 = vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done)
v2 = vec_td_lambda_return_estimate(gamma_tensor, lmbda, next_state_value, reward, done)

assert_isclose(v1, v2)

```


Discussion:

Resolving patch:
diff --git a/torchrl/objectives/returns/functional.py b/torchrl/objectives/returns/functional.py
--- a/torchrl/objectives/returns/functional.py
+++ b/torchrl/objectives/returns/functional.py
@@ -195,9 +195,30 @@ def _custom_conv1d(tensor, filter):
         -1,
     )
 
-    # shape = val.shape
-    filter = filter.squeeze(-1).unsqueeze(0).unsqueeze(0)  # 1 x 1 x T
-    out = torch.conv1d(val_pad, filter)
+    if filter.ndimension() > 2:
+        # filter will have shape batch_dims x timesteps x filter_dim x 1
+        # reshape to batch_dims x timesteps x 1 x filter_dim ready for convolving
+        filter = filter.view(*filter.shape[:-2], 1, filter.shape[-2])
+
+        # because time is represented on two different dimensions, we don't
+        # need all convolutions, just those lying along a diagonal
+        # rather than compute them all and discard, we stack just the slices
+        # of val_pad that we care about, and apply the filter manually
+        batched_val_pad = torch.stack(
+            [val_pad[..., i : i + filter.shape[-1]] for i in range(tensor.shape[-1])],
+            dim=1,
+        )
+
+        # this is just a batched matrix multiplication, but einsum makes it
+        # easy to keep the many dimensions under control. Here b = batch,
+        # t = timestep, s = singleton, j is the filter dimension that should
+        # get summed out. we swap the order of s and t here rather than
+        # reshape / create a view later
+        out = torch.einsum("btsj,btsj->bst", batched_val_pad, filter)
+    else:
+        # shape = val.shape
+        filter = filter.squeeze(-1).unsqueeze(0).unsqueeze(0)  # 1 x 1 x T
+        out = torch.conv1d(val_pad, filter)
     # out = out.view(shape)
     if not out.shape == tensor.shape:
         raise RuntimeError("wrong output shape")
@@ -210,7 +231,8 @@ def vec_td_lambda_advantage_estimate(
     """Vectorized TD(lambda) advantage estimate.
 
     Args:
-        gamma (scalar): exponential mean discount.
+        gamma (scalar, Tensor): exponential mean discount. If tensor-valued,
+            must be a [Batch x TimeSteps x 1] tensor.
         lmbda (scalar): trajectory discount.
         state_value (Tensor): value function result with old_state input.
             must be a [Batch x TimeSteps x 1] or [Batch x TimeSteps] tensor
@@ -230,10 +252,11 @@ def vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done):
     """Vectorized TD(lambda) return estimate.
 
     Args:
-        gamma (scalar): exponential mean discount.
+        gamma (scalar, Tensor): exponential mean discount. If tensor-valued,
+            must be a [Batch x TimeSteps x 1] tensor.
         lmbda (scalar): trajectory discount.
         next_state_value (Tensor): value function result with new_state input.
-            must be a [Batch x TimeSteps x 1] or [Batch x TimeSteps] tensor
+            must be a [Batch x TimeSteps x 1] tensor
         reward (Tensor): reward of taking actions in the environment.
             must be a [Batch x TimeSteps x 1] or [Batch x TimeSteps] tensor
         done (Tensor): boolean flag for end of episode.
@@ -243,21 +266,27 @@ def vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done):
     if not shape[-1] == 1:
         raise RuntimeError("last dimension of inputs shape must be singleton")
 
-    next_state_value = next_state_value.view(-1, 1, shape[-2])
-    reward = reward.view(-1, 1, shape[-2])
-    done = done.view(-1, 1, shape[-2])
+    T = shape[-2]
+
+    next_state_value = next_state_value.view(-1, 1, T)
+    reward = reward.view(-1, 1, T)
+    done = done.view(-1, 1, T)
 
     """Vectorized version of td_lambda_advantage_estimate"""
     device = reward.device
     not_done = 1 - done.to(next_state_value.dtype)
     next_state_value = not_done * next_state_value
 
-    T = shape[-2]
-
     first_below_thr_gamma = None
 
-    gammas = torch.ones(T + 1, 1, device=device)
-    gammas[1:] = gamma
+    if isinstance(gamma, torch.Tensor) and gamma.ndimension() > 0:
+        gamma = gamma.view(-1, T)
+        gammas = torch.ones(*gamma.shape, T + 1, 1, device=device)
+        gammas[..., 1:, :] = gamma[..., None, None]
+    else:
+        gammas = torch.ones(T + 1, 1, device=device)
+        gammas[1:] = gamma
+
     gammas = torch.cumprod(gammas, -2)
 
     lambdas = torch.ones(T + 1, 1, device=device)
@@ -265,15 +294,19 @@ def vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done):
     lambdas = torch.cumprod(lambdas, -2)
 
     first_below_thr = gammas < 1e-7
+    while first_below_thr.ndimension() > 2:
+        # if we have multiple gammas, we only want to truncate if _all_ of
+        # the geometric sequences fall below the threshold
+        first_below_thr = first_below_thr.all(axis=0)
     if first_below_thr.any():
         first_below_thr_gamma = first_below_thr.nonzero()[0, 0]
     first_below_thr = lambdas < 1e-7
     if first_below_thr.any() and first_below_thr_gamma is not None:
         first_below_thr = max(first_below_thr_gamma, first_below_thr.nonzero()[0, 0])
-        gammas = gammas[:first_below_thr]
+        gammas = gammas[..., :first_below_thr, :]
         lambdas = lambdas[:first_below_thr]
 
-    gammas, gammas_prime = gammas[:-1], gammas[1:]
+    gammas, gammas_prime = gammas[..., :-1, :], gammas[..., 1:, :]
     lambdas, lambdas_prime = lambdas[:-1], lambdas[1:]
 
     rs = _custom_conv1d(reward, gammas * lambdas)
@@ -285,14 +318,23 @@ def vec_td_lambda_return_estimate(gamma, lmbda, next_state_value, reward, done):
             # [torch.zeros_like(next_state_value[..., : -mask.shape[-2], :]), mask], -2
             [
                 torch.zeros(
-                    next_state_value.shape[-1] - mask.shape[-2], 1, device=device
+                    *mask.shape[:-2],
+                    next_state_value.shape[-1] - mask.shape[-2],
+                    1,
+                    device=device
                 ),
                 mask,
             ],
             -2,
         )
-    vs2 = (
-        _custom_conv1d(next_state_value, gam_lam)
-        - mask.squeeze(-1) * next_state_value[..., -1:]
-    )
+    if gammas.ndimension() > 2:
+        vs2 = (
+            _custom_conv1d(next_state_value, gam_lam)
+            - mask.squeeze(-1)[..., -1:, :] * next_state_value[..., -1:]
+        )
+    else:
+        vs2 = (
+            _custom_conv1d(next_state_value, gam_lam)
+            - mask.squeeze(-1) * next_state_value[..., -1:]
+        )
     return (rs + vs - vs2).view(shape)


Test patch:
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -55,6 +55,7 @@
 from torchrl.objectives.costs.utils import hold_out_net, HardUpdate, SoftUpdate
 from torchrl.objectives.returns.advantages import TDEstimate, GAE, TDLambdaEstimate
 from torchrl.objectives.returns.functional import (
+    _custom_conv1d,
     vec_td_lambda_advantage_estimate,
     td_lambda_advantage_estimate,
 )
@@ -1713,6 +1714,74 @@ def test_tdlambda(device, gamma, lmbda, N, T):
     torch.testing.assert_close(r1, r2, rtol=1e-4, atol=1e-4)
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("gamma", [0.1, 0.5, 0.99])
+@pytest.mark.parametrize("lmbda", [0.1, 0.5, 0.99])
+@pytest.mark.parametrize("N", [(3,), (7, 3)])
+@pytest.mark.parametrize("T", [3, 5, 200])
+def test_tdlambda_tensor_gamma(device, gamma, lmbda, N, T):
+    torch.manual_seed(0)
+
+    done = torch.zeros(*N, T, 1, device=device, dtype=torch.bool)
+    reward = torch.randn(*N, T, 1, device=device)
+    state_value = torch.randn(*N, T, 1, device=device)
+    next_state_value = torch.randn(*N, T, 1, device=device)
+
+    gamma_tensor = torch.full((*N, T, 1), gamma, device=device)
+
+    v1 = vec_td_lambda_advantage_estimate(
+        gamma, lmbda, state_value, next_state_value, reward, done
+    )
+    v2 = vec_td_lambda_advantage_estimate(
+        gamma_tensor, lmbda, state_value, next_state_value, reward, done
+    )
+
+    torch.testing.assert_close(v1, v2, rtol=1e-4, atol=1e-4)
+
+    # # same with last done being true
+    done[..., -1, :] = True  # terminating trajectory
+    gamma_tensor[..., -1, :] = 0.0
+
+    v1 = vec_td_lambda_advantage_estimate(
+        gamma, lmbda, state_value, next_state_value, reward, done
+    )
+    v2 = vec_td_lambda_advantage_estimate(
+        gamma_tensor, lmbda, state_value, next_state_value, reward, done
+    )
+
+    torch.testing.assert_close(v1, v2, rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("gamma", [0.1, 0.5, 0.99])
+@pytest.mark.parametrize("N", [(3,), (3, 7)])
+@pytest.mark.parametrize("T", [3, 5, 200])
+def test_custom_conv1d_tensor(device, gamma, N, T):
+    """
+    Tests the _custom_conv1d logic against a manual for-loop implementation
+    """
+    torch.manual_seed(0)
+
+    gamma = torch.rand(*N, T, 1, device=device)
+    values = torch.randn(*N, 1, T, device=device)
+
+    out = torch.zeros(*N, 1, T, device=device)
+    for i in range(T):
+        for j in reversed(range(i, T)):
+            out[..., i] = out[..., i] * gamma[..., i, :] + values[..., j]
+
+    # some reshaping code vendored from vec_td_lambda_return_estimate
+    gamma = gamma.view(-1, T)
+    gammas = torch.ones(*gamma.shape, T + 1, 1, device=device)
+    gammas[..., 1:, :] = gamma[..., None, None]
+    gammas = torch.cumprod(gammas, -2)
+    filter = gammas[..., :-1, :]
+
+    out_custom = _custom_conv1d(values.view(-1, 1, T), filter).reshape(values.shape)
+
+    torch.testing.assert_close(out, out_custom, rtol=1e-4, atol=1e-4)
+
+
 @pytest.mark.parametrize(
     "dest,expected_dtype,expected_device",
     list(