class_docstr: An online Q-learning deep RL agent with PopArt.
class_docstr_tok: ["An","online","Q-learning","deep","RL","agent","with","PopArt","."]
docstr: Loss function.
docstr_tok: ["Loss","function","."]
code: def _loss(
      self,
      params,
      pop_art_state,
      obs_tm1,
      a_tm1,
      r_t,
      discount_t,
      obs_t,
  ):
      """Loss function."""
      indices = jnp.array(0)  # Only one output for normalization.
  
      # Calculate targets by unnormalizing Q-values output by network.
      norm_q_t = self._network.apply(params, obs_t)
      q_t = rlax.unnormalize(pop_art_state, norm_q_t, indices)
      target_tm1 = r_t + discount_t * jnp.max(q_t)
  
      # Update PopArt statistics and use them to update the network parameters to
      # POP (preserve outputs precisely). If there were target networks, the
      # parameters for these would also need to be updated.
      final_linear_module_name = "mlp/~/linear_1"
      mutable_params = hk.data_structures.to_mutable_dict(params)
      linear_params = mutable_params[final_linear_module_name]
      popped_linear_params, new_pop_art_state = self._pop_art_update(
          params=linear_params,
          state=pop_art_state,
          targets=target_tm1,
          indices=indices,
      )
      mutable_params[final_linear_module_name] = popped_linear_params
      popped_params = hk.data_structures.to_immutable_dict(
          mutable_params
      )
  
      # Normalize target with updated PopArt statistics.
      norm_target_tm1 = rlax.normalize(
          new_pop_art_state, target_tm1, indices
      )
  
      # Calculate parameter update with normalized target and popped parameters.
      norm_q_t = self._network.apply(popped_params, obs_t)
      norm_q_tm1 = self._network.apply(popped_params, obs_tm1)
      td_error = (
          jax.lax.stop_gradient(norm_target_tm1) - norm_q_tm1[a_tm1]
      )
      return rlax.l2_loss(td_error), new_pop_art_state
  
imports: import collections
  from absl import app
  from absl import flags
  from bsuite.environments import catch
  from bsuite.utils import wrappers
  import haiku
  from haiku import nets
  import jax
  import jax.numpy
  import optax
  import rlax
  from rlax.examples import experiment
label_desc: Use image node_id 7 for calling the PopArtAgent obj's underlying member method code with example usage: obj._loss(params, pop_art_state, obs_tm1, a_tm1, r_t, discount_t, obs_t) and returns: rlax,  new_pop_art_state