code: def target_output(
      model,
      target_network,
      next_states,
      rewards,
      terminals,
      support,
      cumulative_gamma,
      double_dqn,
      distributional,
      rng,
  ):
      """Builds the C51 target distribution or DQN target Q-values."""
  
      is_terminal_multiplier = 1.0 - terminals.astype(jnp.float32)
      # Incorporate terminal state to discount factor.
      gamma_with_terminal = cumulative_gamma * is_terminal_multiplier
  
      target_network_dist = target_network(next_states, key=rng)
      if double_dqn:
          # Use the current network for the action selection
          next_state_target_outputs = model(next_states, key=rng)
      else:
          next_state_target_outputs = target_network_dist
      # Action selection using Q-values for next-state
      q_values = jnp.squeeze(next_state_target_outputs.q_values)
      next_qt_argmax = jnp.argmax(q_values)
  
      if distributional:
          # Compute the target Q-value distribution
          probabilities = jnp.squeeze(target_network_dist.probabilities)
          next_probabilities = probabilities[next_qt_argmax]
          target_support = rewards + gamma_with_terminal * support
          target = rainbow_agent.project_distribution(
              target_support, next_probabilities, support
          )
      else:
          # Compute the target Q-value
          next_q_values = jnp.squeeze(target_network_dist.q_values)
          replay_next_qt_max = next_q_values[next_qt_argmax]
          target = rewards + gamma_with_terminal * replay_next_qt_max
  
      return jax.lax.stop_gradient(target)
  
imports: import functools
  from absl import logging
  from dopamine.jax import losses
  from dopamine.jax import networks
  from dopamine.jax.agents.dqn import dqn_agent
  from dopamine.jax.agents.rainbow import rainbow_agent
  from dopamine.metrics import statistics_instance
  from dopamine.replay_memory import prioritized_replay_buffer
  import gin
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: target_output(model, target_network, next_states, rewards, terminals, support, cumulative_gamma, double_dqn, distributional, rng) and returns: jax