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

        }

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

        var loss = importanceRatio * advantages

        // Importance ratio clipping loss term.
        if let c = self.clip {
          let ε = Tensor<Float>(c.epsilon)
          let importanceRatioClipped = importanceRatio.clipped(min: 1 - ε, max: 1 + ε)
          loss = -min(loss, importanceRatioClipped * advantages).mean()
