class_docstr: Distributional (MoG) MPO learner.
code: def _step(self, inputs: reverb.ReplaySample) -> types.NestedTensor:
      # Get data from replay (dropping extras if any). Note there is no
      # extra data here because we do not insert any into Reverb.
      o_tm1, a_tm1, r_t, d_t, o_t = (
          inputs.data.observation,
          inputs.data.action,
          inputs.data.reward,
          inputs.data.discount,
          inputs.data.next_observation,
      )
  
      # Cast the additional discount to match the environment discount dtype.
      discount = tf.cast(self._discount, dtype=d_t.dtype)
  
      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(o_tm1)
          # 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(o_t))
  
          # 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_distributions = 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_distribution = self._critic_network(
              o_tm1, a_tm1
          )  # [B, ...]
  
          # Get the return distributions used in the policy evaluation bootstrap.
          if self._policy_evaluation_config.evaluate_stochastic_policy:
              z_distributions = sampled_q_t_distributions
              num_joint_samples = self._num_samples
          else:
              z_distributions = self._target_critic_network(
                  o_t, target_action_distribution.mean()
              )
              num_joint_samples = 1
  
          num_value_samples = (
              self._policy_evaluation_config.num_value_samples
          )
          num_joint_samples *= num_value_samples
          z_samples = z_distributions.sample(num_value_samples)
          z_samples = tf.reshape(z_samples, (num_joint_samples, -1, 1))
  
          # Expand dims of reward and discount tensors.
          reward = r_t[..., tf.newaxis]  # [B, 1]
          full_discount = discount * d_t[..., tf.newaxis]
          target_q = reward + full_discount * z_samples  # [N, B, 1]
          target_q = tf.stop_gradient(target_q)
  
          # Compute sample-based cross-entropy.
          log_probs_q = q_tm1_distribution.log_prob(
              target_q
          )  # [N, B, 1]
          critic_loss = -tf.reduce_mean(log_probs_q, axis=0)  # [B, 1]
          critic_loss = tf.reduce_mean(critic_loss)
  
          # Compute Q-values of sampled actions and reshape to [N, B].
          sampled_q_values = sampled_q_t_distributions.mean()
          sampled_q_values = tf.reshape(
              sampled_q_values, (self._num_samples, -1)
          )
  
          # 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_values,
          )
          policy_loss = tf.reduce_mean(policy_loss)
  
      # 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,
      }
      # Log MPO stats.
      fetches.update(policy_stats)
  
      return fetches
  
imports: import dataclasses
  import time
  from typing import List, Optional
  import acme
  from acme import types
  from acme.tf import losses
  from acme.tf import savers
  from acme.tf import utils
  from acme.utils import counting
  from acme.utils import loggers
  import numpy
  import reverb
  import sonnet
  import tensorflow
  import tensorflow_probability
label_desc: Use image node_id 2 for calling the MoGMPOLearner obj's underlying member method code with example usage: obj._step(inputs) and returns: fetches