class_docstr: SVG0 Agent with prior.
  
  This implements a single-process SVG0 agent. This is an actor-critic algorithm
  that generates data via a behavior policy, inserts N-step transitions into
  a replay buffer, and periodically updates the policy (and as a result the
  behavior) by sampling uniformly from this buffer.
class_docstr_tok: ["SVG0","Agent","with","prior",".","This","implements","a","single-process","SVG0","agent",".","This","is","an","actor-critic","algorithm","that","generates","data","via","a","behavior","policy",",","inserts","N-step","transitions","into","a","replay","buffer",",","and","periodically","updates","the","policy","(","and","as","a","result","the","behavior",")","by","sampling","uniformly","from","this","buffer","."]
docstr: Initialize the agent.
  
  Args:
    environment_spec: description of the actions, observations, etc.
    policy_network: the online (optimized) policy.
    critic_network: the online critic.
    discount: discount to use for TD updates.
    batch_size: batch size for updates.
    prefetch_size: size to prefetch from replay.
    target_update_period: number of learner steps to perform before updating
      the target networks.
    prior_network: an optional `behavior prior` to regularize against.
    policy_optimizer: optimizer for the policy network updates.
    critic_optimizer: optimizer for the critic network updates.
    prior_optimizer: optimizer for the prior network updates.
    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.
    min_replay_size: minimum replay size before updating.
    max_replay_size: maximum replay size.
    samples_per_insert: number of samples to take from replay for every insert
      that is made.
    sequence_length: number of timesteps to store for each trajectory.
    sigma: standard deviation of zero-mean, Gaussian exploration noise.
    replay_table_name: string indicating what name to give the replay table.
    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: ["Initialize","the","agent",".","Args",":","environment_spec",":","description","of","the","actions",",","observations",",","etc",".","policy_network",":","the","online","(","optimized",")","policy",".","critic_network",":","the","online","critic",".","discount",":","discount","to","use","for","TD","updates",".","batch_size",":","batch","size","for","updates",".","prefetch_size",":","size","to","prefetch","from","replay",".","target_update_period",":","number","of","learner","steps","to","perform","before","updating","the","target","networks",".","prior_network",":","an","optional","`","behavior","prior","`","to","regularize","against",".","policy_optimizer",":","optimizer","for","the","policy","network","updates",".","critic_optimizer",":","optimizer","for","the","critic","network","updates",".","prior_optimizer",":","optimizer","for","the","prior","network","updates",".","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",".","min_replay_size",":","minimum","replay","size","before","updating",".","max_replay_size",":","maximum","replay","size",".","samples_per_insert",":","number","of","samples","to","take","from","replay","for","every","insert","that","is","made",".","sequence_length",":","number","of","timesteps","to","store","for","each","trajectory",".","sigma",":","standard","deviation","of","zero-mean",",","Gaussian","exploration","noise",".","replay_table_name",":","string","indicating","what","name","to","give","the","replay","table",".","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,
      environment_spec: specs.EnvironmentSpec,
      policy_network: snt.Module,
      critic_network: snt.Module,
      discount: float = 0.99,
      batch_size: int = 256,
      prefetch_size: int = 4,
      target_update_period: int = 100,
      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,
      min_replay_size: int = 1000,
      max_replay_size: int = 1000000,
      samples_per_insert: float = 32.0,
      sequence_length: int = 10,
      sigma: float = 0.3,
      replay_table_name: str = reverb_adders.DEFAULT_PRIORITY_TABLE,
      counter: Optional[counting.Counter] = None,
      logger: Optional[loggers.Logger] = None,
      checkpoint: bool = True,
  ):
      """Initialize the agent.
  
      Args:
        environment_spec: description of the actions, observations, etc.
        policy_network: the online (optimized) policy.
        critic_network: the online critic.
        discount: discount to use for TD updates.
        batch_size: batch size for updates.
        prefetch_size: size to prefetch from replay.
        target_update_period: number of learner steps to perform before updating
          the target networks.
        prior_network: an optional `behavior prior` to regularize against.
        policy_optimizer: optimizer for the policy network updates.
        critic_optimizer: optimizer for the critic network updates.
        prior_optimizer: optimizer for the prior network updates.
        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.
        min_replay_size: minimum replay size before updating.
        max_replay_size: maximum replay size.
        samples_per_insert: number of samples to take from replay for every insert
          that is made.
        sequence_length: number of timesteps to store for each trajectory.
        sigma: standard deviation of zero-mean, Gaussian exploration noise.
        replay_table_name: string indicating what name to give the replay table.
        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.
      """
      # Create the Builder object which will internally create agent components.
      builder = SVG0Builder(
          # TODO(mwhoffman): pass the config dataclass in directly.
          # TODO(mwhoffman): use the limiter rather than the workaround below.
          # Right now this modifies min_replay_size and samples_per_insert so that
          # they are not controlled by a limiter and are instead handled by the
          # Agent base class (the above TODO directly references this behavior).
          SVG0Config(
              discount=discount,
              batch_size=batch_size,
              prefetch_size=prefetch_size,
              target_update_period=target_update_period,
              policy_optimizer=policy_optimizer,
              critic_optimizer=critic_optimizer,
              prior_optimizer=prior_optimizer,
              distillation_cost=distillation_cost,
              entropy_regularizer_cost=entropy_regularizer_cost,
              min_replay_size=1,  # Let the Agent class handle this.
              max_replay_size=max_replay_size,
              samples_per_insert=None,  # Let the Agent class handle this.
              sequence_length=sequence_length,
              sigma=sigma,
              replay_table_name=replay_table_name,
          )
      )
  
      # TODO(mwhoffman): pass the network dataclass in directly.
      online_networks = SVG0Networks(
          policy_network=policy_network,
          critic_network=critic_network,
          prior_network=prior_network,
      )
  
      # Target networks are just a copy of the online networks.
      target_networks = copy.deepcopy(online_networks)
  
      # Initialize the networks.
      online_networks.init(environment_spec)
      target_networks.init(environment_spec)
  
      # TODO(mwhoffman): either make this Dataclass or pass only one struct.
      # The network struct passed to make_learner is just a tuple for the
      # time-being (for backwards compatibility).
      networks = (online_networks, target_networks)
  
      # Create the behavior policy.
      policy_network = online_networks.make_policy()
  
      # Create the replay server and grab its address.
      replay_tables = builder.make_replay_tables(
          environment_spec, sequence_length
      )
      replay_server = reverb.Server(replay_tables, port=None)
      replay_client = reverb.Client(f"localhost:{replay_server.port}")
  
      # Create actor, dataset, and learner for generating, storing, and consuming
      # data respectively.
      adder = builder.make_adder(replay_client)
      actor = builder.make_actor(policy_network, adder)
      dataset = builder.make_dataset_iterator(replay_client)
      learner = builder.make_learner(
          networks, dataset, counter, logger, checkpoint
      )
  
      super().__init__(
          actor=actor,
          learner=learner,
          min_observations=max(batch_size, min_replay_size),
          observations_per_step=float(batch_size) / samples_per_insert,
      )
  
      # Save the replay so we don't garbage collect it.
      self._replay_server = replay_server
  
imports: import copy
  import dataclasses
  from typing import Iterator, List, Optional, Tuple
  from acme import adders
  from acme import core
  from acme import datasets
  from acme import specs
  from acme.adders import reverb
  from acme.agents import agent
  from acme.agents.tf.svg0_prior import acting
  from acme.agents.tf.svg0_prior import learning
  from acme.tf import utils
  from acme.tf import variable_utils
  from acme.utils import counting
  from acme.utils import loggers
  import reverb
  import sonnet
  import tensorflow
label_desc: Use image node_id 1 to create a new SVG0 object from inherited base classes: agent with example: obj = SVG0(environment_spec, policy_network, critic_network, discount, batch_size, prefetch_size, target_update_period, prior_network, policy_optimizer, critic_optimizer, prior_optimizer, distillation_cost, entropy_regularizer_cost, min_replay_size, max_replay_size, samples_per_insert, sequence_length, sigma, replay_table_name, counter, logger, checkpoint)