class_docstr: SVG0 learner with optional prior.
  
  This is the learning component of an SVG0 agent. IE it takes a dataset as
  input and implements update functionality to learn from this dataset.
class_docstr_tok: ["SVG0","learner","with","optional","prior",".","This","is","the","learning","component","of","an","SVG0","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).
    prior_network: the online (optimized) prior.
    target_prior_network: the target prior (which lags behind the online
      prior).
    policy_optimizer: the optimizer to be applied to the SVG-0 (policy) loss.
    critic_optimizer: the optimizer to be applied to the distributional
      Bellman loss.
    prior_optimizer: the optimizer to be applied to the prior (distillation)
      loss.
    distillation_cost: a multiplier to be used when adding distillation
      against the prior to the losses.
    entropy_regularizer_cost: a multiplier used for per state sample based
      entropy added to the actor loss.
    num_action_samples: the number of action samples to use for estimating the
      value function and sample based entropy.
    lambda_: the `lambda` value to be used with retrace.
    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",")",".","prior_network",":","the","online","(","optimized",")","prior",".","target_prior_network",":","the","target","prior","(","which","lags","behind","the","online","prior",")",".","policy_optimizer",":","the","optimizer","to","be","applied","to","the","SVG-0","(","policy",")","loss",".","critic_optimizer",":","the","optimizer","to","be","applied","to","the","distributional","Bellman","loss",".","prior_optimizer",":","the","optimizer","to","be","applied","to","the","prior","(","distillation",")","loss",".","distillation_cost",":","a","multiplier","to","be","used","when","adding","distillation","against","the","prior","to","the","losses",".","entropy_regularizer_cost",":","a","multiplier","used","for","per","state","sample","based","entropy","added","to","the","actor","loss",".","num_action_samples",":","the","number","of","action","samples","to","use","for","estimating","the","value","function","and","sample","based","entropy",".","lambda_",":","the","`","lambda","`","value","to","be","used","with","retrace",".","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],
      prior_network: Optional[snt.Module] = None,
      target_prior_network: Optional[snt.Module] = None,
      policy_optimizer: Optional[snt.Optimizer] = None,
      critic_optimizer: Optional[snt.Optimizer] = None,
      prior_optimizer: Optional[snt.Optimizer] = None,
      distillation_cost: Optional[float] = 1e-3,
      entropy_regularizer_cost: Optional[float] = 1e-3,
      num_action_samples: int = 10,
      lambda_: float = 1.0,
      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).
        prior_network: the online (optimized) prior.
        target_prior_network: the target prior (which lags behind the online
          prior).
        policy_optimizer: the optimizer to be applied to the SVG-0 (policy) loss.
        critic_optimizer: the optimizer to be applied to the distributional
          Bellman loss.
        prior_optimizer: the optimizer to be applied to the prior (distillation)
          loss.
        distillation_cost: a multiplier to be used when adding distillation
          against the prior to the losses.
        entropy_regularizer_cost: a multiplier used for per state sample based
          entropy added to the actor loss.
        num_action_samples: the number of action samples to use for estimating the
          value function and sample based entropy.
        lambda_: the `lambda` value to be used with retrace.
        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
  
      self._prior_network = prior_network
      self._target_prior_network = target_prior_network
  
      self._lambda = lambda_
      self._num_action_samples = num_action_samples
      self._distillation_cost = distillation_cost
      self._entropy_regularizer_cost = entropy_regularizer_cost
  
      # 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
  
      # Necessary to track when to update target networks.
      self._num_steps = tf.Variable(0, dtype=tf.int32)
      self._target_update_period = target_update_period
  
      # Batch dataset and create iterator.
      self._iterator = dataset_iterator
  
      # 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
      )
      self._prior_optimizer = prior_optimizer or snt.optimizers.Adam(
          1e-4
      )
  
      # Expose the variables.
      self._variables = {
          "critic": self._critic_network.variables,
          "policy": self._policy_network.variables,
      }
      if self._prior_network is not None:
          self._variables["prior"] = self._prior_network.variables
  
      # Create a checkpointer and snapshotter objects.
      self._checkpointer = None
      self._snapshotter = None
  
      if checkpoint:
          objects_to_save = {
              "counter": self._counter,
              "policy": self._policy_network,
              "critic": self._critic_network,
              "target_policy": self._target_policy_network,
              "target_critic": self._target_critic_network,
              "policy_optimizer": self._policy_optimizer,
              "critic_optimizer": self._critic_optimizer,
              "num_steps": self._num_steps,
          }
          if self._prior_network is not None:
              objects_to_save["prior"] = self._prior_network
              objects_to_save[
                  "target_prior"
              ] = self._target_prior_network
              objects_to_save["prior_optimizer"] = self._prior_optimizer
  
          self._checkpointer = tf2_savers.Checkpointer(
              subdirectory="svg0_learner",
              objects_to_save=objects_to_save,
          )
          objects_to_snapshot = {
              "policy": self._policy_network,
              "critic": self._critic_network,
          }
          if self._prior_network is not None:
              objects_to_snapshot["prior"] = self._prior_network
  
          self._snapshotter = tf2_savers.Snapshotter(
              objects_to_save=objects_to_snapshot
          )
  
      # 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
  
file_path: acme/acme/agents/tf/svg0_prior/learning.py
imports: import time
  from typing import Dict, Iterator, List, Optional
  import acme
  from acme.agents.tf.svg0_prior import utils
  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
  from trfl import continuous_retrace_ops
label_desc: Use image node_id 1 to create a new SVG0Learner object from inherited base classes: acme with example: obj = SVG0Learner(policy_network, critic_network, target_policy_network, target_critic_network, discount, target_update_period, dataset_iterator, prior_network, target_prior_network, policy_optimizer, critic_optimizer, prior_optimizer, distillation_cost, entropy_regularizer_cost, num_action_samples, lambda_, counter, logger, checkpoint)