class_docstr: A JAX implementation of the Offline DQN agent with DR3 regularizer.
class_docstr_tok: ["A","JAX","implementation","of","the","Offline","DQN","agent","with","DR3","regularizer","."]
docstr: Runs a single training step.
docstr_tok: ["Runs","a","single","training","step","."]
code: def train_step(self):
      """Runs a single training step."""
      if self.training_steps % self.update_period == 0:
          self._sample_from_replay_buffer()
          (
              self.optimizer_state,
              self.online_params,
              loss,
              bellman_loss,
              dr3_loss,
              cql_loss,
              actual_bellman_loss,
              q_data,
              q_pi,
          ) = train(
              self.network_def,
              self.online_params,
              self.target_network_params,
              self.optimizer,
              self.optimizer_state,
              self.preprocess_fn(self.replay_elements["state"]),
              self.replay_elements["action"],
              self.preprocess_fn(self.replay_elements["next_state"]),
              self.replay_elements["reward"],
              self.replay_elements["terminal"],
              self.cumulative_gamma,
              self._dr3_coefficient,
              self._cql_coefficient,
          )
  
          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",
                          loss,
                          step=self.training_steps,
                      )
                      tf.summary.scalar(
                          "Losses/Bellman",
                          bellman_loss,
                          step=self.training_steps,
                      )
                      tf.summary.scalar(
                          "Losses/DR3",
                          dr3_loss,
                          step=self.training_steps,
                      )
                      tf.summary.scalar(
                          "Losses/CQLLoss",
                          cql_loss,
                          step=self.training_steps,
                      )
                      tf.summary.scalar(
                          "Losses/BellmanActual",
                          actual_bellman_loss,
                          step=self.training_steps,
                      )
                      tf.summary.scalar(
                          "Losses/q_vals_data",
                          q_data,
                          step=self.training_steps,
                      )
                      tf.summary.scalar(
                          "Losses/q_vals_pi",
                          q_pi,
                          step=self.training_steps,
                      )
                  self.summary_writer.flush()
              if hasattr(self, "collector_dispatcher"):
                  self.collector_dispatcher.write(
                      [
                          statistics_instance.StatisticsInstance(
                              "Loss",
                              np.asarray(loss),
                              step=self.training_steps,
                          ),
                      ],
                      collector_allowlist=self._collector_allowlist,
                  )
      if self.training_steps % self.target_update_period == 0:
          self._sync_weights()
      self.training_steps += 1
  
file_path: dopamine/dopamine/labs/offline_rl/jax/offline_dr3_agent.py
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.labs.offline_rl.jax import offline_dqn_agent
  from dopamine.metrics import statistics_instance
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 2 for calling the OfflineJaxDR3Agent obj's underlying member method code with example usage: obj.train_step() without return types