class_docstr: A JAX implementation of the SAC agent.
class_docstr_tok: ["A","JAX","implementation","of","the","SAC","agent","."]
code: def _build_networks_and_optimizer(self):
      self._rng, init_key = jax.random.split(self._rng)
  
      # We can reuse init_key safely for the action selection key
      # since it is only used for shape inference during initialization.
      self.network_params = self.network_def.init(
          init_key, self.state, init_key
      )
      self.network_optimizer = dqn_agent.create_optimizer(
          self._optimizer_name
      )
      self.optimizer_state = self.network_optimizer.init(
          self.network_params
      )
  
      # TODO(joshgreaves): Find a way to just copy the critic params
      self.target_params = self.network_params
  
      # \alpha network
      self.log_alpha = jnp.zeros(1)
      self.alpha_optimizer = dqn_agent.create_optimizer(
          self._optimizer_name
      )
      self.alpha_optimizer_state = self.alpha_optimizer.init(
          self.log_alpha
      )
  
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 2 for calling the SACAgent obj's underlying member method code with example usage: obj._build_networks_and_optimizer() without return types