class_docstr: Offline Rainbow agent with BC regularization (akin to CQL).
docstr: Runs a single training step.
docstr_tok: ["Runs","a","single","training","step","."]
code: def _training_step_update(self):
      """Runs a single training step."""
      self._sample_from_replay_buffer()
      states = self.preprocess_fn(self.replay_elements["state"])
      next_states = self.preprocess_fn(
          self.replay_elements["next_state"]
      )
  
      (
          self.optimizer_state,
          self.online_params,
          mean_loss,
          aux_info,
          self._rng,
      ) = train(
          self.network_def,
          self.online_params,
          self.target_network_params,
          self.optimizer,
          self.optimizer_state,
          states,
          self.replay_elements["action"],
          next_states,
          self.replay_elements["reward"],
          self.replay_elements["terminal"],
          self._support,
          self.cumulative_gamma,
          self._double_dqn,
          self._distributional,
          self._rng,
          bc_coefficient=self._bc_coefficient,
          td_coefficient=self._td_coefficient,
      )
  
      td_loss, bc_loss = aux_info
      if (
          self.training_steps > 0
          and self.training_steps % self.summary_writing_frequency == 0
      ):
          if self.summary_writer is not None:
              with self.summary_writer.as_default():
                  tf.summary.scalar(
                      "Losses/Aggregate",
                      mean_loss,
                      step=self.training_steps,
                  )
                  tf.summary.scalar(
                      "Losses/TD", td_loss, step=self.training_steps
                  )
                  tf.summary.scalar(
                      "Losses/BCLoss", bc_loss, step=self.training_steps
                  )
              self.summary_writer.flush()
      if self._use_tfds:
          self.log_gradient_steps_per_epoch()
  
      if self.training_steps % self.target_update_period == 0:
          self._sync_weights()
      self.training_steps += 1
  
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax.agents.full_rainbow import full_rainbow_agent
  from dopamine.jax.agents.rainbow import rainbow_agent
  from dopamine.labs.offline_rl import fixed_replay
  from dopamine.labs.offline_rl.rlu_tfds import tfds_replay
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 2 for calling the OfflineJaxRainbowAgent obj's underlying member method code with example usage: obj._training_step_update() without return types