class_docstr: A JAX implementation of the SAC agent.
class_docstr_tok: ["A","JAX","implementation","of","the","SAC","agent","."]
docstr: Runs a single training step.
  
  Runs training if both:
    (1) A minimum number of frames have been added to the replay buffer.
    (2) `training_steps` is a multiple of `update_period`.
  
  Also, syncs weights from online_network to target_network if training steps
  is a multiple of target update period.
docstr_tok: ["Runs","a","single","training","step",".","Runs","training","if","both",":","(","1",")","A","minimum","number","of","frames","have","been","added","to","the","replay","buffer",".","(","2",")","`","training_steps","`","is","a","multiple","of","`","update_period","`",".","Also",",","syncs","weights","from","online_network","to","target_network","if","training","steps","is","a","multiple","of","target","update","period","."]
code: def _train_step(self):
      """Runs a single training step.
  
      Runs training if both:
        (1) A minimum number of frames have been added to the replay buffer.
        (2) `training_steps` is a multiple of `update_period`.
  
      Also, syncs weights from online_network to target_network if training steps
      is a multiple of target update period.
      """
      if self._replay.add_count > self.min_replay_history:
          if self.training_steps % self.update_period == 0:
              self._sample_from_replay_buffer()
              self._rng, key = jax.random.split(self._rng)
  
              train_returns = train(
                  self.network_def,
                  self.network_optimizer,
                  self.alpha_optimizer,
                  self.optimizer_state,
                  self.alpha_optimizer_state,
                  self.network_params,
                  self.target_params,
                  self.log_alpha,
                  key,
                  self.replay_elements["state"],
                  self.replay_elements["action"],
                  self.replay_elements["next_state"],
                  self.replay_elements["reward"],
                  self.replay_elements["terminal"],
                  self.cumulative_gamma,
                  self.target_entropy,
                  self.reward_scale_factor,
              )
  
              self.network_params = train_returns["network_params"]
              self.optimizer_state = train_returns["optimizer_state"]
              self.log_alpha = train_returns["log_alpha"]
              self.alpha_optimizer_state = train_returns[
                  "alpha_optimizer_state"
              ]
  
              if (
                  self.summary_writer is not None
                  and self.training_steps > 0
                  and self.training_steps
                  % self.summary_writing_frequency
                  == 0
              ):
                  statistics = []
                  for k in train_returns:
                      if k.startswith("Losses") or k.startswith(
                          "Values"
                      ):
                          self.summary_writer.scalar(
                              k, train_returns[k], self.training_steps
                          )
                          if hasattr(self, "collector_dispatcher"):
                              statistics.append(
                                  statistics_instance.StatisticsInstance(
                                      k,
                                      onp.asarray(train_returns[k]),
                                      step=self.training_steps,
                                  )
                              )
                  if hasattr(self, "collector_dispatcher"):
                      self.collector_dispatcher.write(
                          statistics,
                          collector_allowlist=self._collector_allowlist,
                      )
                  self.summary_writer.flush()
              self._maybe_sync_weights()
      self.training_steps += 1
  
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 7 for calling the SACAgent obj's underlying member method code with example usage: obj._train_step() without return types