Repository: pytorch/rl

Issue:
[WIP] Accelerate GAE
## Description

In a recent effort to support consecutive trajectories, we refactored GAE such that the following trajectories are supported

```
reward = [r00, r01, r02, r03, r10, r11]
done = [False, False, False, True, False, False]
```
The way GAE and other vectorized advantages work is that we build a cumprod of the gamma decay
```
gamma_cumprod = [1, gamma, gamma ** 2, ...]
```
and apply a conv1d on the reward (or another tensor) along the time dimension.
This results in 
```
[r00 + gamma r01 + gamma ** 2 r02 + ....,
r01 + gamma r02 + gamma ** 2 r03 + ....,
]
```

With consecutive trajectories, we can't really do that because one trajectory will see its rewards carried over the previous trajectory. The "done" signal should stop such propagation.
The way `vec_generalized_advantage_estiamate` deals with this is by building a "giant" gamma tensor of size `[Batch, T, T]` with the decayed gamma that suits each trajectory and only this, and multiply this by the reward (or other signal) to apply the sum-decay over. This implies two heavy operations: allocating the big `[B, T, T]` tensor and performing the heavy matrix multiplication.
The result is that the vectorized GAE can be slower than its loop counterpart when B or T is high.
This PR proposes a solution: if there is a single gamma and lambda, we can build a tensor
```
r_transformed = [[r00, r01, r02, r03],
                              [r10, r11, 0, 0]]
```
and apply the gamma filter presented above, before remapping the resulting r_transformed to its original shape.

The file `functional_gae_wip.py` has this solution implemented, but the bits need to be linked to the current pipeline. 
Something like
```
def vec_generalized_advantage_estiamte(...):
    if isscalar(gamma) and isscalar(lambda):
        return fast_gae(...)
    # else do what we did before as we need the "big" gamma tensor if there is one gamma associated with each state
    if not gamma.shape == reward.shape:
        raise ValueError("if gamma is not a scalar, its shape must match the reward shape)

```
Also, we should test whether we actually have a better performance with this solution on CPU and GPU.

Finally, we need to add more tests for various B and T sized in the benchmarking tool to check in which condition one outperforms another.



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
@@ -3,10 +3,10 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 from functools import wraps
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import torch
-from tensordict import MemmapTensor, TensorDictBase
+from tensordict import MemmapTensor
 
 __all__ = [
     "generalized_advantage_estimate",
@@ -23,7 +23,13 @@
     "vec_td_lambda_advantage_estimate",
 ]
 
-from torchrl.objectives.value.utils import _custom_conv1d, _make_gammas_tensor
+from torchrl.objectives.value.utils import (
+    _custom_conv1d,
+    _get_num_per_traj,
+    _inv_pad_sequence,
+    _make_gammas_tensor,
+    _split_and_pad_sequence,
+)
 
 
 def _transpose_time(fun):
@@ -114,10 +120,65 @@ def generalized_advantage_estimate(
     return advantage, value_target
 
 
-@_transpose_time
-def vec_generalized_advantage_estimate(
+def _fast_vec_gae(
+    reward: torch.Tensor,
+    state_value: torch.Tensor,
+    next_state_value: torch.Tensor,
+    done: torch.Tensor,
     gamma: float,
     lmbda: float,
+):
+    """Fast vectorized Generalized Advantage Estimate when gamma and lmbda are scalars.
+
+    In contrast to `vec_generalized_advantage_estimate` this function does not need
+    to allocate a big tensor of the form [B, T, T].
+
+    Args:
+        reward (torch.Tensor): a [*B, T, F] tensor containing rewards
+        state_value (torch.Tensor): a [*B, T, F] tensor containing state values (value function)
+        next_state_value (torch.Tensor): a [*B, T, F] tensor containing next state values (value function)
+        done (torch.Tensor): a [B, T] boolean tensor containing the done states
+        gamma (scalar): the gamma decay (trajectory discount)
+        lmbda (scalar): the lambda decay (exponential mean discount)
+
+    All tensors (values, reward and done) must have shape
+    ``[*Batch x TimeSteps x F]``, with ``F`` feature dimensions.
+    """
+    # _gen_num_per_traj and _split_and_pad_sequence need
+    # time dimension at last position
+    done = done.transpose(-2, -1)
+    reward = reward.transpose(-2, -1)
+    state_value = state_value.transpose(-2, -1)
+    next_state_value = next_state_value.transpose(-2, -1)
+
+    gammalmbda = gamma * lmbda
+    not_done = 1 - done.int()
+    td0 = reward + not_done * gamma * next_state_value - state_value
+
+    num_per_traj = _get_num_per_traj(done)
+    td0_flat = _split_and_pad_sequence(td0, num_per_traj)
+
+    gammalmbdas = torch.ones_like(td0_flat[0])
+    gammalmbdas[1:] = gammalmbda
+    gammalmbdas[1:] = gammalmbdas[1:].cumprod(0)
+    gammalmbdas = gammalmbdas.unsqueeze(-1)
+
+    advantage = _custom_conv1d(td0_flat.unsqueeze(1), gammalmbdas)
+    advantage = advantage.squeeze(1)
+    advantage = _inv_pad_sequence(advantage, num_per_traj).view_as(reward)
+
+    value_target = advantage + state_value
+
+    advantage = advantage.transpose(-1, -2)
+    value_target = value_target.transpose(-1, -2)
+
+    return advantage, value_target
+
+
+@_transpose_time
+def vec_generalized_advantage_estimate(
+    gamma: Union[float, torch.Tensor],
+    lmbda: Union[float, torch.Tensor],
     state_value: torch.Tensor,
     next_state_value: torch.Tensor,
     reward: torch.Tensor,
@@ -151,19 +212,31 @@ def vec_generalized_advantage_estimate(
     *batch_size, time_steps, lastdim = not_done.shape
 
     value = gamma * lmbda
-    if isinstance(value, torch.Tensor):
+
+    if isinstance(value, torch.Tensor) and value.numel() > 1:
         # create tensor while ensuring that gradients are passed
-        gammalmbdas = torch.ones_like(not_done) * not_done * value
+        gammalmbdas = not_done * value
     else:
-        gammalmbdas = torch.full_like(not_done, value) * not_done
+        # when gamma and lmbda are scalars, use fast_vec_gae implementation
+        return _fast_vec_gae(
+            reward=reward,
+            state_value=state_value,
+            next_state_value=next_state_value,
+            done=done,
+            gamma=gamma,
+            lmbda=lmbda,
+        )
+
     gammalmbdas = _make_gammas_tensor(gammalmbdas, time_steps, True)
     gammalmbdas = gammalmbdas.cumprod(-2)
-    # first_below_thr = gammalmbdas < 1e-7
-    # # 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():
-    #     gammalmbdas = gammalmbdas[..., :first_below_thr, :]
+
+    first_below_thr = gammalmbdas < 1e-7
+    # 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.flatten(0, 1).all(0).all(-1)
+    if first_below_thr.any():
+        first_below_thr = torch.where(first_below_thr)[0][0].item()
+        gammalmbdas = gammalmbdas[..., :first_below_thr, :]
 
     td0 = reward + not_done * gamma * next_state_value - state_value
 
@@ -900,32 +973,6 @@ def vec_td_lambda_advantage_estimate(
 # ------------
 
 
-def _flatten_batch(tensor):
-    """Because we mark the end of each batch with a truncated signal, we can concatenate them.
-
-    Args:
-        tensor (torch.Tensor): a tensor of shape [B, T]
-    """
-    return tensor.flatten(0, 1)
-
-
-def _get_num_per_traj(dones_and_truncated):
-    """Because we mark the end of each batch with a truncated signal, we can concatenate them.
-
-    Args:
-        dones_and_truncated (torch.Tensor): A done or truncated mark of shape [B, T]
-
-    Returns:
-        A list of integers representing the number of steps in each trajectories
-    """
-    dones_and_truncated = dones_and_truncated.clone()
-    dones_and_truncated[..., -1] = 1
-    dones_and_truncated = _flatten_batch(dones_and_truncated)
-    num_per_traj = torch.ones_like(dones_and_truncated).cumsum(0)[dones_and_truncated]
-    num_per_traj[1:] -= num_per_traj[:-1].clone()
-    return num_per_traj
-
-
 def _get_num_per_traj_init(is_init):
     """Like _get_num_per_traj, but with is_init signal."""
     done = torch.zeros_like(is_init)
@@ -933,145 +980,6 @@ def _get_num_per_traj_init(is_init):
     return _get_num_per_traj(done)
 
 
-def _split_and_pad_sequence(tensor, splits):
-    """Given a tensor of size [B, T, *other] and the corresponding traj lengths (flattened), returns the padded trajectories [NPad, Tmax, *other].
-
-    Compatible with tensordict inputs.
-
-    Examples:
-        >>> from tensordict import TensorDict
-        >>> is_init = torch.zeros(4, 5, dtype=torch.bool)
-        >>> is_init[:, 0] = True
-        >>> is_init[0, 3] = True
-        >>> is_init[1, 2] = True
-        >>> tensordict = TensorDict({
-        ...     "is_init": is_init,
-        ...     "obs": torch.arange(20).view(4, 5).unsqueeze(-1).expand(4, 5, 3),
-        ... }, [4, 5])
-        >>> splits = _get_num_per_traj_init(is_init)
-        >>> print(splits)
-        tensor([3, 2, 2, 3, 5, 5])
-        >>> td = _split_and_pad_sequence(tensordict, splits)
-        >>> print(td)
-        TensorDict(
-            fields={
-                is_init: Tensor(shape=torch.Size([6, 5]), device=cpu, dtype=torch.bool, is_shared=False),
-                obs: Tensor(shape=torch.Size([6, 5, 3]), device=cpu, dtype=torch.int64, is_shared=False)},
-            batch_size=torch.Size([6, 5]),
-            device=None,
-            is_shared=False)
-        >>> print(td["obs"])
-        tensor([[[ 0,  0,  0],
-                 [ 1,  1,  1],
-                 [ 2,  2,  2],
-                 [ 0,  0,  0],
-                 [ 0,  0,  0]],
-        <BLANKLINE>
-                [[ 3,  3,  3],
-                 [ 4,  4,  4],
-                 [ 0,  0,  0],
-                 [ 0,  0,  0],
-                 [ 0,  0,  0]],
-        <BLANKLINE>
-                [[ 5,  5,  5],
-                 [ 6,  6,  6],
-                 [ 0,  0,  0],
-                 [ 0,  0,  0],
-                 [ 0,  0,  0]],
-        <BLANKLINE>
-                [[ 7,  7,  7],
-                 [ 8,  8,  8],
-                 [ 9,  9,  9],
-                 [ 0,  0,  0],
-                 [ 0,  0,  0]],
-        <BLANKLINE>
-                [[10, 10, 10],
-                 [11, 11, 11],
-                 [12, 12, 12],
-                 [13, 13, 13],
-                 [14, 14, 14]],
-        <BLANKLINE>
-                [[15, 15, 15],
-                 [16, 16, 16],
-                 [17, 17, 17],
-                 [18, 18, 18],
-                 [19, 19, 19]]])
-
-    """
-    tensor = _flatten_batch(tensor)
-    max_val = max(splits)
-    mask = torch.zeros(len(splits), max_val, dtype=torch.bool, device=tensor.device)
-    mask.scatter_(
-        index=max_val - torch.tensor(splits, device=tensor.device).unsqueeze(-1),
-        dim=1,
-        value=1,
-    )
-    mask = mask.cumsum(-1).flip(-1).bool()
-
-    def _fill_tensor(tensor):
-        empty_tensor = torch.zeros(
-            len(splits),
-            max_val,
-            *tensor.shape[1:],
-            dtype=tensor.dtype,
-            device=tensor.device,
-        )
-        empty_tensor[mask] = tensor
-        return empty_tensor
-
-    if isinstance(tensor, TensorDictBase):
-        tensor = tensor.apply(_fill_tensor, batch_size=[len(splits), max_val])
-    else:
-        tensor = _fill_tensor(tensor)
-    return tensor
-
-
-def _inv_pad_sequence(tensor, splits):
-    """Inverses a pad_sequence operation.
-
-    Examples:
-        >>> rewards = torch.randn(100, 20)
-        >>> num_per_traj = _get_num_per_traj(torch.zeros(100, 20).bernoulli_(0.1))
-        >>> padded = _split_and_pad_sequence(rewards, num_per_traj.tolist())
-        >>> reconstructed = _inv_pad_sequence(padded, num_per_traj)
-        >>> assert (reconstructed==rewards).all()
-
-
-    Compatible with tensordict inputs.
-
-    Examples:
-        >>> from tensordict import TensorDict
-        >>> is_init = torch.zeros(4, 5, dtype=torch.bool)
-        >>> is_init[:, 0] = True
-        >>> is_init[0, 3] = True
-        >>> is_init[1, 2] = True
-        >>> tensordict = TensorDict({
-        ...     "is_init": is_init,
-        ...     "obs": torch.arange(20).view(4, 5).unsqueeze(-1).expand(4, 5, 3),
-        ... }, [4, 5])
-        >>> splits = _get_num_per_traj_init(is_init)
-        >>> td = _split_and_pad_sequence(tensordict, splits)
-        >>> assert (_inv_pad_sequence(td, splits).view(tensordict.shape) == tensordict).all()
-
-    """
-    offset = torch.ones_like(splits) * tensor.shape[-1]
-    offset[0] = 0
-    offset = offset.cumsum(0)
-    z = torch.zeros(tensor.numel(), dtype=torch.bool, device=offset.device)
-
-    ones = offset + splits
-    ones = ones[ones < tensor.numel()]
-    # while ones[-1] == tensor.numel():
-    #     ones = ones[:-1]
-    z[ones] = 1
-    z_idx = z[offset[1:]]
-    z[offset[1:]] = torch.bitwise_xor(
-        z_idx, torch.ones_like(z_idx)
-    )  # make sure that the longest is accounted for
-    idx = z.cumsum(0) % 2 == 0
-    return tensor.reshape(-1)[idx]
-
-
 @_transpose_time
 def reward2go(
     reward,
@@ -1130,6 +1038,7 @@ def reward2go(
     gammas[1:] = gammas[1:].cumprod(0)
     gammas = gammas.unsqueeze(-1)
     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)
     if cumsum.shape != shape:
diff --git a/torchrl/objectives/value/utils.py b/torchrl/objectives/value/utils.py
--- a/torchrl/objectives/value/utils.py
+++ b/torchrl/objectives/value/utils.py
@@ -5,6 +5,8 @@
 
 import torch
 
+from tensordict import TensorDictBase
+
 
 def _custom_conv1d(tensor: torch.Tensor, filter: torch.Tensor):
     """Computes a conv1d filter over a value.
@@ -183,3 +185,146 @@ def _make_gammas_tensor(gamma: torch.Tensor, T: int, rolling_gamma: bool):
         gammas = torch.ones(*gamma.shape, T + 1, 1, device=device, dtype=dtype)
         gammas[..., 1:, :] = gamma[..., None, None]
     return gammas
+
+
+def _flatten_batch(tensor):
+    """Because we mark the end of each batch with a truncated signal, we can concatenate them.
+
+    Args:
+        tensor (torch.Tensor): a tensor of shape [*B, T]
+
+    """
+    return tensor.flatten(0, -1)
+
+
+def _get_num_per_traj(dones_and_truncated):
+    """Because we mark the end of each batch with a truncated signal, we can concatenate them.
+
+    Args:
+        dones_and_truncated (torch.Tensor): A done or truncated mark of shape [*B, T]
+
+    Returns:
+        A list of integers representing the number of steps in each trajectories
+
+    """
+    dones_and_truncated = dones_and_truncated.clone()
+    dones_and_truncated[..., -1] = 1
+    dones_and_truncated = _flatten_batch(dones_and_truncated)
+    num_per_traj = torch.ones_like(dones_and_truncated).cumsum(0)[dones_and_truncated]
+    num_per_traj[1:] -= num_per_traj[:-1].clone()
+    return num_per_traj
+
+
+def _split_and_pad_sequence(tensor, splits):
+    """Given a tensor of size [B, T, *other] and the corresponding traj lengths (flattened), returns the padded trajectories [NPad, Tmax, *other].
+
+    Compatible with tensordict inputs.
+
+    Examples:
+        >>> from tensordict import TensorDict
+        >>> is_init = torch.zeros(4, 5, dtype=torch.bool)
+        >>> is_init[:, 0] = True
+        >>> is_init[0, 3] = True
+        >>> is_init[1, 2] = True
+        >>> tensordict = TensorDict({
+        ...     "is_init": is_init,
+        ...     "obs": torch.arange(20).view(4, 5).unsqueeze(-1).expand(4, 5, 3),
+        ... }, [4, 5])
+        >>> splits = _get_num_per_traj_init(is_init)
+        >>> print(splits)
+        tensor([3, 2, 2, 3, 5, 5])
+        >>> td = _split_and_pad_sequence(tensordict, splits)
+        >>> print(td)
+        TensorDict(
+            fields={
+                is_init: Tensor(shape=torch.Size([6, 5]), device=cpu, dtype=torch.bool, is_shared=False),
+                obs: Tensor(shape=torch.Size([6, 5, 3]), device=cpu, dtype=torch.int64, is_shared=False)},
+            batch_size=torch.Size([6, 5]),
+            device=None,
+            is_shared=False)
+        >>> print(td["obs"])
+        tensor([[[ 0,  0,  0],
+                 [ 1,  1,  1],
+                 [ 2,  2,  2],
+                 [ 0,  0,  0],
+                 [ 0,  0,  0]],
+        <BLANKLINE>
+                [[ 3,  3,  3],
+                 [ 4,  4,  4],
+                 [ 0,  0,  0],
+                 [ 0,  0,  0],
+                 [ 0,  0,  0]],
+        <BLANKLINE>
+                [[ 5,  5,  5],
+                 [ 6,  6,  6],
+                 [ 0,  0,  0],
+                 [ 0,  0,  0],
+                 [ 0,  0,  0]],
+        <BLANKLINE>
+                [[ 7,  7,  7],
+                 [ 8,  8,  8],
+                 [ 9,  9,  9],
+                 [ 0,  0,  0],
+                 [ 0,  0,  0]],
+        <BLANKLINE>
+                [[10, 10, 10],
+                 [11, 11, 11],
+                 [12, 12, 12],
+                 [13, 13, 13],
+                 [14, 14, 14]],
+        <BLANKLINE>
+                [[15, 15, 15],
+                 [16, 16, 16],
+                 [17, 17, 17],
+                 [18, 18, 18],
+                 [19, 19, 19]]])
+
+    """
+    tensor = _flatten_batch(tensor)
+    max_val = max(splits)
+    mask = torch.zeros(len(splits), max_val, dtype=torch.bool, device=tensor.device)
+    mask.scatter_(
+        index=max_val - splits.unsqueeze(-1),
+        dim=1,
+        value=1,
+    )
+    mask = mask.cumsum(-1).flip(-1).bool()
+
+    def _fill_tensor(tensor):
+        empty_tensor = torch.zeros(
+            len(splits),
+            max_val,
+            *tensor.shape[1:],
+            dtype=tensor.dtype,
+            device=tensor.device,
+        )
+        empty_tensor[mask] = tensor
+        return empty_tensor
+
+    if isinstance(tensor, TensorDictBase):
+        tensor = tensor.apply(_fill_tensor, batch_size=[len(splits), max_val])
+    else:
+        tensor = _fill_tensor(tensor)
+    return tensor
+
+
+def _inv_pad_sequence(tensor, splits):
+    """Inverse a pad_sequence operation.
+
+    If tensor is of shape [N, T], than splits must be of of shape [N] with all elements
+    and integer between [1, T].
+
+    Examples:
+        >>> rewards = torch.randn(100, 20)
+        >>> num_per_traj = _get_num_per_traj(torch.zeros(100, 20).bernoulli_(0.1))
+        >>> padded = _split_and_pad_sequence(rewards, num_per_traj.tolist())
+        >>> reconstructed = _inv_pad_sequence(padded, num_per_traj)
+        >>> assert (reconstructed==rewards).all()
+    """
+    if splits.numel() == 1:
+        return tensor
+
+    splits = splits.expand(tensor.shape[::-1]).T
+    decay = splits - torch.arange(0, tensor.shape[-1], device=tensor.device)
+    idx = (decay >= 1).reshape(-1)
+    return tensor.reshape(-1)[idx]


Test patch:
diff --git a/benchmarks/test_objectives_benchmarks.py b/benchmarks/test_objectives_benchmarks.py
--- a/benchmarks/test_objectives_benchmarks.py
+++ b/benchmarks/test_objectives_benchmarks.py
@@ -75,6 +75,43 @@ def test_values(benchmark, val_fn, has_lmbda, has_state_value):
     )
 
 
+@pytest.mark.parametrize(
+    "gae_fn,gamma_tensor,batches,timesteps",
+    [
+        [generalized_advantage_estimate, False, 1, 512],
+        [vec_generalized_advantage_estimate, True, 1, 512],
+        [vec_generalized_advantage_estimate, False, 1, 512],
+        [vec_generalized_advantage_estimate, True, 32, 512],
+        [vec_generalized_advantage_estimate, False, 32, 512],
+    ],
+)
+def test_gae_speed(benchmark, gae_fn, gamma_tensor, batches, timesteps):
+    size = (batches, timesteps, 1)
+    print(size)
+
+    torch.manual_seed(0)
+    device = "cuda:0" if torch.cuda.device_count() else "cpu"
+    values = torch.randn(*size, device=device)
+    next_values = torch.randn(*size, device=device)
+    reward = torch.randn(*size, device=device)
+    done = torch.zeros(*size, dtype=torch.bool, device=device).bernoulli_(0.1)
+
+    gamma = 0.99
+    if gamma_tensor:
+        gamma = torch.full(size, gamma)
+    lmbda = 0.95
+
+    benchmark(
+        gae_fn,
+        gamma=gamma,
+        lmbda=lmbda,
+        state_value=values,
+        next_state_value=next_values,
+        reward=reward,
+        done=done,
+    )
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -4076,7 +4076,7 @@ def test_td1_multi(
     @pytest.mark.parametrize("device", get_available_devices())
     @pytest.mark.parametrize("gamma", [0.99, 0.5, 0.1])
     @pytest.mark.parametrize("lmbda", [0.99, 0.5, 0.1])
-    @pytest.mark.parametrize("N", [(3,), (7, 3)])
+    @pytest.mark.parametrize("N", [(1,), (3,), (7, 3)])
     @pytest.mark.parametrize("T", [200, 5, 3])
     @pytest.mark.parametrize("dtype", [torch.float, torch.double])
     @pytest.mark.parametrize("has_done", [True, False])
@@ -4098,6 +4098,54 @@ def test_gae(self, device, gamma, lmbda, N, T, dtype, has_done):
         )
         torch.testing.assert_close(r1, r2, rtol=1e-4, atol=1e-4)
 
+    @pytest.mark.parametrize("device", get_available_devices())
+    @pytest.mark.parametrize("N", [(1,), (8,), (7, 3)])
+    @pytest.mark.parametrize("dtype", [torch.float, torch.double])
+    @pytest.mark.parametrize("has_done", [True, False])
+    @pytest.mark.parametrize(
+        "gamma_tensor", ["scalar", "tensor", "tensor_single_element"]
+    )
+    @pytest.mark.parametrize(
+        "lmbda_tensor", ["scalar", "tensor", "tensor_single_element"]
+    )
+    def test_gae_param_as_tensor(
+        self, device, N, dtype, has_done, gamma_tensor, lmbda_tensor
+    ):
+        torch.manual_seed(0)
+
+        gamma = 0.95
+        lmbda = 0.90
+        T = 200
+
+        done = torch.zeros(*N, T, 1, device=device, dtype=torch.bool)
+        if has_done:
+            done = done.bernoulli_(0.1)
+        reward = torch.randn(*N, T, 1, device=device, dtype=dtype)
+        state_value = torch.randn(*N, T, 1, device=device, dtype=dtype)
+        next_state_value = torch.randn(*N, T, 1, device=device, dtype=dtype)
+
+        if gamma_tensor == "tensor":
+            gamma_vec = torch.full_like(reward, gamma)
+        elif gamma_tensor == "tensor_single_element":
+            gamma_vec = torch.as_tensor([gamma], device=device)
+        else:
+            gamma_vec = gamma
+
+        if lmbda_tensor == "tensor":
+            lmbda_vec = torch.full_like(reward, lmbda)
+        elif gamma_tensor == "tensor_single_element":
+            lmbda_vec = torch.as_tensor([lmbda], device=device)
+        else:
+            lmbda_vec = lmbda
+
+        r1 = vec_generalized_advantage_estimate(
+            gamma_vec, lmbda_vec, state_value, next_state_value, reward, done
+        )
+        r2 = generalized_advantage_estimate(
+            gamma, lmbda, state_value, next_state_value, reward, done
+        )
+        torch.testing.assert_close(r1, r2, rtol=1e-4, atol=1e-4)
+
     @pytest.mark.parametrize("device", get_available_devices())
     @pytest.mark.parametrize("gamma", [0.99, 0.5, 0.1])
     @pytest.mark.parametrize("lmbda", [0.99, 0.5, 0.1])