class_docstr: A JAX implementation of the SAC agent.
class_docstr_tok: ["A","JAX","implementation","of","the","SAC","agent","."]
docstr: Initializes the agent and constructs the necessary components.
  
  Args:
    action_shape: int or tuple, dimensionality of the action space.
    action_limits: pair of lower and higher bounds for actions.
    observation_shape: tuple of ints describing the observation shape.
    action_dtype: jnp.dtype, specifies the type of the actions.
    observation_dtype: jnp.dtype, specifies the type of the observations.
    reward_scale_factor: float, factor by which to scale rewards.
    stack_size: int, number of frames to use in state stack.
    network: Jax network to use for training.
    num_layers: int, number of layers in the network.
    hidden_units: int, number of hidden units in the network.
    gamma: float, discount factor with the usual RL meaning.
    update_horizon: int, horizon at which updates are performed, the 'n' in
      n-step update.
    min_replay_history: int, number of transitions that should be experienced
      before the agent begins training its value function.
    update_period: int, period between DQN updates.
    target_update_type: str, if 'hard', will perform a hard update of the
      target network every target_update_period training steps; if 'soft',
      will use target_smoothing_coefficient to update the target network at
      every training step.
    target_update_period: int, frequency with which to update target network
      when in 'hard' mode.
    target_smoothing_coefficient: float, smoothing coefficient for target
      network updates (\tau in paper) when in 'soft' mode.
    target_entropy: float or None, the target entropy for training alpha. If
      None, it will default to the half the negative of the number of action
      dimensions.
    eval_mode: bool, True for evaluation and False for training.
    optimizer: str, name of optimizer to use.
    summary_writer: SummaryWriter object for outputting training statistics.
    summary_writing_frequency: int, frequency with which summaries will be
      written. Lower values will result in slower training.
    allow_partial_reload: bool, whether we allow reloading a partial agent
      (for instance, only the network parameters).
    seed: int, a seed for SAC's internal RNG, used for initialization and
      sampling actions.
    collector_allowlist: list of str, if using CollectorDispatcher, this can
      be used to specify which Collectors to log to.
docstr_tok: ["Initializes","the","agent","and","constructs","the","necessary","components",".","Args",":","action_shape",":","int","or","tuple",",","dimensionality","of","the","action","space",".","action_limits",":","pair","of","lower","and","higher","bounds","for","actions",".","observation_shape",":","tuple","of","ints","describing","the","observation","shape",".","action_dtype",":","jnp.dtype",",","specifies","the","type","of","the","actions",".","observation_dtype",":","jnp.dtype",",","specifies","the","type","of","the","observations",".","reward_scale_factor",":","float",",","factor","by","which","to","scale","rewards",".","stack_size",":","int",",","number","of","frames","to","use","in","state","stack",".","network",":","Jax","network","to","use","for","training",".","num_layers",":","int",",","number","of","layers","in","the","network",".","hidden_units",":","int",",","number","of","hidden","units","in","the","network",".","gamma",":","float",",","discount","factor","with","the","usual","RL","meaning",".","update_horizon",":","int",",","horizon","at","which","updates","are","performed",",","the","'n","'","in","n-step","update",".","min_replay_history",":","int",",","number","of","transitions","that","should","be","experienced","before","the","agent","begins","training","its","value","function",".","update_period",":","int",",","period","between","DQN","updates",".","target_update_type",":","str",",","if","'hard","'",",","will","perform","a","hard","update","of","the","target","network","every","target_update_period","training","steps",";","if","'soft","'",",","will","use","target_smoothing_coefficient","to","update","the","target","network","at","every","training","step",".","target_update_period",":","int",",","frequency","with","which","to","update","target","network","when","in","'hard","'","mode",".","target_smoothing_coefficient",":","float",",","smoothing","coefficient","for","target","network","updates","(","\\tau","in","paper",")","when","in","'soft","'","mode",".","target_entropy",":","float","or","None",",","the","target","entropy","for","training","alpha",".","If","None",",","it","will","default","to","the","half","the","negative","of","the","number","of","action","dimensions",".","eval_mode",":","bool",",","True","for","evaluation","and","False","for","training",".","optimizer",":","str",",","name","of","optimizer","to","use",".","summary_writer",":","SummaryWriter","object","for","outputting","training","statistics",".","summary_writing_frequency",":","int",",","frequency","with","which","summaries","will","be","written",".","Lower","values","will","result","in","slower","training",".","allow_partial_reload",":","bool",",","whether","we","allow","reloading","a","partial","agent","(","for","instance",",","only","the","network","parameters",")",".","seed",":","int",",","a","seed","for","SAC","'s","internal","RNG",",","used","for","initialization","and","sampling","actions",".","collector_allowlist",":","list","of","str",",","if","using","CollectorDispatcher",",","this","can","be","used","to","specify","which","Collectors","to","log","to","."]
code: def __init__(
      self,
      action_shape,
      action_limits,
      observation_shape,
      action_dtype=jnp.float32,
      observation_dtype=jnp.float32,
      reward_scale_factor=1.0,
      stack_size=1,
      network=continuous_networks.SACNetwork,
      num_layers=2,
      hidden_units=256,
      gamma=0.99,
      update_horizon=1,
      min_replay_history=20000,
      update_period=1,
      target_update_type="soft",
      target_update_period=1000,
      target_smoothing_coefficient=0.005,
      target_entropy=None,
      eval_mode=False,
      optimizer="adam",
      summary_writer=None,
      summary_writing_frequency=500,
      allow_partial_reload=False,
      seed=None,
      collector_allowlist=("tensorboard"),
  ):
      r"""Initializes the agent and constructs the necessary components.
  
      Args:
        action_shape: int or tuple, dimensionality of the action space.
        action_limits: pair of lower and higher bounds for actions.
        observation_shape: tuple of ints describing the observation shape.
        action_dtype: jnp.dtype, specifies the type of the actions.
        observation_dtype: jnp.dtype, specifies the type of the observations.
        reward_scale_factor: float, factor by which to scale rewards.
        stack_size: int, number of frames to use in state stack.
        network: Jax network to use for training.
        num_layers: int, number of layers in the network.
        hidden_units: int, number of hidden units in the network.
        gamma: float, discount factor with the usual RL meaning.
        update_horizon: int, horizon at which updates are performed, the 'n' in
          n-step update.
        min_replay_history: int, number of transitions that should be experienced
          before the agent begins training its value function.
        update_period: int, period between DQN updates.
        target_update_type: str, if 'hard', will perform a hard update of the
          target network every target_update_period training steps; if 'soft',
          will use target_smoothing_coefficient to update the target network at
          every training step.
        target_update_period: int, frequency with which to update target network
          when in 'hard' mode.
        target_smoothing_coefficient: float, smoothing coefficient for target
          network updates (\tau in paper) when in 'soft' mode.
        target_entropy: float or None, the target entropy for training alpha. If
          None, it will default to the half the negative of the number of action
          dimensions.
        eval_mode: bool, True for evaluation and False for training.
        optimizer: str, name of optimizer to use.
        summary_writer: SummaryWriter object for outputting training statistics.
        summary_writing_frequency: int, frequency with which summaries will be
          written. Lower values will result in slower training.
        allow_partial_reload: bool, whether we allow reloading a partial agent
          (for instance, only the network parameters).
        seed: int, a seed for SAC's internal RNG, used for initialization and
          sampling actions.
        collector_allowlist: list of str, if using CollectorDispatcher, this can
          be used to specify which Collectors to log to.
      """
      assert isinstance(observation_shape, tuple)
      # If we're performing hard updates, we force the smoothing coefficient to 1.
      if target_update_type == "hard":
          target_smoothing_coefficient = 1.0
  
      if isinstance(action_shape, int):
          action_shape = (action_shape,)
  
      # If target_entropy is None, set to default value.
      if target_entropy is None:
          action_dim = functools.reduce(operator.mul, action_shape, 1.0)
          target_entropy = -0.5 * action_dim
      seed = int(time.time() * 1e6) if seed is None else seed
      logging.info(
          "Creating %s agent with the following parameters:",
          self.__class__.__name__,
      )
      logging.info("\t action_shape: %s", action_shape)
      logging.info("\t action_dtype: %s", action_dtype)
      logging.info("\t action_limits: %s", action_limits)
      logging.info("\t observation_shape: %s", observation_shape)
      logging.info("\t observation_dtype: %s", observation_dtype)
      logging.info("\t reward_scale_factor: %f", reward_scale_factor)
      logging.info("\t num_layers: %d", num_layers)
      logging.info("\t hidden_units: %d", hidden_units)
      logging.info("\t gamma: %f", gamma)
      logging.info("\t update_horizon: %f", update_horizon)
      logging.info("\t min_replay_history: %d", min_replay_history)
      logging.info("\t update_period: %d", update_period)
      logging.info("\t target_update_type: %s", target_update_type)
      logging.info("\t target_update_period: %d", target_update_period)
      logging.info(
          "\t target_smoothing_coefficient: %f",
          target_smoothing_coefficient,
      )
      logging.info("\t target_entropy: %f", target_entropy)
      logging.info("\t optimizer: %s", optimizer)
      logging.info("\t seed: %d", seed)
  
      self.action_shape = action_shape
      self.action_dtype = action_dtype
      self.observation_shape = tuple(observation_shape)
      self.observation_dtype = observation_dtype
      self.reward_scale_factor = reward_scale_factor
      self.stack_size = stack_size
      self.action_limits = action_limits
      action_limits = tuple(tuple(x.reshape(-1)) for x in action_limits)
      self.network_def = network(
          action_shape, num_layers, hidden_units, action_limits
      )
      self.gamma = gamma
      self.update_horizon = update_horizon
      self.cumulative_gamma = math.pow(gamma, update_horizon)
      self.min_replay_history = min_replay_history
      self.update_period = update_period
      self.target_update_type = target_update_type
      self.target_update_period = target_update_period
      self.target_smoothing_coefficient = target_smoothing_coefficient
      self.target_entropy = target_entropy
      self.eval_mode = eval_mode
      self.training_steps = 0
      self.summary_writer = summary_writer
      self.summary_writing_frequency = summary_writing_frequency
      self.allow_partial_reload = allow_partial_reload
      self._collector_allowlist = collector_allowlist
  
      self._rng = jax.random.PRNGKey(seed)
      state_shape = self.observation_shape + (stack_size,)
      self.state = onp.zeros(state_shape)
      self._replay = self._build_replay_buffer()
      self._optimizer_name = optimizer
      self._build_networks_and_optimizer()
  
      # Variables to be initialized by the agent once it interacts with the
      # environment.
      self._observation = None
      self._last_observation = None
  
imports: import functools
  import math
  import operator
  import time
  from typing import Any, Mapping, Tuple
  from absl import logging
  from dopamine.jax import continuous_networks
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  import dopamine.labs.sac_from_pixels.continuous_networks
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import circular_replay_buffer
  import flax
  from flax import linen
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 1 to create a new SACAgent object from inherited base classes: dqn_agent with example: obj = SACAgent(action_shape, action_limits, observation_shape, action_dtype, observation_dtype, reward_scale_factor, stack_size, network, num_layers, hidden_units, gamma, update_horizon, min_replay_history, update_period, target_update_type, target_update_period, target_smoothing_coefficient, target_entropy, eval_mode, optimizer, summary_writer, summary_writing_frequency, allow_partial_reload, seed, collector_allowlist)