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 __init__(
      self,
      reward_objectives: Sequence[RewardObjective],
      qvalue_objectives: Sequence[QValueObjective],
      policy_network: snt.Module,
      critic_network: snt.Module,
      target_policy_network: snt.Module,
      target_critic_network: snt.Module,
      discount: float,
      num_samples: int,
      target_policy_update_period: int,
      target_critic_update_period: int,
      dataset: tf.data.Dataset,
      observation_network: types.TensorTransformation = tf.identity,
      target_observation_network: types.TensorTransformation = tf.identity,
      policy_loss_module: Optional[losses.MultiObjectiveMPO] = None,
      policy_optimizer: Optional[snt.Optimizer] = None,
      critic_optimizer: Optional[snt.Optimizer] = None,
      dual_optimizer: Optional[snt.Optimizer] = None,
      clipping: bool = True,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      checkpoint: bool = True,
  ):
      # Store online and target networks.
      self._policy_network = policy_network
      self._critic_network = critic_network
      self._target_policy_network = target_policy_network
      self._target_critic_network = target_critic_network
  
      # Make sure observation networks are snt.Module's so they have variables.
      self._observation_network = tf2_utils.to_sonnet_module(
          observation_network
      )
      self._target_observation_network = tf2_utils.to_sonnet_module(
          target_observation_network
      )
  
      # General learner book-keeping and loggers.
      self._counter = counter or counting.Counter()
      self._logger = logger or loggers.make_default_logger("learner")
  
      # Other learner parameters.
      self._discount = discount
      self._num_samples = num_samples
      self._clipping = clipping
  
      # Necessary to track when to update target networks.
      self._num_steps = tf.Variable(0, dtype=tf.int32)
      self._target_policy_update_period = target_policy_update_period
      self._target_critic_update_period = target_critic_update_period
  
      # Batch dataset and create iterator.
      # TODO(b/155086959): Fix type stubs and remove.
      self._iterator = iter(dataset)  # pytype: disable=wrong-arg-types
  
      # Store objectives
      self._reward_objectives = reward_objectives
      self._qvalue_objectives = qvalue_objectives
      if self._qvalue_objectives is None:
          self._qvalue_objectives = []
      self._num_critic_heads = len(self._reward_objectives)  # C
      self._objective_names = [
          x.name for x in self._reward_objectives
      ] + [x.name for x in self._qvalue_objectives]
  
      self._policy_loss_module = (
          policy_loss_module
          or losses.MultiObjectiveMPO(
              epsilons=[
                  losses.KLConstraint(name, _DEFAULT_EPSILON)
                  for name in self._objective_names
              ],
              epsilon_mean=_DEFAULT_EPSILON_MEAN,
              epsilon_stddev=_DEFAULT_EPSILON_STDDEV,
              init_log_temperature=_DEFAULT_INIT_LOG_TEMPERATURE,
              init_log_alpha_mean=_DEFAULT_INIT_LOG_ALPHA_MEAN,
              init_log_alpha_stddev=_DEFAULT_INIT_LOG_ALPHA_STDDEV,
          )
      )
  
      # Check that ordering of objectives matches the policy_loss_module's
      if self._objective_names != list(
          self._policy_loss_module.objective_names
      ):
          raise ValueError(
              "Agent's ordering of objectives doesn't match "
              "the policy loss module's ordering of epsilons."
          )
  
      # Create the optimizers.
      self._critic_optimizer = critic_optimizer or snt.optimizers.Adam(
          1e-4
      )
      self._policy_optimizer = policy_optimizer or snt.optimizers.Adam(
          1e-4
      )
      self._dual_optimizer = dual_optimizer or snt.optimizers.Adam(1e-2)
  
      # Expose the variables.
      policy_network_to_expose = snt.Sequential(
          [
              self._target_observation_network,
              self._target_policy_network,
          ]
      )
      self._variables = {
          "critic": self._target_critic_network.variables,
          "policy": policy_network_to_expose.variables,
      }
  
      # Create a checkpointer and snapshotter object.
      self._checkpointer = None
      self._snapshotter = None
  
      if checkpoint:
          self._checkpointer = tf2_savers.Checkpointer(
              subdirectory="mompo_learner",
              objects_to_save={
                  "counter": self._counter,
                  "policy": self._policy_network,
                  "critic": self._critic_network,
                  "observation": self._observation_network,
                  "target_policy": self._target_policy_network,
                  "target_critic": self._target_critic_network,
                  "target_observation": self._target_observation_network,
                  "policy_optimizer": self._policy_optimizer,
                  "critic_optimizer": self._critic_optimizer,
                  "dual_optimizer": self._dual_optimizer,
                  "policy_loss_module": self._policy_loss_module,
                  "num_steps": self._num_steps,
              },
          )
  
          self._snapshotter = tf2_savers.Snapshotter(
              objects_to_save={
                  "policy": snt.Sequential(
                      [
                          self._target_observation_network,
                          self._target_policy_network,
                      ]
                  ),
              }
          )
  
      # Do not record timestamps until after the first learning step is done.
      # This is to avoid including the time it takes for actors to come online and
      # fill the replay buffer.
      self._timestamp: float = None
  
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 1 to create a new MultiObjectiveMPOLearner object from inherited base classes: acme with example: obj = MultiObjectiveMPOLearner(reward_objectives, qvalue_objectives, policy_network, critic_network, target_policy_network, target_critic_network, discount, num_samples, target_policy_update_period, target_critic_update_period, dataset, observation_network, target_observation_network, policy_loss_module, policy_optimizer, critic_optimizer, dual_optimizer, clipping, counter, logger, checkpoint)