class_docstr: An Acme learner based around SGD on batches.
  
  This learner currently supports optional prioritized replay and assumes a
  TrainingState as described above.
docstr: Initialize the SGD learner.
docstr_tok: ["Initialize","the","SGD","learner","."]
code: def __init__(
      self,
      network: networks_lib.TypedFeedForwardNetwork,
      loss_fn: LossFn,
      optimizer: optax.GradientTransformation,
      data_iterator: Iterator[utils.PrefetchingSplit],
      target_update_period: int,
      random_key: networks_lib.PRNGKey,
      replay_client: Optional[reverb.Client] = None,
      replay_table_name: str = adders.DEFAULT_PRIORITY_TABLE,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      num_sgd_steps_per_step: int = 1,
  ):
      """Initialize the SGD learner."""
      self.network = network
  
      # Internalize the loss_fn with network.
      self._loss = jax.jit(functools.partial(loss_fn, self.network))
  
      # SGD performs the loss, optimizer update and periodic target net update.
      def sgd_step(
          state: TrainingState, batch: reverb.ReplaySample
      ) -> Tuple[TrainingState, LossExtra]:
          next_rng_key, rng_key = jax.random.split(state.rng_key)
          # Implements one SGD step of the loss and updates training state
          (loss, extra), grads = jax.value_and_grad(
              self._loss, has_aux=True
          )(state.params, state.target_params, batch, rng_key)
  
          loss = jax.lax.pmean(loss, axis_name=PMAP_AXIS_NAME)
          # Average gradients over pmap replicas before optimizer update.
          grads = jax.lax.pmean(grads, axis_name=PMAP_AXIS_NAME)
          # Apply the optimizer updates
          updates, new_opt_state = optimizer.update(
              grads, state.opt_state
          )
          new_params = optax.apply_updates(state.params, updates)
  
          extra.metrics.update({"total_loss": loss})
  
          # Periodically update target networks.
          steps = state.steps + 1
          target_params = optax.periodic_update(
              new_params,
              state.target_params,  # pytype: disable=wrong-arg-types  # numpy-scalars
              steps,
              target_update_period,
          )
  
          new_training_state = TrainingState(
              new_params,
              target_params,
              new_opt_state,
              steps,
              next_rng_key,
          )
          return new_training_state, extra
  
      def postprocess_aux(extra: LossExtra) -> LossExtra:
          reverb_priorities = jax.tree_util.tree_map(
              lambda a: jnp.reshape(a, (-1, *a.shape[2:])),
              extra.reverb_priorities,
          )
          return extra._replace(
              metrics=jax.tree_util.tree_map(jnp.mean, extra.metrics),
              reverb_priorities=reverb_priorities,
          )
  
      self._num_sgd_steps_per_step = num_sgd_steps_per_step
      sgd_step = utils.process_multiple_batches(
          sgd_step, num_sgd_steps_per_step, postprocess_aux
      )
      self._sgd_step = jax.pmap(
          sgd_step, axis_name=PMAP_AXIS_NAME, devices=jax.devices()
      )
  
      # Internalise agent components
      self._data_iterator = data_iterator
      self._target_update_period = target_update_period
      self._counter = counter or counting.Counter()
      self._logger = logger or loggers.TerminalLogger(
          "learner", time_delta=1.0
      )
  
      # 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
  
      # Initialize the network parameters
      key_params, key_target, key_state = jax.random.split(
          random_key, 3
      )
      initial_params = self.network.init(key_params)
      initial_target_params = self.network.init(key_target)
      state = TrainingState(
          params=initial_params,
          target_params=initial_target_params,
          opt_state=optimizer.init(initial_params),
          steps=0,
          rng_key=key_state,
      )
      self._state = utils.replicate_in_all_devices(
          state, jax.local_devices()
      )
  
      # Update replay priorities
      def update_priorities(reverb_update: ReverbUpdate) -> None:
          if replay_client is None:
              return
          keys, priorities = tree.map_structure(
              # Fetch array and combine device and batch dimensions.
              lambda x: utils.fetch_devicearray(x).reshape(
                  (-1,) + x.shape[2:]
              ),
              (reverb_update.keys, reverb_update.priorities),
          )
          replay_client.mutate_priorities(
              table=replay_table_name,
              updates=dict(zip(keys, priorities)),
          )
  
      self._replay_client = replay_client
      self._async_priority_updater = async_utils.AsyncExecutor(
          update_priorities
      )
  
      self._current_step = 0
  
imports: import functools
  import time
  from typing import Dict, Iterator, List, NamedTuple, Optional, Tuple
  import acme
  from acme.adders import reverb
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import async_utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import jax.numpy
  import optax
  import reverb
  import tree
  import typing_extensions
label_desc: Use image node_id 1 to create a new SGDLearner object from inherited base classes: acme with example: obj = SGDLearner(network, loss_fn, optimizer, data_iterator, target_update_period, random_key, replay_client, replay_table_name, counter, logger, num_sgd_steps_per_step)