import torch


def compute_target(
    value_final: torch.Tensor,
    rewards: torch.Tensor,
    masks: torch.Tensor,
    gamma: float,
) -> torch.Tensor:
    """
    Compute target (sum of total discounted rewards) for rollout.

    Parameters
    -----------
    value_final: torch.Tensor
        state values from final time step of rollout, size (num_env,)
    rewards: torch.Tensor
        rewards across rollout, size (rollout_steps, num_env)
    masks: torch.Tensor
        masks for episode end states, 0 if end state, 1 otherwise,
        size (rollout_steps, num_env)
    gamma: float
        discount factor for rollout

    Returns
    -------
    torch.Tensor: targets, size (rollout_steps, num_env)
    """
    G = value_final
    T = rewards.shape[0]
    targets = torch.zeros(rewards.shape)

    for i in range(T - 1, -1, -1):
        G = rewards[i] + gamma * G * masks[i]
        targets[i] = G

    return targets