   public let discountWeight: Float
 
   /// - Parameters:
-  ///   - discountFactor: Return discount factor value, which must be between `0.0` and `1.0`.
+  ///   - discountFactor: Reward discount factor value, which must be between `0.0` and `1.0`.
   ///   - discountWeight: A weight between `0.0` and `1.0` that is used for variance reduction in 
   ///     the temporal differences.
   public init(discountFactor: Float, discountWeight: Float = 1) {

        advantages.append(delta)
      }
    }
    return Tensor<Scalar>(advantages.reversed())
  }
}

      where: isLast[T - t - 1])
    discountedReturns.append(discountedReturn)
  }
  return Tensor(stacking: discountedReturns.reversed())
}

/// Advantage estimation result, which contains two tensors:
///   - `advantages`: Estimated advantages that are typically used to train actor networks.
///   - `discountedReturns`: Discounted returns that are typically used to train value networks.
public struct AdvantageEstimate<Scalar: TensorFlowFloatingPoint> {
  public let advantages: Tensor<Scalar>
  public let discountedReturns: Tensor<Scalar>

  public init(advantages: Tensor<Scalar>, discountedReturns: Tensor<Scalar>) {
    self.advantages = advantages
    self.discountedReturns = discountedReturns
  }
}

public protocol AdvantageFunction {
  /// - Parameters:
  ///   - stepKinds: Contains the step kinds (represented using their integer values) for each step.
  ///   - rewards: Contains the rewards obtained at each step.
  ///   - values: Contains the value estimates for each step.
  ///   - finalValue: Estimated value at the final step.
  func callAsFunction<Scalar: TensorFlowFloatingPoint>(
    stepKinds: StepKind,
    rewards: Tensor<Scalar>,
    values: Tensor<Scalar>,
    finalValue: Tensor<Scalar>
  ) -> AdvantageEstimate<Scalar>
}

/// Performs empirical advantage estimation.
///
/// The empirical advantage estimate at step `t` is defined as:
/// `advantage[t] = returns[t] - value[t]`, where the returns are computed using
/// `discountedReturns(discountFactor:stepKinds:rewards:finalValue:)`.
public struct EmpiricalAdvantageEstimation: AdvantageFunction {
  public let discountFactor: Float

  /// - Parameters:
  ///   - discountFactor: Reward discount factor value, which must be between `0.0` and `1.0`.
  public init(discountFactor: Float) {
    self.discountFactor = discountFactor
  }

  /// - Parameters:
  ///   - stepKinds: Contains the step kinds (represented using their integer values) for each step.
  ///   - rewards: Contains the rewards obtained at each step.
  ///   - values: Contains the value estimates for each step.
  ///   - finalValue: Estimated value at the final step.
  @inlinable
  public func callAsFunction<Scalar: TensorFlowFloatingPoint>(
    stepKinds: StepKind,
    rewards: Tensor<Scalar>,
    values: Tensor<Scalar>,
    finalValue: Tensor<Scalar>
  ) -> AdvantageEstimate<Scalar> {
    let returns = discountedReturns(
      discountFactor: Scalar(discountFactor),
      stepKinds: stepKinds,
      rewards: rewards,
      finalValue: finalValue)
    return AdvantageEstimate(
      advantages: returns - values,
      discountedReturns: returns)
  }
}

