def returns_minus_base(returns, b, masks):
    """
    Subtract optimal baseline from returns
    :param returns: returns from each step on or for whole trajectory. Size: batch_size x 1
    :param b: Optimal baseline per time step for every trajectory. Size: length_time_horizon
    :param masks: indicate the final step of each trajectory. 0 for last step and 1 for others. Size: batch_size x 1
    :return: Returns after subtraction
    """

    returns_sub_base = returns.new(returns.size(0), 1).fill_(0)

    traj_start = 0
    for i in range(len(returns_sub_base)):
        returns_sub_base[i] = returns[i] - b[i - traj_start]
        if masks[i] == 0:
            traj_start = i + 1

    return returns_sub_base