code: def define_ppo_step(
      data_points,
      hparams,
      action_space,
      lr,
      epoch=-1,
      distributional_size=1,
      distributional_subscale=0.04,
  ):
      """Define ppo step."""
      del distributional_subscale
      (
          observation,
          action,
          discounted_reward,
          discounted_reward_probs,
          norm_advantage,
          old_pdf,
      ) = data_points
  
      obs_shape = common_layers.shape_list(observation)
      observation = tf.reshape(
          observation, [obs_shape[0] * obs_shape[1]] + obs_shape[2:]
      )
      (logits, new_value) = get_policy(
          observation,
          hparams,
          action_space,
          epoch=epoch,
          distributional_size=distributional_size,
      )
      logits = tf.reshape(logits, obs_shape[:2] + [action_space.n])
      new_policy_dist = tfp.distributions.Categorical(logits=logits)
  
      new_pdf = new_policy_dist.prob(action)
  
      ratio = new_pdf / old_pdf
      clipped_ratio = tf.clip_by_value(
          ratio, 1 - hparams.clipping_coef, 1 + hparams.clipping_coef
      )
  
      surrogate_objective = tf.minimum(
          clipped_ratio * norm_advantage, ratio * norm_advantage
      )
      policy_loss = -tf.reduce_mean(surrogate_objective)
  
      if distributional_size > 1:
          new_value = tf.reshape(
              new_value, obs_shape[:2] + [distributional_size]
          )
          new_value = tf.nn.log_softmax(new_value, axis=-1)
          value_shape = common_layers.shape_list(new_value)
          # The above is the new value distribution. We are also given as discounted
          # reward the value distribution and the corresponding probabilities.
          # The given discounted reward is already rounded to integers but in range
          # increased by 2x for greater fidelity. Increase range of new_values here.
          new_value_shifted = tf.concat(
              [new_value[1:], new_value[-1:]], axis=0
          )
          new_value_mean = (new_value + new_value_shifted) / 2
          new_value = tf.concat(
              [
                  tf.expand_dims(new_value, axis=-1),
                  tf.expand_dims(new_value_mean, axis=-1),
              ],
              -1,
          )
          new_value = tf.reshape(
              new_value, value_shape[:-1] + [2 * value_shape[-1]]
          )
          # Cast discounted reward to integers and gather the new log-probs for them.
          discounted_reward = tf.cast(discounted_reward, tf.int32)
          value_loss = tf.batch_gather(new_value, discounted_reward)
          # Weight the gathered (new) log-probs by the old probabilities.
          discounted_reward_probs = tf.expand_dims(
              discounted_reward_probs, axis=1
          )
          value_loss = -tf.reduce_sum(
              value_loss * discounted_reward_probs, axis=-1
          )
          # Take the mean over batch and time as final loss, multiply by coefficient.
          value_loss = hparams.value_loss_coef * tf.reduce_mean(
              value_loss
          )
      else:
          new_value = tf.reshape(new_value, obs_shape[:2])
          value_error = new_value - discounted_reward
          value_loss = hparams.value_loss_coef * tf.reduce_mean(
              value_error**2
          )
  
      entropy = new_policy_dist.entropy()
      entropy_loss = -hparams.entropy_loss_coef * tf.reduce_mean(
          entropy
      )
  
      losses = [policy_loss, value_loss, entropy_loss]
      loss = sum(losses)
      variables = tf.global_variables(hparams.policy_network + "/.*")
      train_op = optimize.optimize(
          loss, lr, hparams, variables=variables
      )
  
      with tf.control_dependencies([train_op]):
          return [tf.identity(x) for x in losses]
  
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 1 for calling a global function with example usage: define_ppo_step(data_points, hparams, action_space, lr, epoch, distributional_size, distributional_subscale) and returns: unknown