code: def define_ppo_epoch(
      memory,
      hparams,
      action_space,
      batch_size,
      distributional_size=1,
      distributional_subscale=0.04,
      distributional_threshold=0.0,
      epoch=-1,
  ):
      """PPO epoch."""
      observation, reward, done, action, old_pdf, value_sm = memory
  
      # This is to avoid propagating gradients through simulated environment.
      observation = tf.stop_gradient(observation)
      action = tf.stop_gradient(action)
      reward = tf.stop_gradient(reward)
      if hasattr(hparams, "rewards_preprocessing_fun"):
          reward = hparams.rewards_preprocessing_fun(reward)
      done = tf.stop_gradient(done)
      value_sm = tf.stop_gradient(value_sm)
      old_pdf = tf.stop_gradient(old_pdf)
  
      value = value_sm
      if distributional_size > 1:
          value = _distributional_to_value(
              value_sm,
              distributional_size,
              distributional_subscale,
              distributional_threshold,
          )
  
      advantage = calculate_generalized_advantage_estimator(
          reward, value, done, hparams.gae_gamma, hparams.gae_lambda
      )
  
      if distributional_size > 1:
          # Create discounted reward values range.
          half = distributional_size // 2
          value_range = (
              tf.to_float(tf.range(-half, half)) + 0.5
          )  # Mid-bucket value.
          value_range *= distributional_subscale
          # Acquire new discounted rewards by using the above range as end-values.
          end_values = tf.expand_dims(value_range, 0)
          discounted_reward = discounted_rewards(
              reward, done, hparams.gae_gamma, end_values
          )
          # Re-normalize the discounted rewards to integers, in [0, dist_size] range.
          discounted_reward /= distributional_subscale
          discounted_reward += half
          discounted_reward = tf.maximum(discounted_reward, 0.0)
          discounted_reward = tf.minimum(
              discounted_reward, distributional_size
          )
          # Multiply the rewards by 2 for greater fidelity and round to integers.
          discounted_reward = tf.stop_gradient(
              tf.round(2 * discounted_reward)
          )
          # The probabilities corresponding to the end values from old predictions.
          discounted_reward_prob = tf.stop_gradient(value_sm[-1])
          discounted_reward_prob = tf.nn.softmax(
              discounted_reward_prob, axis=-1
          )
      else:
          discounted_reward = tf.stop_gradient(advantage + value[:-1])
          discounted_reward_prob = (
              discounted_reward  # Unused in this case.
          )
  
      advantage_mean, advantage_variance = tf.nn.moments(
          advantage, axes=[0, 1], keep_dims=True
      )
      advantage_normalized = tf.stop_gradient(
          (advantage - advantage_mean)
          / (tf.sqrt(advantage_variance) + 1e-8)
      )
  
      add_lists_elementwise = lambda l1, l2: [
          x + y for x, y in zip(l1, l2)
      ]
  
      number_of_batches = (
          (hparams.epoch_length - 1)
          * hparams.optimization_epochs
          // hparams.optimization_batch_size
      )
      epoch_length = hparams.epoch_length
      if hparams.effective_num_agents is not None:
          number_of_batches *= batch_size
          number_of_batches //= hparams.effective_num_agents
          epoch_length //= hparams.effective_num_agents
  
      assert (
          number_of_batches > 0
      ), "Set the paremeters so that number_of_batches>0"
      lr = learning_rate.learning_rate_schedule(hparams)
  
      shuffled_indices = [
          tf.random.shuffle(tf.range(epoch_length - 1))
          for _ in range(hparams.optimization_epochs)
      ]
      shuffled_indices = tf.concat(shuffled_indices, axis=0)
      shuffled_indices = shuffled_indices[
          : number_of_batches * hparams.optimization_batch_size
      ]
      indices_of_batches = tf.reshape(
          shuffled_indices, shape=(-1, hparams.optimization_batch_size)
      )
      input_tensors = [
          observation,
          action,
          discounted_reward,
          discounted_reward_prob,
          advantage_normalized,
          old_pdf,
      ]
  
      ppo_step_rets = tf.scan(
          lambda a, i: add_lists_elementwise(  # pylint: disable=g-long-lambda
              a,
              define_ppo_step(
                  [
                      tf.gather(t, indices_of_batches[i, :])
                      for t in input_tensors
                  ],
                  hparams,
                  action_space,
                  lr,
                  epoch=epoch,
                  distributional_size=distributional_size,
                  distributional_subscale=distributional_subscale,
              ),
          ),
          tf.range(number_of_batches),
          [0.0, 0.0, 0.0],
          parallel_iterations=1,
      )
  
      ppo_summaries = [
          tf.reduce_mean(ret) / number_of_batches
          for ret in ppo_step_rets
      ]
      ppo_summaries.append(lr)
      summaries_names = [
          "policy_loss",
          "value_loss",
          "entropy_loss",
          "learning_rate",
      ]
  
      summaries = [
          tf.summary.scalar(summary_name, summary)
          for summary_name, summary in zip(
              summaries_names, ppo_summaries
          )
      ]
      losses_summary = tf.summary.merge(summaries)
  
      for summary_name, summary in zip(summaries_names, ppo_summaries):
          losses_summary = tf.Print(
              losses_summary, [summary], summary_name + ": "
          )
  
      return losses_summary
  
file_path: tensor2tensor/tensor2tensor/rl/ppo.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_layers
  from tensor2tensor.models.research.rl import get_policy
  from tensor2tensor.utils import learning_rate
  from tensor2tensor.utils import optimize
  import tensorflow.compat.v1
  import tensorflow_probability
label_desc: Use image node_id 3 for calling a global function with example usage: define_ppo_epoch(memory, hparams, action_space, batch_size, distributional_size, distributional_subscale, distributional_threshold, epoch) and returns: losses_summary