     self.optimizer = optimizer(network)
     self.learningRate = learningRate
     self.maxGradientNorm = maxGradientNorm
+    self.rewardsPreprocessor = rewardsPreprocessor
     self.advantageFunction = advantageFunction
     self.advantagesNormalizer = advantagesNormalizer
     self.useTDLambdaReturn = useTDLambdaReturn

    let finalValue = networkOutput.value[sequenceLength]

    // Estimate the advantages for the provided trajectory.
    let rewards = trajectory.reward[0..<sequenceLength]
    let advantageEstimate = advantageFunction(
      stepKinds: stepKinds,
      rewards: rewards,

    optimizer: (Network) -> Optimizer,
    learningRate: LearningRate,
    maxGradientNorm: Float? = 0.5,
    rewardsPreprocessor: @escaping (Tensor<Float>) -> Tensor<Float> = { $0 },
    advantageFunction: AdvantageFunction = GeneralizedAdvantageEstimation(
      discountFactor: 0.99,
      discountWeight: 0.95),
