class_docstr: Recurrent CRR learner.
  
  This is the learning component of a RCRR agent. It takes a dataset as
  input and implements update functionality to learn from this dataset.
class_docstr_tok: ["Recurrent","CRR","learner",".","This","is","the","learning","component","of","a","RCRR","agent",".","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.
    dataset: dataset to learn from, whether fixed or from a replay buffer
      (see `acme.datasets.reverb.make_reverb_dataset` documentation).
    accelerator_strategy: the strategy used to distribute computation,
      whether on a single, or multiple, GPU or TPU; as supported by
      tf.distribute.
    behavior_network: The network to snapshot under `policy` name. If None,
      snapshots `policy_network` instead.
    cwp_network: CWP network to snapshot: samples actions
      from the policy and weighs them with the critic, then returns the action
      by sampling from the softmax distribution using critic values as logits.
      Used only for snapshotting, not training.
    policy_optimizer: the optimizer to be applied to the policy loss.
    critic_optimizer: the optimizer to be applied to the distributional
      Bellman loss.
    discount: discount to use for TD updates.
    target_update_period: number of learner steps to perform before updating
      the target networks.
    num_action_samples_td_learning: number of action samples to use to
      estimate expected value of the critic loss w.r.t. stochastic policy.
    num_action_samples_policy_weight: number of action samples to use to
      estimate the advantage function for the CRR weighting of the policy
      loss.
    baseline_reduce_function: one of 'mean', 'max', 'min'. Way of aggregating
      values from `num_action_samples` estimates of the value function.
    clipping: whether to clip gradients by global norm.
    policy_improvement_modes: one of 'exp', 'binary', 'all'. CRR mode which
      determines how the advantage function is processed before being
      multiplied by the policy loss.
    ratio_upper_bound: if policy_improvement_modes is 'exp', determines
      the upper bound of the weight (i.e. the weight is
        min(exp(advantage / beta), upper_bound)
      ).
    beta: if policy_improvement_modes is 'exp', determines the beta (see
      above).
    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",".","dataset",":","dataset","to","learn","from",",","whether","fixed","or","from","a","replay","buffer","(","see","`","acme.datasets.reverb.make_reverb_dataset","`","documentation",")",".","accelerator_strategy",":","the","strategy","used","to","distribute","computation",",","whether","on","a","single",",","or","multiple",",","GPU","or","TPU",";","as","supported","by","tf.distribute",".","behavior_network",":","The","network","to","snapshot","under","`","policy","`","name",".","If","None",",","snapshots","`","policy_network","`","instead",".","cwp_network",":","CWP","network","to","snapshot",":","samples","actions","from","the","policy","and","weighs","them","with","the","critic",",","then","returns","the","action","by","sampling","from","the","softmax","distribution","using","critic","values","as","logits",".","Used","only","for","snapshotting",",","not","training",".","policy_optimizer",":","the","optimizer","to","be","applied","to","the","policy","loss",".","critic_optimizer",":","the","optimizer","to","be","applied","to","the","distributional","Bellman","loss",".","discount",":","discount","to","use","for","TD","updates",".","target_update_period",":","number","of","learner","steps","to","perform","before","updating","the","target","networks",".","num_action_samples_td_learning",":","number","of","action","samples","to","use","to","estimate","expected","value","of","the","critic","loss","w.r.t",".","stochastic","policy",".","num_action_samples_policy_weight",":","number","of","action","samples","to","use","to","estimate","the","advantage","function","for","the","CRR","weighting","of","the","policy","loss",".","baseline_reduce_function",":","one","of","'mean","'",",","'max","'",",","'min","'",".","Way","of","aggregating","values","from","`","num_action_samples","`","estimates","of","the","value","function",".","clipping",":","whether","to","clip","gradients","by","global","norm",".","policy_improvement_modes",":","one","of","'exp","'",",","'binary","'",",","'all","'",".","CRR","mode","which","determines","how","the","advantage","function","is","processed","before","being","multiplied","by","the","policy","loss",".","ratio_upper_bound",":","if","policy_improvement_modes","is","'exp","'",",","determines","the","upper","bound","of","the","weight","(","i.e",".","the","weight","is","min","(","exp","(","advantage","\/","beta",")",",","upper_bound",")",")",".","beta",":","if","policy_improvement_modes","is","'exp","'",",","determines","the","beta","(","see","above",")",".","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.RNNCore,
      critic_network: networks.CriticDeepRNN,
      target_policy_network: snt.RNNCore,
      target_critic_network: networks.CriticDeepRNN,
      dataset: tf.data.Dataset,
      accelerator_strategy: Optional[tf.distribute.Strategy] = None,
      behavior_network: Optional[snt.Module] = None,
      cwp_network: Optional[snt.Module] = None,
      policy_optimizer: Optional[snt.Optimizer] = None,
      critic_optimizer: Optional[snt.Optimizer] = None,
      discount: float = 0.99,
      target_update_period: int = 100,
      num_action_samples_td_learning: int = 1,
      num_action_samples_policy_weight: int = 4,
      baseline_reduce_function: str = "mean",
      clipping: bool = True,
      policy_improvement_modes: str = "exp",
      ratio_upper_bound: float = 20.0,
      beta: float = 1.0,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      checkpoint: bool = False,
  ):
      """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.
        dataset: dataset to learn from, whether fixed or from a replay buffer
          (see `acme.datasets.reverb.make_reverb_dataset` documentation).
        accelerator_strategy: the strategy used to distribute computation,
          whether on a single, or multiple, GPU or TPU; as supported by
          tf.distribute.
        behavior_network: The network to snapshot under `policy` name. If None,
          snapshots `policy_network` instead.
        cwp_network: CWP network to snapshot: samples actions
          from the policy and weighs them with the critic, then returns the action
          by sampling from the softmax distribution using critic values as logits.
          Used only for snapshotting, not training.
        policy_optimizer: the optimizer to be applied to the policy loss.
        critic_optimizer: the optimizer to be applied to the distributional
          Bellman loss.
        discount: discount to use for TD updates.
        target_update_period: number of learner steps to perform before updating
          the target networks.
        num_action_samples_td_learning: number of action samples to use to
          estimate expected value of the critic loss w.r.t. stochastic policy.
        num_action_samples_policy_weight: number of action samples to use to
          estimate the advantage function for the CRR weighting of the policy
          loss.
        baseline_reduce_function: one of 'mean', 'max', 'min'. Way of aggregating
          values from `num_action_samples` estimates of the value function.
        clipping: whether to clip gradients by global norm.
        policy_improvement_modes: one of 'exp', 'binary', 'all'. CRR mode which
          determines how the advantage function is processed before being
          multiplied by the policy loss.
        ratio_upper_bound: if policy_improvement_modes is 'exp', determines
          the upper bound of the weight (i.e. the weight is
            min(exp(advantage / beta), upper_bound)
          ).
        beta: if policy_improvement_modes is 'exp', determines the beta (see
          above).
        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.
      """
  
      if accelerator_strategy is None:
          accelerator_strategy = snt.distribute.Replicator()
      self._accelerator_strategy = accelerator_strategy
      self._policy_improvement_modes = policy_improvement_modes
      self._ratio_upper_bound = ratio_upper_bound
      self._num_action_samples_td_learning = (
          num_action_samples_td_learning
      )
      self._num_action_samples_policy_weight = (
          num_action_samples_policy_weight
      )
      self._baseline_reduce_function = baseline_reduce_function
      self._beta = beta
  
      # When running on TPUs we have to know the amount of memory required (and
      # thus the sequence length) at the graph compilation stage. At the moment,
      # the only way to get it is to sample from the dataset, since the dataset
      # does not have any metadata, see b/160672927 to track this upcoming
      # feature.
      sample = next(dataset.as_numpy_iterator())
      self._sequence_length = sample.action.shape[1]
  
      self._counter = counter or counting.Counter()
      self._logger = logger or loggers.TerminalLogger(
          "learner", time_delta=1.0
      )
      self._discount = discount
      self._clipping = clipping
  
      self._target_update_period = target_update_period
  
      with self._accelerator_strategy.scope():
          # Necessary to track when to update target networks.
          self._num_steps = tf.Variable(0, dtype=tf.int32)
  
          # (Maybe) distributing the dataset across multiple accelerators.
          distributed_dataset = self._accelerator_strategy.experimental_distribute_dataset(
              dataset
          )
          self._iterator = iter(distributed_dataset)
  
          # Create the optimizers.
          self._critic_optimizer = (
              critic_optimizer or snt.optimizers.Adam(1e-4)
          )
          self._policy_optimizer = (
              policy_optimizer or snt.optimizers.Adam(1e-4)
          )
  
      # 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
  
      # Expose the variables.
      self._variables = {
          "critic": self._target_critic_network.variables,
          "policy": self._target_policy_network.variables,
      }
  
      # Create a checkpointer object.
      self._checkpointer = None
      self._snapshotter = None
      if checkpoint:
          self._checkpointer = tf2_savers.Checkpointer(
              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,
              },
              time_delta_minutes=30.0,
          )
  
          raw_policy = snt.DeepRNN(
              [policy_network, networks.StochasticSamplingHead()]
          )
          critic_mean = networks.CriticDeepRNN(
              [critic_network, networks.StochasticMeanHead()]
          )
          objects_to_save = {
              "raw_policy": raw_policy,
              "critic": critic_mean,
          }
          if behavior_network is not None:
              objects_to_save["policy"] = behavior_network
          if cwp_network is not None:
              objects_to_save["cwp_policy"] = cwp_network
          self._snapshotter = tf2_savers.Snapshotter(
              objects_to_save=objects_to_save, time_delta_minutes=30
          )
      # Timestamp to keep track of the wall time.
      self._walltime_timestamp = time.time()
  
imports: import operator
  import time
  from typing import Dict, List, Optional
  from acme import core
  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 RCRRLearner object from inherited base classes: core with example: obj = RCRRLearner(policy_network, critic_network, target_policy_network, target_critic_network, dataset, accelerator_strategy, behavior_network, cwp_network, policy_optimizer, critic_optimizer, discount, target_update_period, num_action_samples_td_learning, num_action_samples_policy_weight, baseline_reduce_function, clipping, policy_improvement_modes, ratio_upper_bound, beta, counter, logger, checkpoint)