         }
 
         // Entropy regularization loss term.
-        if let e = entropyRegularization {
+        if let e = self.entropyRegularization {
           let entropy = newActionDistribution.entropy()[0..<sequenceLength]
           loss = loss - e.weight * entropy.mean()
         }

        }

        // KL penalty loss term.
        if let p = penalty {
          let klDivergence = actionDistribution.klDivergence(to: newActionDistribution)
          let klMean = klDivergence.mean()
          let klCutoffLoss = max(klMean - p.klCutoffFactor * p.adaptiveKLTarget, 0).squared()

        // Value estimation loss term.
        let newValues = newNetworkOutput.value[0..<sequenceLength]
        var valueLoss = (newValues - returns).squared()
        if let c = self.valueEstimationLoss.clipThreshold {
          let ε = Tensor<Float>(c)
          let clippedValues = values + (newValues - values).clipped(min: -ε, max: ε)
          let clippedValueLoss = (clippedValues - returns).squared()
          valueLoss = max(valueLoss, clippedValueLoss)
        }
        return loss + self.valueEstimationLoss.weight * valueLoss.mean()
      }
      if let clipNorm = maxGradientNorm {
        gradient.clipByGlobalNorm(clipNorm: clipNorm)
