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

      discountFactor: discountFactor,
      stepKinds: trajectory.stepKind,
      rewards: trajectory.reward)
    let (loss, gradient) = network.valueWithGradient { [
      // TODO: !!!!
      // inout returnsNormalizer,
      entropyRegularizationWeight
    ] network -> Tensor<Float> in
      let actionDistribution = network(trajectory.observation)
      // returnsNormalizer?.update(using: returns)
      // if let normalizer = 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

        }

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