class_docstr: Class to build CEM Policy.
  
  Some notations used in the comments below are:
  
  B: batch_size
  A: action_size
  N: num_samples
  M: num_elites
  
  Supports the following three cases.
  
  1. Nested action_spec containing 1d continuous action.
  2. Nested action_spec containing 1d continuous or discrete action.
  3. Nested action_spec containing 1d continuous action and 1 one_hot action.
  
  Note: all actions need to be 1-d array.
docstr: Returns an action to perform using CEM given the q network.
  
  Args:
    observation: Observation for which we need to find a CEM based action.
    step_type: A `StepType` enum value.
    policy_state: A Tensor, or a nested dict, list or tuple of Tensors
      representing the previous policy_state.
  
  Returns:
    A [B, A] tensor representing the action taken by the actor,
    that is best in the CEM sense
docstr_tok: ["Returns","an","action","to","perform","using","CEM","given","the","q","network",".","Args",":","observation",":","Observation","for","which","we","need","to","find","a","CEM","based","action",".","step_type",":","A","`","StepType","`","enum","value",".","policy_state",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","representing","the","previous","policy_state",".","Returns",":","A","[","B",",","A","]","tensor","representing","the","action","taken","by","the","actor",",","that","is","best","in","the","CEM","sense"]
code: def actor_func(
      self,
      observation: types.NestedTensorOrArray,
      step_type: Optional[tf.Tensor],
      policy_state: Sequence[tf.Tensor],
  ) -> Tuple[
      types.NestedTensor, types.NestedTensor, Sequence[tf.Tensor]
  ]:
      """Returns an action to perform using CEM given the q network.
  
      Args:
        observation: Observation for which we need to find a CEM based action.
        step_type: A `StepType` enum value.
        policy_state: A Tensor, or a nested dict, list or tuple of Tensors
          representing the previous policy_state.
  
      Returns:
        A [B, A] tensor representing the action taken by the actor,
        that is best in the CEM sense
      """
      batch_size = nest_utils.get_outer_shape(
          observation, self._observation_spec
      )[0]
      outer_rank = nest_utils.get_outer_rank(
          observation, self._observation_spec
      )
      if outer_rank == 2:
          seq_size = nest_utils.get_outer_shape(
              observation, self._observation_spec
          )[1]
          batch_size = batch_size * seq_size
      # Run the sampler to sample N actions from a distribution
      # Run the q_func for each of those actions and choose the best M (<N)
      # Fit another distribution to the best M actions and repeat k times.
  
      def body(
          mean,
          var,
          i,
          iters,
          best_actions,
          best_scores,
          best_next_policy_state,
      ):
          """Defines the body of the while loop in graph.
  
          Args:
            mean: A [B, A] sized tensor, where, a is the size of the action space,
              indicating the mean value of the sample distribution var : [B, A]
              sized tensor, where, a is the size of the action space, indicating the
              variance value of the sample distribution
            i: tensor indicating the iteration index
            iters: tensor representing the max number of iteration to run
            best_actions: the best action so far as per the CEM iteration [B, A]
            best_scores: best score.
            best_next_policy_state: A Tensor, or a nested dict, list or tuple of
              Tensors representing the best next policy_state.
  
          Returns:
            list of tensors [new_mean, new_var, i+1, iters, new_best_actions,
              new_best_scores, new_best_next_policy_state]
          """
          del best_actions, best_scores, best_next_policy_state
  
          # Prevent variance to collapse to 0.0 (stuck in local minimum).
          var = tf.nest.map_structure(
              lambda v: tf.maximum(v, self._minimal_var), var
          )
  
          # Sample a batch of actions with the shape of [B, N, A] or [BxT, N, A]
          actions = self._actions_sampler.sample_batch_and_clip(
              self._num_samples, mean, var, observation
          )  # pytype: disable=wrong-arg-count  # trace-all-classes
  
          if outer_rank == 2:
              scores, next_policy_state = self._score_with_time(
                  observation,
                  actions,
                  step_type,
                  policy_state,
                  seq_size,
              )  # [BxT, N]
          else:
              scores, next_policy_state = self._score(
                  observation, actions, step_type, policy_state
              )  # [B, N]
  
          best_scores, ind = tf.nn.top_k(
              scores, self._num_elites
          )  # ind: [B, M]
  
          actions_float = tf.nest.map_structure(
              lambda t: tf.cast(t, tf.float32), actions
          )
          mean, var = self._actions_sampler.refit_distribution_to(
              ind, actions_float
          )
  
          best_next_policy_state = next_policy_state
  
          def select_best_actions(actions):
              best_actions = tf.gather(actions, ind, batch_dims=1)
              return best_actions
  
          best_actions = tf.nest.map_structure(
              select_best_actions, actions
          )
          return [
              mean,
              var,
              tf.add(i, 1),
              iters,
              best_actions,
              best_scores,
              best_next_policy_state,
          ]
  
      def cond(
          mean,
          var,
          i,
          iters,
          best_actions,
          best_scores,
          best_next_policy_state,
      ):
          del (
              mean,
              var,
              best_actions,
              best_scores,
              best_next_policy_state,
          )
          return tf.less(i, iters)
  
      mean, var = self._initial_params(batch_size)
      iters = tf.constant(self._num_iterations)
  
      def init_best_actions(action_spec):
          best_actions = tf.zeros(
              [
                  batch_size,
                  self._num_elites,
                  action_spec.shape.as_list()[0],
              ],
              dtype=action_spec.dtype,
          )
          return best_actions
  
      best_actions = tf.nest.map_structure(
          init_best_actions, self._action_spec
      )
      best_scores = tf.zeros(
          [batch_size, self._num_elites], dtype=tf.float32
      )
  
      # Run the while loop for CEM in-graph.
      mean_shape = tf.nest.map_structure(
          lambda m: [None] + m.get_shape()[1:], mean
      )
      var_shape = tf.nest.map_structure(
          lambda v: [None] + v.get_shape()[1:], var
      )
      best_action_shape = tf.nest.map_structure(
          lambda a: [None] + a.get_shape()[1:], best_actions
      )
      elites_shape = tf.TensorShape([None, self._num_elites])
      policy_state_shape = ()
      if common.safe_has_state(policy_state):
          policy_state_shape = tf.nest.map_structure(
              lambda state: state.get_shape(), policy_state
          )
  
      (
          _,
          _,
          _,
          _,
          best_actions,
          best_scores,
          best_next_policy_state,
      ) = tf.while_loop(
          cond=cond,
          body=body,
          loop_vars=[
              mean,
              var,
              0,
              iters,
              best_actions,
              best_scores,
              policy_state,
          ],
          shape_invariants=[
              mean_shape,
              var_shape,
              tf.TensorShape(()),
              iters.get_shape(),
              best_action_shape,
              elites_shape,
              policy_state_shape,
          ],
      )
  
      if outer_rank == 2:
          best_actions = tf.nest.map_structure(
              lambda x: tf.reshape(  # pylint: disable=g-long-lambda
                  x, [-1, seq_size, self._num_elites, tf.shape(x)[-1]]
              ),
              best_actions,
          )
          best_scores = tf.reshape(
              best_scores, [-1, seq_size, self._num_elites]
          )
  
      return best_actions, best_scores, best_next_policy_state
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  from typing import Any, Optional, Sequence, Tuple, Union
  import gin
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.networks import network
  from tf_agents.networks import utils
  from tf_agents.policies import tf_policy
  from tf_agents.policies.samplers import qtopt_cem_actions_sampler
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 3 for calling the CEMPolicy obj's underlying member method code with example usage: obj.actor_func(observation, step_type, policy_state) and returns: best_actions,  best_scores,  best_next_policy_state, list, tf, best_actions, best_actions