code: def target_q(
      online_network,
      target_network,
      next_states,
      rewards,
      terminals,
      cumulative_gamma,
      double_dqn,
  ):
      """Compute the target Q-value."""
      if double_dqn:
          next_state_q_vals_for_argmax = jax.vmap(
              online_network, in_axes=(0)
          )(next_states).q_values
      else:
          next_state_q_vals_for_argmax = jax.vmap(
              target_network, in_axes=(0)
          )(next_states).q_values
      next_state_q_vals_for_argmax = jnp.squeeze(
          next_state_q_vals_for_argmax
      )
      next_argmax = jnp.argmax(next_state_q_vals_for_argmax, axis=1)
      q_values = jax.vmap(target_network, in_axes=(0))(
          next_states
      ).q_values
      replay_next_qt_max = jax.vmap(lambda t, u: t[u])(
          q_values, next_argmax
      )
      return jax.lax.stop_gradient(
          rewards
          + cumulative_gamma * replay_next_qt_max * (1.0 - terminals)
      )
  
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax.agents.dqn import dqn_agent
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: target_q(online_network, target_network, next_states, rewards, terminals, cumulative_gamma, double_dqn) and returns: jax