class_docstr: D4PG learner.
  
  This is the learning component of a D4PG agent. IE it takes a dataset as input
  and implements update functionality to learn from this dataset.
class_docstr_tok: ["D4PG","learner",".","This","is","the","learning","component","of","a","D4PG","agent",".","IE","it","takes","a","dataset","as","input","and","implements","update","functionality","to","learn","from","this","dataset","."]
docstr: Initializes the learner.
  
  Args:
    policy_network: the online (optimized) policy.
    critic_network: the online critic.
    target_policy_network: the target policy (which lags behind the online
      policy).
    target_critic_network: the target critic.
    discount: discount to use for TD updates.
    target_update_period: number of learner steps to perform before updating
      the target networks.
    dataset_iterator: dataset to learn from, whether fixed or from a replay
      buffer (see `acme.datasets.reverb.make_reverb_dataset` documentation).
    replicator: Replicates variables and their update methods over multiple
      accelerators, such as the multiple chips in a TPU.
    observation_network: an optional online network to process observations
      before the policy and the critic.
    target_observation_network: the target observation network.
    policy_optimizer: the optimizer to be applied to the DPG (policy) loss.
    critic_optimizer: the optimizer to be applied to the distributional
      Bellman loss.
    clipping: whether to clip gradients by global norm.
    counter: counter object used to keep track of steps.
    logger: logger object to be used by learner.
    checkpoint: boolean indicating whether to checkpoint the learner.
docstr_tok: ["Initializes","the","learner",".","Args",":","policy_network",":","the","online","(","optimized",")","policy",".","critic_network",":","the","online","critic",".","target_policy_network",":","the","target","policy","(","which","lags","behind","the","online","policy",")",".","target_critic_network",":","the","target","critic",".","discount",":","discount","to","use","for","TD","updates",".","target_update_period",":","number","of","learner","steps","to","perform","before","updating","the","target","networks",".","dataset_iterator",":","dataset","to","learn","from",",","whether","fixed","or","from","a","replay","buffer","(","see","`","acme.datasets.reverb.make_reverb_dataset","`","documentation",")",".","replicator",":","Replicates","variables","and","their","update","methods","over","multiple","accelerators",",","such","as","the","multiple","chips","in","a","TPU",".","observation_network",":","an","optional","online","network","to","process","observations","before","the","policy","and","the","critic",".","target_observation_network",":","the","target","observation","network",".","policy_optimizer",":","the","optimizer","to","be","applied","to","the","DPG","(","policy",")","loss",".","critic_optimizer",":","the","optimizer","to","be","applied","to","the","distributional","Bellman","loss",".","clipping",":","whether","to","clip","gradients","by","global","norm",".","counter",":","counter","object","used","to","keep","track","of","steps",".","logger",":","logger","object","to","be","used","by","learner",".","checkpoint",":","boolean","indicating","whether","to","checkpoint","the","learner","."]
code: def __init__(
      self,
      policy_network: snt.Module,
      critic_network: snt.Module,
      target_policy_network: snt.Module,
      target_critic_network: snt.Module,
      discount: float,
      target_update_period: int,
      dataset_iterator: Iterator[reverb.ReplaySample],
      replicator: Optional[Replicator] = None,
      observation_network: types.TensorTransformation = lambda x: x,
      target_observation_network: types.TensorTransformation = lambda x: x,
      policy_optimizer: Optional[snt.Optimizer] = None,
      critic_optimizer: Optional[snt.Optimizer] = None,
      clipping: bool = True,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      checkpoint: bool = True,
  ):
      """Initializes the learner.
  
      Args:
        policy_network: the online (optimized) policy.
        critic_network: the online critic.
        target_policy_network: the target policy (which lags behind the online
          policy).
        target_critic_network: the target critic.
        discount: discount to use for TD updates.
        target_update_period: number of learner steps to perform before updating
          the target networks.
        dataset_iterator: dataset to learn from, whether fixed or from a replay
          buffer (see `acme.datasets.reverb.make_reverb_dataset` documentation).
        replicator: Replicates variables and their update methods over multiple
          accelerators, such as the multiple chips in a TPU.
        observation_network: an optional online network to process observations
          before the policy and the critic.
        target_observation_network: the target observation network.
        policy_optimizer: the optimizer to be applied to the DPG (policy) loss.
        critic_optimizer: the optimizer to be applied to the distributional
          Bellman loss.
        clipping: whether to clip gradients by global norm.
        counter: counter object used to keep track of steps.
        logger: logger object to be used by learner.
        checkpoint: boolean indicating whether to checkpoint the learner.
      """
  
      # 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._clipping = clipping
  
      # Replicates Variables across multiple accelerators
      if not replicator:
          accelerator = _get_first_available_accelerator_type()
          if accelerator == "TPU":
              replicator = snt.distribute.TpuReplicator()
          else:
              replicator = snt.distribute.Replicator()
  
      self._replicator = replicator
  
      with replicator.scope():
          # Necessary to track when to update target networks.
          self._num_steps = tf.Variable(0, dtype=tf.int32)
          self._target_update_period = target_update_period
  
          # Create optimizers if they aren't given.
          self._critic_optimizer = (
              critic_optimizer or snt.optimizers.Adam(1e-4)
          )
          self._policy_optimizer = (
              policy_optimizer or snt.optimizers.Adam(1e-4)
          )
  
      # Batch dataset and create iterator.
      self._iterator = dataset_iterator
  
      # 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 objects.
      self._checkpointer = None
      self._snapshotter = None
  
      if checkpoint:
          self._checkpointer = tf2_savers.Checkpointer(
              subdirectory="d4pg_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,
                  "num_steps": self._num_steps,
              },
          )
          critic_mean = snt.Sequential(
              [self._critic_network, acme_nets.StochasticMeanHead()]
          )
          self._snapshotter = tf2_savers.Snapshotter(
              objects_to_save={
                  "policy": self._policy_network,
                  "critic": critic_mean,
              }
          )
  
      # 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 time
  from typing import Dict, Iterator, List, Optional, Union, 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 reverb
  import sonnet
  import tensorflow
  import tree
label_desc: Use image node_id 1 to create a new D4PGLearner object from inherited base classes: acme with example: obj = D4PGLearner(policy_network, critic_network, target_policy_network, target_critic_network, discount, target_update_period, dataset_iterator, replicator, observation_network, target_observation_network, policy_optimizer, critic_optimizer, clipping, counter, logger, checkpoint)