class_docstr: MPO learner (discrete or continuous, distributional or not).
class_docstr_tok: ["MPO","learner","(","discrete","or","continuous",",","distributional","or","not",")","."]
code: def __init__(  # pytype: disable=annotation-type-mismatch  # numpy-scalars
      self,
      critic_type: CriticType,
      discrete_policy: bool,
      environment_spec: specs.EnvironmentSpec,
      networks: mpo_networks.MPONetworks,
      random_key: jax_types.PRNGKey,
      discount: float,
      num_samples: int,
      iterator: Iterator[reverb.ReplaySample],
      experience_type: mpo_types.ExperienceType,
      loss_scales: mpo_types.LossScalesConfig,
      target_update_period: Optional[int] = 100,
      target_update_rate: Optional[float] = None,
      sgd_steps_per_learner_step: int = 20,
      policy_eval_stochastic: bool = True,
      policy_eval_num_val_samples: int = 128,
      policy_loss_config: Optional[mpo_types.PolicyLossConfig] = None,
      use_online_policy_to_bootstrap: bool = False,
      use_stale_state: bool = False,
      use_retrace: bool = False,
      retrace_lambda: float = 0.95,
      model_rollout_length: int = 0,
      optimizer: Optional[optax.GradientTransformation] = None,
      learning_rate: optax.ScalarOrSchedule = 1e-4,
      dual_optimizer: Optional[optax.GradientTransformation] = None,
      dual_learning_rate: optax.ScalarOrSchedule = 1e-2,
      grad_norm_clip: float = 40.0,
      reward_clip: float = np.float32("inf"),
      value_tx_pair: rlax.TxPair = rlax.IDENTITY_PAIR,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      devices: Optional[Sequence[jax.Device]] = None,
  ):
      self._critic_type = critic_type
      self._discrete_policy = discrete_policy
  
      process_id = jax.process_index()
      local_devices = jax.local_devices()
      self._devices = devices or local_devices
      logging.info(
          "Learner process id: %s. Devices passed: %s",
          process_id,
          devices,
      )
      logging.info(
          "Learner process id: %s. Local devices from JAX API: %s",
          process_id,
          local_devices,
      )
      self._local_devices = [
          d for d in self._devices if d in local_devices
      ]
  
      # Store networks.
      self._networks = networks
  
      # General learner book-keeping and loggers.
      self._counter = counter or counting.Counter()
      self._logger = logger
  
      # Other learner parameters.
      self._discount = discount
      self._num_samples = num_samples
      self._sgd_steps_per_learner_step = sgd_steps_per_learner_step
  
      self._policy_eval_stochastic = policy_eval_stochastic
      self._policy_eval_num_val_samples = policy_eval_num_val_samples
  
      self._reward_clip_range = sorted([-reward_clip, reward_clip])
      self._tx_pair = value_tx_pair
      self._loss_scales = loss_scales
      self._use_online_policy_to_bootstrap = (
          use_online_policy_to_bootstrap
      )
      self._model_rollout_length = model_rollout_length
  
      self._use_retrace = use_retrace
      self._retrace_lambda = retrace_lambda
      if use_retrace and critic_type == CriticType.MIXTURE_OF_GAUSSIANS:
          logging.warning(
              "Warning! Retrace has not been tested with the MoG critic."
          )
      self._use_stale_state = use_stale_state
  
      self._experience_type = experience_type
      if isinstance(self._experience_type, mpo_types.FromTransitions):
          # Each n=5-step transition will be converted to a length 2 sequence before
          # being passed to the loss, so we do n=1 step bootstrapping on the
          # resulting sequence to get n=5-step bootstrapping as intended.
          self._n_step_for_sequence_bootstrap = 1
          self._td_lambda = 1.0
      elif isinstance(self._experience_type, mpo_types.FromSequences):
          self._n_step_for_sequence_bootstrap = (
              self._experience_type.n_step
          )
          self._td_lambda = self._experience_type.td_lambda
  
      # Necessary to track when to update target networks.
      self._target_update_period = target_update_period
      self._target_update_rate = target_update_rate
      # Assert one and only one of target update period or rate is defined.
      if (target_update_period and target_update_rate) or (
          target_update_period is None and target_update_rate is None
      ):
          raise ValueError(
              "Exactly one of target_update_{period|rate} must be set."
              f" Received target_update_period={target_update_period} and"
              f" target_update_rate={target_update_rate}."
          )
  
      # Create policy loss.
      if self._discrete_policy:
          policy_loss_config = (
              policy_loss_config
              or mpo_types.CategoricalPolicyLossConfig()
          )
          self._policy_loss_module = discrete_losses.CategoricalMPO(
              **dataclasses.asdict(policy_loss_config)
          )
      else:
          policy_loss_config = (
              policy_loss_config or mpo_types.GaussianPolicyLossConfig()
          )
          self._policy_loss_module = continuous_losses.MPO(
              **dataclasses.asdict(policy_loss_config)
          )
  
      self._policy_loss_module.__call__ = jax.named_call(
          self._policy_loss_module.__call__, name="policy_loss"
      )
  
      # Create the dynamics model rollout loss.
      if model_rollout_length > 0:
          if not discrete_policy and (
              self._loss_scales.rollout.policy
              or self._loss_scales.rollout.bc_policy
          ):
              raise ValueError(
                  "Policy rollout losses are only supported in the "
                  "discrete policy case."
              )
          self._model_rollout_loss_fn = rollout_loss.RolloutLoss(
              dynamics_model=networks.dynamics_model,
              model_rollout_length=model_rollout_length,
              loss_scales=loss_scales,
              distributional_loss_fn=self._distributional_loss,
          )
  
      # Create optimizers if they aren't given.
      self._optimizer = optimizer or _get_default_optimizer(
          learning_rate, grad_norm_clip
      )
      self._dual_optimizer = dual_optimizer or _get_default_optimizer(
          dual_learning_rate, grad_norm_clip
      )
  
      self._action_spec = environment_spec.actions
  
      # Initialize random key for the rest of training.
      random_key, key = jax.random.split(random_key)
  
      # Initialize network parameters, ignoring the dummy initial state.
      network_params, _ = mpo_networks.init_params(
          self._networks,
          environment_spec,
          key,
          add_batch_dim=True,
          dynamics_rollout_length=self._model_rollout_length,
      )
  
      # Get action dims (unused in the discrete case).
      dummy_action = utils.zeros_like(environment_spec.actions)
      dummy_action_concat = utils.batch_concat(
          dummy_action, num_batch_dims=0
      )
  
      if isinstance(
          self._policy_loss_module, discrete_losses.CategoricalMPO
      ):
          self._dual_clip_fn = (
              discrete_losses.clip_categorical_mpo_params
          )
      elif isinstance(self._policy_loss_module, continuous_losses.MPO):
          is_constraining = (
              self._policy_loss_module.per_dim_constraining
          )
          self._dual_clip_fn = lambda dp: continuous_losses.clip_mpo_params(  # pylint: disable=g-long-lambda  # pytype: disable=wrong-arg-types  # numpy-scalars
              dp, per_dim_constraining=is_constraining
          )
  
      # Create dual parameters. In the discrete case, the action dim is unused.
      dual_params = self._policy_loss_module.init_params(
          action_dim=dummy_action_concat.shape[-1], dtype=jnp.float32
      )
  
      # Initialize optimizers.
      opt_state = self._optimizer.init(network_params)
      dual_opt_state = self._dual_optimizer.init(dual_params)
  
      # Initialise training state (parameters and optimiser state).
      state = TrainingState(
          params=network_params,
          target_params=network_params,
          dual_params=dual_params,
          opt_state=opt_state,
          dual_opt_state=dual_opt_state,
          steps=0,
          random_key=random_key,
      )
      self._state = utils.replicate_in_all_devices(
          state, self._local_devices
      )
  
      # Log how many parameters the network has.
      sizes = tree.map_structure(jnp.size, network_params)._asdict()
      num_params_by_component_str = " | ".join(
          [
              f"{key}: {sum(tree.flatten(size))}"
              for key, size in sizes.items()
          ]
      )
      logging.info(
          "Number of params by network component: %s",
          num_params_by_component_str,
      )
      logging.info(
          "Total number of params: %d",
          sum(tree.flatten(sizes.values())),
      )
  
      # Combine multiple SGD steps and pmap across devices.
      sgd_steps = utils.process_multiple_batches(
          self._sgd_step, self._sgd_steps_per_learner_step
      )
      self._sgd_steps = jax.pmap(
          sgd_steps, axis_name=_PMAP_AXIS_NAME, devices=self._devices
      )
  
      self._iterator = iterator
  
      # 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 = None
      self._current_step = 0
  
imports: import dataclasses
  import functools
  import time
  from typing import Any, Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple, Union
  from absl import logging
  import acme
  from acme import specs
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import categorical_mpo
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import rollout_loss
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  from acme.jax import types
  from acme.jax import utils
  import acme.jax.losses.mpo
  from acme.utils import counting
  from acme.utils import loggers
  import chex
  import jax
  import jax.numpy
  import numpy
  import optax
  import reverb
  import rlax
  import tree
label_desc: Use image node_id 1 to create a new MPOLearner object from inherited base classes: acme with example: obj = MPOLearner(critic_type, discrete_policy, environment_spec, networks, random_key, discount, num_samples, iterator, experience_type, loss_scales, target_update_period, target_update_rate, sgd_steps_per_learner_step, policy_eval_stochastic, policy_eval_num_val_samples, policy_loss_config, use_online_policy_to_bootstrap, use_stale_state, use_retrace, retrace_lambda, model_rollout_length, optimizer, learning_rate, dual_optimizer, dual_learning_rate, grad_norm_clip, reward_clip, value_tx_pair, counter, logger, devices)