         advantages.append(delta)
       }
     }
-    return Tensor<Scalar>(advantages.reversed())
+    // Compute discounted returns.
+    let returns = discountedReturns(
+      discountFactor: Scalar(discountFactor),
+      stepKinds: stepKinds,
+      rewards: rewards,
+      finalValue: finalValue)
+    return AdvantageEstimate(
+      advantages: Tensor(advantages.reversed()),
+      discountedReturns: returns)
   }
 }


  /// - Parameters:
  ///   - stepKinds: Contains the step kinds (represented using their integer values) for each step.
  ///   - returns: Contains the returns for each step. These are typically already discounted
  ///     (i.e., the result of calling `discount(discountFactor:stepKinds:returns:finalValue:)`).
  ///   - values: Contains the value estimates for each step.
  ///   - finalValue: Estimated value at the final step.
  @inlinable
  public func callAsFunction<Scalar: TensorFlowFloatingPoint>(
    stepKinds: StepKind,
    returns: Tensor<Scalar>,
    values: Tensor<Scalar>,
    finalValue: Tensor<Scalar>
  ) -> Tensor<Scalar> {
    let discountWeight = Scalar(self.discountWeight)
    let discountFactor = Scalar(self.discountFactor)
    let isLast = stepKinds.isLast()
    let T = stepKinds.rawValue.shape[0]
    var advantages = [Tensor<Scalar>]()
    for t in 0..<T {
      let nextValue = T - t < T ? values[T - t] : finalValue
      let delta = returns[T - t - 1] + discountFactor * nextValue - values[T - t - 1]
      if T - t < T {
        let nextAdvantage = advantages[t - 1].replacing(
          with: Tensor<Scalar>(zerosLike: advantages[t - 1]),

      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)
  }
}

