         // Value estimation loss term.
         let newValues = newNetworkOutput.value[0..<sequenceLength]
         var valueLoss = (newValues - returns).squared()
-        if let c = valueEstimationLoss.clipThreshold {
+        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 + valueEstimationLoss.weight * valueLoss.mean()
+        return loss + self.valueEstimationLoss.weight * valueLoss.mean()
       }
       if let clipNorm = maxGradientNorm {
         gradient.clipByGlobalNorm(clipNorm: clipNorm)

    
    var lastEpochLoss: Float = 0.0
    for _ in 0..<iterationCountPerUpdate {
      var (loss, gradient) = network.valueWithGradient { [
        clip, penalty, entropyRegularization, valueEstimationLoss
      ] network -> Tensor<Float> in
        let newNetworkOutput = network(trajectory.observation)

        // Compute the new action log probabilities.

      discountFactor: discountFactor,
      stepKinds: trajectory.stepKind,
      rewards: trajectory.reward)
    let (loss, gradient) = network.valueWithGradient { network -> Tensor<Float> in
      let actionDistribution = network(trajectory.observation)
      self.returnsNormalizer?.update(using: returns)
      if let normalizer = self.returnsNormalizer {
        returns = normalizer.normalize(returns)
      }
      let actionLogProbs = actionDistribution.logProbability(of: trajectory.action)

      // The policy gradient loss is defined as the sum, over time steps, of action
