         initialState: network.initialState(batchSize: batchSize),
         optimizer: { AMSGrad(for: $0) },
         learningRate: learningRate,
+        maxGradientNorm: 0.5,
         advantageFunction: advantageFunction,
         advantagesNormalizer: nil,
         useTDLambdaReturn: true,

    observation: Observation,
    batchSize: Int
   ) -> AnyAgent<JellyBeanWorld.Environment, LSTMCell<Float>.State> {
    let learningRate = ExponentiallyDecayedLearningRate(
      baseLearningRate: FixedLearningRate(Float(1e-4)),
      decayRate: 0.999,
      decayStepCount: 1)
    let advantageFunction = GeneralizedAdvantageEstimation(
      discountFactor: 0.99,
      discountWeight: 0.95)
    let ppoClip = PPOClip(epsilon: 0.1)
    let ppoPenalty = PPOPenalty(klCutoffFactor: 0.5)
    let ppoValueEstimationLoss = PPOValueEstimationLoss(weight: 0.5, clipThreshold: 0.1)

        initialState: network.initialState(batchSize: batchSize),
        optimizer: { AMSGrad(for: $0) },
        learningRate: learningRate,
        maxGradientNorm: 0.5,
        advantageFunction: advantageFunction,
        advantagesNormalizer: nil,
        useTDLambdaReturn: true,
