code: def _assert_correctness(
      self,
      res,
      sampler_type,
      anchor_point=0.0,
      categorical_action_returns=None,
  ):
      logging.info("anchor = %.2f", anchor_point)
      logging.info(
          "Output best actions (size [batch, action])\n %s", res
      )
  
      if sampler_type == "continuous_and_one_hot":
          num_categorical_actions = len(categorical_action_returns)
          argmax_action = np.argmax(categorical_action_returns)
  
          res = self.evaluate(res)
  
          # First ensure all continuous parts are 0 for whole batch.
          self.assertEqual(res["continuous"].all(), 0.0)
          # Ensure the right categorical action is chosen for whole batch.
          expected = [
              1 if i == argmax_action else 0
              for i in range(num_categorical_actions)
          ]
          expected = np.stack([expected for i in range(len(res))])
          self.assertAllEqual(expected, res["discrete"][:, 0, :])
      elif sampler_type == "hybrid":
          # All actions must be close to the anchor point.
          self.assertAllLessEqual(
              tf.abs(
                  tf.cast(res["discrete"], tf.float32) - anchor_point
              ),
              DELTA,
          )
          self.assertAllLessEqual(
              tf.abs(res["continuous"] - anchor_point), DELTA
          )
      elif sampler_type == "continuous":
          # All actions must be close to the anchor point.
          self.assertAllLessEqual(tf.abs(res - anchor_point), DELTA)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl import flags
  from absl import logging
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.networks import network
  from tf_agents.policies import qtopt_cem_policy
  from tf_agents.policies.samplers import qtopt_cem_actions_sampler_continuous
  from tf_agents.policies.samplers import qtopt_cem_actions_sampler_continuous_and_one_hot
  from tf_agents.policies.samplers import qtopt_cem_actions_sampler_hybrid
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
label_desc: Use image node_id 5 for calling the CEMPolicyTest obj's underlying member method code with example usage: obj._assert_correctness(res, sampler_type, anchor_point, categorical_action_returns) without return types