       let klDivergence = network(trajectory.observation).actionDistribution.klDivergence(
         to: actionDistribution)
       let klMean = klDivergence.mean().scalarized()
-      if klMean < p.adaptiveKLTarget * (1 - p.adaptiveKLTolerance) {
-        p.adaptiveKLBeta = max(beta / 1.5, 1e-16)
-      } else if klMean > p.adaptiveKLTarget * (1 + p.adaptiveKLTolerance) {
-        p.adaptiveKLBeta = max(beta * 1.5, 1e-16)
+      if klMean < p.adaptiveKLTarget / p.adaptiveKLToleranceFactor {
+        p.adaptiveKLBeta = max(beta / p.adaptiveKLBetaScalingFactor, 1e-16)
+      } else if klMean > p.adaptiveKLTarget * p.adaptiveKLToleranceFactor {
+        p.adaptiveKLBeta = beta * p.adaptiveKLBetaScalingFactor
       }
     }
 

    klCutoffFactor: Float = 0.2,
    klCutoffCoefficient: Float = 1000.0,
    adaptiveKLTarget: Float = 0.01,
    adaptiveKLTolerance: Float = 0.3,
    adaptiveKLBeta: Float? = 1.0
  ) {
    self.klCutoffFactor = klCutoffFactor
    self.klCutoffCoefficient = klCutoffCoefficient
    self.adaptiveKLTarget = adaptiveKLTarget
    self.adaptiveKLTolerance = adaptiveKLTolerance
    self.adaptiveKLBeta = adaptiveKLBeta
  }
}

  public let klCutoffFactor: Float
  public let klCutoffCoefficient: Float
  public let adaptiveKLTarget: Float
  public let adaptiveKLToleranceFactor: Float
  public let adaptiveKLBetaScalingFactor: Float

  public fileprivate(set) var adaptiveKLBeta: Float?

