class_docstr: R2D2 learner.
  
  This is the learning component of the R2D2 agent. It takes a dataset as input
  and implements update functionality to learn from this dataset.
class_docstr_tok: ["R2D2","learner",".","This","is","the","learning","component","of","the","R2D2","agent",".","It","takes","a","dataset","as","input","and","implements","update","functionality","to","learn","from","this","dataset","."]
code: def __init__(
      self,
      environment_spec: specs.EnvironmentSpec,
      network: Union[networks.RNNCore, snt.RNNCore],
      target_network: Union[networks.RNNCore, snt.RNNCore],
      burn_in_length: int,
      sequence_length: int,
      dataset: tf.data.Dataset,
      reverb_client: Optional[reverb.TFClient] = None,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      discount: float = 0.99,
      target_update_period: int = 100,
      importance_sampling_exponent: float = 0.2,
      max_replay_size: int = 1_000_000,
      learning_rate: float = 1e-3,
      # TODO(sergomez): rename to use_core_state for consistency with JAX agent.
      store_lstm_state: bool = True,
      max_priority_weight: float = 0.9,
      n_step: int = 5,
      clip_grad_norm: Optional[float] = None,
  ):
      if not isinstance(network, networks.RNNCore):
          network.unroll = functools.partial(snt.static_unroll, network)
          target_network.unroll = functools.partial(
              snt.static_unroll, target_network
          )
  
      # Internalise agent components (replay buffer, networks, optimizer).
      # TODO(b/155086959): Fix type stubs and remove.
      self._iterator: Iterator[reverb.ReplaySample] = iter(
          dataset
      )  # pytype: disable=wrong-arg-types
      self._network = network
      self._target_network = target_network
      self._optimizer = snt.optimizers.Adam(learning_rate, epsilon=1e-3)
      self._reverb_client = reverb_client
  
      # Internalise the hyperparameters.
      self._store_lstm_state = store_lstm_state
      self._burn_in_length = burn_in_length
      self._discount = discount
      self._max_replay_size = max_replay_size
      self._importance_sampling_exponent = importance_sampling_exponent
      self._max_priority_weight = max_priority_weight
      self._target_update_period = target_update_period
      self._num_actions = environment_spec.actions.num_values
      self._sequence_length = sequence_length
      self._n_step = n_step
      self._clip_grad_norm = clip_grad_norm
  
      if burn_in_length:
          self._burn_in = lambda o, s: self._network.unroll(
              o, s, burn_in_length
          )
      else:
          self._burn_in = lambda o, s: (
              o,
              s,
          )  # pylint: disable=unnecessary-lambda
  
      # Learner state.
      self._variables = network.variables
      self._num_steps = tf.Variable(
          0.0, dtype=tf.float32, trainable=False, name="step"
      )
  
      # Internalise logging/counting objects.
      self._counter = counting.Counter(counter, "learner")
      self._logger = logger or loggers.TerminalLogger(
          "learner", time_delta=100.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
  
imports: import functools
  import time
  from typing import Dict, Iterator, List, Mapping, Union, Optional
  import acme
  from acme import specs
  from acme.adders import reverb
  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 reverb
  import sonnet
  import tensorflow
  import tree
label_desc: Use image node_id 1 to create a new R2D2Learner object from inherited base classes: acme, tf2_savers with example: obj = R2D2Learner(environment_spec, network, target_network, burn_in_length, sequence_length, dataset, reverb_client, counter, logger, discount, target_update_period, importance_sampling_exponent, max_replay_size, learning_rate, store_lstm_state, max_priority_weight, n_step, clip_grad_norm)