class_docstr: RND learner.
class_docstr_tok: ["RND","learner","."]
code: def __init__(
      self,
      direct_rl_learner_factory: Callable[
          [Any, Iterator[reverb.ReplaySample]], acme.Learner
      ],
      iterator: Iterator[reverb.ReplaySample],
      optimizer: optax.GradientTransformation,
      rnd_network: rnd_networks.RNDNetworks,
      rng_key: jnp.ndarray,
      grad_updates_per_batch: int,
      is_sequence_based: bool,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
  ):
      self._is_sequence_based = is_sequence_based
  
      target_key, predictor_key = jax.random.split(rng_key)
      target_params = rnd_network.target.init(target_key)
      predictor_params = rnd_network.predictor.init(predictor_key)
      optimizer_state = optimizer.init(predictor_params)
  
      self._state = RNDTrainingState(
          optimizer_state=optimizer_state,
          params=predictor_params,
          target_params=target_params,
          steps=0,
      )
  
      # General learner book-keeping and loggers.
      self._counter = counter or counting.Counter()
      self._logger = logger or loggers.make_default_logger(
          "learner",
          asynchronous=True,
          serialize_fn=utils.fetch_devicearray,
          steps_key=self._counter.get_steps_key(),
      )
  
      loss = functools.partial(rnd_loss, networks=rnd_network)
      self._update = functools.partial(
          rnd_update_step, loss_fn=loss, optimizer=optimizer
      )
      self._update = utils.process_multiple_batches(
          self._update, grad_updates_per_batch
      )
      self._update = jax.jit(self._update)
  
      self._get_reward = jax.jit(
          functools.partial(
              rnd_networks.compute_rnd_reward, networks=rnd_network
          )
      )
  
      # Generator expression that works the same as an iterator.
      # https://pymbook.readthedocs.io/en/latest/igd.html#generator-expressions
      updated_iterator = (
          self._process_sample(sample) for sample in iterator
      )
  
      self._direct_rl_learner = direct_rl_learner_factory(
          rnd_network.direct_rl_networks, updated_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
  
imports: import functools
  import time
  from typing import Any, Callable, Dict, Iterator, List, NamedTuple, Optional, Tuple
  import acme
  from acme import types
  from acme.agents.jax.rnd import networks
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  from acme.utils import reverb_utils
  import jax
  import jax.numpy
  import optax
  import reverb
label_desc: Use image node_id 1 to create a new RNDLearner object from inherited base classes: acme with example: obj = RNDLearner(direct_rl_learner_factory, iterator, optimizer, rnd_network, rng_key, grad_updates_per_batch, is_sequence_based, counter, logger)