code: def test_q_head_loss():
      loss_fn = QHeadLoss()
      # example with batch_size of 3, and num_actions of 2
      target_q_values = mx.nd.array(((3, 5), (-1, -2), (0, 2)))
      pred_q_values_worse = mx.nd.array(((6, 5), (-1, -2), (0, 2)))
      pred_q_values_better = mx.nd.array(((4, 5), (-2, -2), (1, 2)))
      loss_worse = loss_fn(pred_q_values_worse, target_q_values)
      loss_better = loss_fn(pred_q_values_better, target_q_values)
      assert len(loss_worse) == 1  # (LOSS)
      loss_worse_val = loss_worse[0]
      assert loss_worse_val.ndim == 1
      assert loss_worse_val.shape[0] == 1
      assert len(loss_better) == 1  # (LOSS)
      loss_better_val = loss_better[0]
      assert loss_better_val.ndim == 1
      assert loss_better_val.shape[0] == 1
      assert loss_worse_val > loss_better_val
  
file_path: coach/rl_coach/tests/architectures/mxnet_components/heads/test_q_head.py
imports: import mxnet
  import os
  import pytest
  import sys
  from rl_coach.architectures.mxnet_components.heads.q_head import QHead, QHeadLoss
  from rl_coach.agents.clipped_ppo_agent import ClippedPPOAgentParameters
  from rl_coach.spaces import SpacesDefinition, DiscreteActionSpace
label_desc: Use image node_id 1 for calling a global function with example usage: test_q_head_loss() without return types