     let finalValue = networkOutput.value[sequenceLength]
 
     // Estimate the advantages for the provided trajectory.
-    let rewards = trajectory.reward[0..<sequenceLength]
+    let rewards = rewardsPreprocessor(trajectory.reward[0..<sequenceLength])
     let advantageEstimate = advantageFunction(
       stepKinds: stepKinds,
       rewards: rewards,

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

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