class_docstr: The Proximal Policy Optimization Algorithm aka PPO.
  
  Trains policy and value models using the PPO algortithm.
class_docstr_tok: ["The","Proximal","Policy","Optimization","Algorithm","aka","PPO",".","Trains","policy","and","value","models","using","the","PPO","algortithm","."]
docstr: Joint policy and value loss.
docstr_tok: ["Joint","policy","and","value","loss","."]
code: def joint_loss(self):
      """Joint policy and value loss."""
  
      def f(
          dist_inputs,
          values,
          returns,
          dones,
          rewards,
          actions,
          old_log_probs,
          mask,
      ):
          """Definition of the Proximal Policy Optimization loss."""
          del mask  # TODO(lukaszkaiser): make PPO work with Transformer
          # We have dist_inputs of the shape float32[128,1,18]
          assert len(dist_inputs.shape) == 3, (
              f"dist_inputs.shape was {dist_inputs.shape}"
              f"but expected length of the tensor shape is 3"
          )
          # values of the shape float32[128,1,1]
          # returns of the shape float32[128,1,1]
          # dones of the shape int32[128,1,1]
          # rewards of the shape float32[128,1,1]
          # and old_log_probs of the shape float32[128,1]
          assert values.shape == returns.shape, (
              f"values.shape was {values.shape}"
              f"returns.shape was {returns.shape}"
          )
          assert values.shape == dones.shape, (
              f"values.shape was {values.shape}"
              f"returns.shape was {dones.shape}"
          )
          assert rewards.shape == dones.shape, (
              f"values.shape was {values.shape}"
              f"returns.shape was {dones.shape}"
          )
          assert returns.shape[0:2] == old_log_probs.shape, (
              f"returns.shape was {returns.shape}"
              f"old_log_probs.shape was {old_log_probs.shape}"
          )
  
          # actions is a tensor of the shape int32[128,1] in the case
          # of discrete actions and float32[128,1,6] in the case of
          # half-cheetah and other continuous actions
          # actions agree with returns/values on the first two coordinates
          # meaning batch and time
          assert actions.shape[0:2] == returns.shape[0:2], (
              f"actions.shape was {actions.shape} and "
              f"returns.shape was {returns.shape}"
          )
  
          ppo_objective = rl_layers.PPOObjective(
              dist_inputs,
              stop_gradient(values),
              returns,
              dones,
              rewards,
              actions,
              old_log_probs,
              log_prob_fun=self._policy_dist.log_prob,
              epsilon=self._epsilon,
              normalize_advantages=self._normalize_advantages,
          )
  
          # we insist that ppo_objective is a vector of shape [128,1]
          assert (
              len(ppo_objective.shape) == 2
          ), f"ppo_objective was {ppo_objective}"
          # which agrees with returns/values/actions on the first two coordinates
          assert ppo_objective.shape[0:2] == values.shape[0:2], (
              f"ppo_objective.shape was {ppo_objective.shape} and "
              f"values.shape was {values.shape}"
          )
  
          entropy_loss = rl_layers.EntropyLoss(
              dist_inputs,
              distribution=self._policy_dist,
              coeff=self._entropy_coeff,
          )
  
          assert (
              jnp.ndim(entropy_loss) == 0
          ), f"entropy_loss was {entropy_loss}"
  
          l2_value_loss = rl_layers.ValueLoss(
              values, returns, value_loss_coeff=self._value_loss_coeff
          )
  
          assert (
              jnp.ndim(l2_value_loss) == 0
          ), f"l2_value_loss was {l2_value_loss}"
  
          return -ppo_objective.mean() + l2_value_loss - entropy_loss
  
      return tl.Fn("PPOJointLoss", f)
  
imports: import functools
  from trax import data
  from trax import layers
  from trax import supervised
  from trax.fastmath import numpy
  from trax.fastmath import stop_gradient
  from trax.rl import actor_critic
  from trax.rl import distributions
  from trax.rl import rl_layers
  from trax.rl import training
  from trax.supervised import lr_schedules
label_desc: Use image node_id 3 for calling the PPOJoint obj's underlying member method code with example usage: obj.joint_loss() and returns: tl, unknown