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: Generates the distribution over next actions given the time_step.
  
  Args:
    time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
    policy_state: A Tensor, or a nested dict, list or tuple of Tensors
      representing the previous policy_state.
  
  Returns:
    A `PolicyStep` named tuple containing:
  
      `action`: A tf.distribution capturing the distribution of next actions.
      `state`: A policy state tensor for the next call to distribution.
      `info`: Optional side information such as action log probabilities.
docstr_tok: ["Generates","the","distribution","over","next","actions","given","the","time_step",".","Args",":","time_step",":","A","`","TimeStep","`","tuple","corresponding","to","`","time_step_spec","(",")","`",".","policy_state",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","representing","the","previous","policy_state",".","Returns",":","A","`","PolicyStep","`","named","tuple","containing",":","`","action","`",":","A","tf.distribution","capturing","the","distribution","of","next","actions",".","`","state","`",":","A","policy","state","tensor","for","the","next","call","to","distribution",".","`","info","`",":","Optional","side","information","such","as","action","log","probabilities","."]
code: def _distribution(
      self, time_step: ts.TimeStep, policy_state: Sequence[tf.Tensor]
  ) -> policy_step.PolicyStep:
      """Generates the distribution over next actions given the time_step.
  
      Args:
        time_step: A `TimeStep` tuple corresponding to `time_step_spec()`.
        policy_state: A Tensor, or a nested dict, list or tuple of Tensors
          representing the previous policy_state.
  
      Returns:
        A `PolicyStep` named tuple containing:
  
          `action`: A tf.distribution capturing the distribution of next actions.
          `state`: A policy state tensor for the next call to distribution.
          `info`: Optional side information such as action log probabilities.
      """
      if (
          not policy_state
          and self._q_network
          and self._q_network.state_spec
      ):
          batch_size = nest_utils.get_outer_shape(
              time_step.observation, self._observation_spec
          )[0]
          policy_state = self.get_initial_state(batch_size=batch_size)
      (
          best_actions,
          best_scores,
          best_next_policy_state,
      ) = self.actor_func(
          time_step.observation, time_step.step_type, policy_state
      )
  
      def select_best_action(actions):
          best_action = actions[..., 0, :]  # [B, A]
          return best_action
  
      best_action = tf.nest.map_structure(
          select_best_action, best_actions
      )
  
      best_action_consider_actor = best_action
      best_score_consider_actor = best_scores[..., 0]
      if self._actor_policy is not None:
          potential_best_action = self._actor_policy.action(
              time_step
          ).action
          potential_best_q, _ = self.compute_target_q(
              time_step.observation, potential_best_action
          )
          use_cem = tf.cast(
              best_score_consider_actor > potential_best_q, tf.float32
          )
          best_score_consider_actor = (
              best_score_consider_actor * use_cem
              + (tf.ones_like(use_cem, tf.float32) - use_cem)
              * potential_best_q
          )
  
          def select_best_action_consider_actor(action1, action2):
              use_cem_expanded = tf.expand_dims(
                  tf.cast(use_cem, action1.dtype), axis=-1
              )
              return action1 * use_cem_expanded + action2 * (
                  tf.ones_like(use_cem_expanded, action1.dtype)
                  - use_cem_expanded
              )
  
          best_action_consider_actor = tf.nest.map_structure(
              select_best_action_consider_actor,
              best_action_consider_actor,
              potential_best_action,
          )
  
      distribution = tf.nest.map_structure(
          tfp.distributions.Deterministic, best_action_consider_actor
      )
      if self._info_spec and "target_q" in self._info_spec:
          batch_size = nest_utils.get_outer_shape(
              time_step, self._time_step_spec
          )[0]
          info = tf.nest.map_structure(
              lambda spec: tf.zeros(
                  tf.concat([[batch_size], spec.shape], axis=-1)
              ),
              self._info_spec,
          )
          info["target_q"] = best_score_consider_actor
      else:
          batch_size = nest_utils.get_outer_shape(
              time_step, self._time_step_spec
          )[0]
          info = tf.nest.map_structure(
              lambda spec: tf.zeros(
                  tf.concat([[batch_size], spec.shape], axis=-1)
              ),
              self._info_spec,
          )
      return policy_step.PolicyStep(
          distribution, best_next_policy_state, info
      )
  
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 8 for calling the CEMPolicy obj's underlying member method code with example usage: obj._distribution(time_step, policy_state) and returns: policy_step, best_action, unknown