def gae(data: tuple, gamma: float = 0.99, lambda_: float = 0.97) -> torch.FloatTensor:
    """
    **Overview**:
        Implementation of the Generalized Advantage Estimator (GAE) as proposed in arXiv:1506.02438.
        This function calculates the advantages, which are used to update policy parameters in reinforcement learning.

    Arguments:
        - data (:obj:`namedtuple`): Tuple containing trajectory data including state values, next state
            values, rewards, done flags, and trajectory flags. Please note that the ``done`` flag signals
            the termination of an episode, whereas the ``traj_flag`` indicates the completion of a trajectory, which represents a segment within an episode.
        - gamma (:obj:`float`): Discount factor for future rewards, should be in the range [0, 1]. Default is 0.99.
        - lambda_ (:obj:`float`): The decay rate for the GAE, should be in the range [0, 1]. Default is 0.97.
            As lambda approaches 0, it introduces bias, and as lambda approaches 1, it increases variance
            due to the cumulative effect of terms.
    Returns:
        - adv (:obj:`torch.FloatTensor`): The calculated advantage estimates.
    Shapes:
        - value (:obj:`torch.FloatTensor`): Size of (T, B), where T is the length of the trajectory and B is the batch size.
        - next_value (:obj:`torch.FloatTensor`):  Size of (T, B)
        - reward (:obj:`torch.FloatTensor`): Size of (T, B)
        - adv (:obj:`torch.FloatTensor`): Size of (T, B)
    """
    # Unpack the input data.
    value, next_value, reward, done, traj_flag = data

    # Convert the done and trajectory flags to tensor format.
    done = torch.tensor(done).float()
    traj_flag = torch.tensor(traj_flag).float()

    # Expand ``done`` for possible broadcast operation in multi-agent cases
    if len(value.shape) == 2:
        done = done.unsqueeze(1)

    # If done equals 1, it indicates the end of an episode, thus the next state value should be 0.
    next_value *= (1 - done)

    # Calculate the temporal difference (TD) error for each time step.
    # $$\delta_t=-V_{\phi}(s_t)+r_t+V_{\phi}(s_{t+1})$$
    delta = reward + gamma * next_value - value

    # Set the GAE decay factor. If traj_flag equals 1, the factor will be 0. Otherwise, the factor is gamma * lambda.
    factor = gamma * lambda_ * (1 - traj_flag)

    # Initialize the advantage tensor.
    adv = torch.zeros_like(value)
    # Calculate ``adv`` in a reversed sequence.
    # Consider the definition of GAE: $$A^{GAE}_t = \sum_{i=1}\gamma^{i-1}\lambda^{i-1}\delta_{t+i-1}$$
    # Rewrite the equation above in a recurrent form, we finally have: $$A^{GAE}_t = \delta_t + \gamma\lambda A^{GAE}_{t+1}$$
    gae_item = torch.zeros_like(value[0])
    for t in reversed(range(reward.shape[0])):
        gae_item = delta[t] + factor[t] * gae_item
        adv[t] = gae_item

    # Return the calculated advantage estimates.
    return adv

    **Overview**:
        Implementation of the Generalized Advantage Estimator (GAE) as proposed in arXiv:1506.02438.
        This function calculates the advantages, which are used to update policy parameters in reinforcement learning.

    Arguments:
        - data (:obj:`namedtuple`): Tuple containing trajectory data including state values, next state
            values, rewards, done flags, and trajectory flags. Please note that the ``done`` flag signals
            the termination of an episode, whereas the ``traj_flag`` indicates the completion of a trajectory, which represents a segment within an episode.
        - gamma (:obj:`float`): Discount factor for future rewards, should be in the range [0, 1]. Default is 0.99.
        - lambda_ (:obj:`float`): The decay rate for the GAE, should be in the range [0, 1]. Default is 0.97.
            As lambda approaches 0, it introduces bias, and as lambda approaches 1, it increases variance
            due to the cumulative effect of terms.
    Returns:
        - adv (:obj:`torch.FloatTensor`): The calculated advantage estimates.
    Shapes:
        - value (:obj:`torch.FloatTensor`): Size of (T, B), where T is the length of the trajectory and B is the batch size.
        - next_value (:obj:`torch.FloatTensor`):  Size of (T, B)
        - reward (:obj:`torch.FloatTensor`): Size of (T, B)
        - adv (:obj:`torch.FloatTensor`): Size of (T, B)
    
chapter7_tricks/gae.py