class_docstr: Distributional MPO learner.
  
  This is the learning component of a multi-objective MPO (MO-MPO) agent. Two
  sequences of objectives must be specified. Otherwise, the inputs are identical
  to those of the MPO / DMPO learners.
  
  Each objective must be defined as either a RewardObjective or an
  QValueObjective. These objectives are provided by the reward_objectives and
  qvalue_objectives parameters, respectively. For each RewardObjective, a critic
  will be trained to estimate Q-values for that objective. Whereas for each
  QValueObjective, the Q-values are computed directly by its qvalue_fn.
  
  A RewardObjective's reward_fn takes the observation, action, and environment
  reward as input, and returns the reward for that objective. For example, if
  the environment reward is a scalar, then an objective corresponding to the =
  task would simply return the environment reward.
  
  A QValueObjective's qvalue_fn takes the actions and reward-based objectives'
  Q-values as input, and outputs the "Q-values" for that objective. For
  instance, in the MO-MPO paper ([Abdolmaleki, Huang et al., 2020]), the action
  norm objective in the Humanoid run task is defined by setting the qvalue_fn
  to be the l2-norm of the actions.
  
  Note: If there is only one objective and that is the task reward, then this
  algorithm becomes exactly the same as (D)MPO.
  
  (Abdolmaleki, Huang et al., 2020): https://arxiv.org/pdf/2005.07513.pdf
code: def _step(self) -> types.NestedTensor:
      # Update target network.
      online_policy_variables = self._policy_network.variables
      target_policy_variables = self._target_policy_network.variables
      online_critic_variables = (
          *self._observation_network.variables,
          *self._critic_network.variables,
      )
      target_critic_variables = (
          *self._target_observation_network.variables,
          *self._target_critic_network.variables,
      )
  
      # Make online policy -> target policy network update ops.
      if (
          tf.math.mod(
              self._num_steps, self._target_policy_update_period
          )
          == 0
      ):
          for src, dest in zip(
              online_policy_variables, target_policy_variables
          ):
              dest.assign(src)
      # Make online critic -> target critic network update ops.
      if (
          tf.math.mod(
              self._num_steps, self._target_critic_update_period
          )
          == 0
      ):
          for src, dest in zip(
              online_critic_variables, target_critic_variables
          ):
              dest.assign(src)
  
      self._num_steps.assign_add(1)
  
      # Get data from replay (dropping extras if any). Note there is no
      # extra data here because we do not insert any into Reverb.
      inputs = next(self._iterator)
      transitions: types.Transition = inputs.data
  
      with tf.GradientTape(persistent=True) as tape:
          # Maybe transform the observation before feeding into policy and critic.
          # Transforming the observations this way at the start of the learning
          # step effectively means that the policy and critic share observation
          # network weights.
          o_tm1 = self._observation_network(transitions.observation)
          # This stop_gradient prevents gradients to propagate into the target
          # observation network. In addition, since the online policy network is
          # evaluated at o_t, this also means the policy loss does not influence
          # the observation network training.
          o_t = tf.stop_gradient(
              self._target_observation_network(
                  transitions.next_observation
              )
          )
  
          # Get online and target action distributions from policy networks.
          online_action_distribution = self._policy_network(o_t)
          target_action_distribution = self._target_policy_network(o_t)
  
          # Sample actions to evaluate policy; of size [N, B, ...].
          sampled_actions = target_action_distribution.sample(
              self._num_samples
          )
  
          # Tile embedded observations to feed into the target critic network.
          # Note: this is more efficient than tiling before the embedding layer.
          tiled_o_t = tf2_utils.tile_tensor(
              o_t, self._num_samples
          )  # [N, B, ...]
  
          # Compute target-estimated distributional value of sampled actions at o_t.
          sampled_q_t_all = self._target_critic_network(
              # Merge batch dimensions; to shape [N*B, ...].
              snt.merge_leading_dims(tiled_o_t, num_dims=2),
              snt.merge_leading_dims(sampled_actions, num_dims=2),
          )
  
          # Compute online critic value distribution of a_tm1 in state o_tm1.
          q_tm1_all = self._critic_network(o_tm1, transitions.action)
  
          # Compute rewards for objectives with defined reward_fn
          reward_stats = {}
          r_t_all = []
          for objective in self._reward_objectives:
              r = objective.reward_fn(
                  o_tm1, transitions.action, transitions.reward
              )
              reward_stats[
                  "{}_reward".format(objective.name)
              ] = tf.reduce_mean(r)
              r_t_all.append(r)
          r_t_all = tf.stack(r_t_all, axis=-1)
          r_t_all.get_shape().assert_has_rank(2)  # [B, C]
  
          if isinstance(
              sampled_q_t_all, list
          ):  # Distributional critics
              (
                  critic_loss,
                  sampled_q_t,
              ) = _compute_distributional_critic_loss(
                  sampled_q_t_all,
                  q_tm1_all,
                  r_t_all,
                  transitions.discount,
                  self._discount,
                  self._num_samples,
              )
          else:
              critic_loss, sampled_q_t = _compute_critic_loss(
                  sampled_q_t_all,
                  q_tm1_all,
                  r_t_all,
                  transitions.discount,
                  self._discount,
                  self._num_samples,
                  self._num_critic_heads,
              )
  
          # Add sampled Q-values for objectives with defined qvalue_fn
          sampled_q_t_k = [sampled_q_t]
          for objective in self._qvalue_objectives:
              sampled_q_t_k.append(
                  tf.expand_dims(
                      tf.stop_gradient(
                          objective.qvalue_fn(
                              sampled_actions, sampled_q_t
                          )
                      ),
                      axis=-1,
                  )
              )
          sampled_q_t_k = tf.concat(sampled_q_t_k, axis=-1)  # [N, B, K]
  
          # Compute MPO policy loss.
          policy_loss, policy_stats = self._policy_loss_module(
              online_action_distribution=online_action_distribution,
              target_action_distribution=target_action_distribution,
              actions=sampled_actions,
              q_values=sampled_q_t_k,
          )
  
      # For clarity, explicitly define which variables are trained by which loss.
      critic_trainable_variables = (
          # In this agent, the critic loss trains the observation network.
          self._observation_network.trainable_variables
          + self._critic_network.trainable_variables
      )
      policy_trainable_variables = (
          self._policy_network.trainable_variables
      )
      # The following are the MPO dual variables, stored in the loss module.
      dual_trainable_variables = (
          self._policy_loss_module.trainable_variables
      )
  
      # Compute gradients.
      critic_gradients = tape.gradient(
          critic_loss, critic_trainable_variables
      )
      policy_gradients, dual_gradients = tape.gradient(
          policy_loss,
          (policy_trainable_variables, dual_trainable_variables),
      )
  
      # Delete the tape manually because of the persistent=True flag.
      del tape
  
      # Maybe clip gradients.
      if self._clipping:
          policy_gradients = tuple(
              tf.clip_by_global_norm(policy_gradients, 40.0)[0]
          )
          critic_gradients = tuple(
              tf.clip_by_global_norm(critic_gradients, 40.0)[0]
          )
  
      # Apply gradients.
      self._critic_optimizer.apply(
          critic_gradients, critic_trainable_variables
      )
      self._policy_optimizer.apply(
          policy_gradients, policy_trainable_variables
      )
      self._dual_optimizer.apply(
          dual_gradients, dual_trainable_variables
      )
  
      # Losses to track.
      fetches = {
          "critic_loss": critic_loss,
          "policy_loss": policy_loss,
      }
      fetches.update(policy_stats)  # Log MPO stats.
      fetches.update(reward_stats)  # Log reward stats.
  
      return fetches
  
imports: import dataclasses
  import time
  from typing import Callable, List, Optional, Sequence
  import acme
  from acme import types
  from acme.tf import losses
  from acme.tf import networks
  from acme.tf import savers
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import sonnet
  import tensorflow
  import trfl
label_desc: Use image node_id 2 for calling the MultiObjectiveMPOLearner obj's underlying member method code with example usage: obj._step() and returns: fetches