class_docstr: The A2C algorithm.
  
  Trains policy and value models using the A2C algortithm.
class_docstr_tok: ["The","A2C","algorithm",".","Trains","policy","and","value","models","using","the","A2C","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 A2C loss."""
          del old_log_probs
  
          # Typically 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]
          assert values.shape == returns.shape, (
              f"values.shape was {values.shape}"
              f"returns.shape was (returns.shape)"
          )
          # actions of the shape int32[128,1] in the case of discrete actions
          # and float32[128,1,6] in the case of of half-cheetah
          # actions agree with returns/values on the first two coordinates
          assert actions.shape[0:2] == returns.shape[0:2], (
              f"actions.shape was {actions.shape}"
              f"returns.shape was (returns.shape)"
          )
          # and mask of the shape float32[128,1]
          assert len(mask.shape) == 2, f"mask.shape was {mask.shape}"
          # which agrees with returns/values/actions on the first two coordinates
          assert mask.shape[0:2] == returns.shape[0:2], (
              f"mask.shape was {mask.shape}"
              f"returns.shape was (returns.shape)"
          )
  
          a2c_objective = rl_layers.A2CObjective(
              dist_inputs,
              stop_gradient(values),
              returns,
              dones,
              rewards,
              actions,
              mask,
              log_prob_fun=self._policy_dist.log_prob,
              normalize_advantages=self._normalize_advantages,
          )
  
          # we insist that a2c_objective is a scalar
          assert (
              jnp.ndim(a2c_objective) == 0
          ), f"a2c_objective was {a2c_objective}"
  
          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}"
  
          combined_loss = a2c_objective + l2_value_loss - entropy_loss
  
          return combined_loss
  
      return tl.Fn("A2CJointLoss", f, n_out=1)
  
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 A2CJoint obj's underlying member method code with example usage: obj.joint_loss() and returns: tl, combined_loss