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: Builds a CEM-Policy given a network and a sampler.
  
  Args:
    time_step_spec: A `TimeStep` spec of the expected time_steps.
    action_spec: A nest of BoundedTensorSpec representing the actions.
    q_network: An instance of a `tf_agents.network.Network`, callable via
      `network(observation, step_type) -> (output, final_state)`.
    sampler: Samples the actions needed for the CEM.
    init_mean: A list or tuple or scalar, reprenting initial mean for actions.
    init_var: A list or tuple or scalar, reprenting initial var for actions.
    actor_policy: Optional actor policy.
    minimal_var: Minimal variance to prevent CEM distributon collapsing.
    info_spec: A policy info spec.
    num_samples: Number of samples to sample each round.
    num_elites: Number of best actions each round to refit the distribution
      with.
    num_iterations: Number of iterations to run the CEM loop.
    emit_log_probability: Whether to emit log-probs in info of `PolicyStep`.
    preprocess_state_action: The shape of state is (B, ...) and the shape of
      action is (B, N, A). When preprocess_state_action is enabled, the state
      will be tile_batched to be (BxN, ...) and the action will be reshaped to
      be (BxN, A). When preprocess_state_action is not enabled, the same
      operation needs to be done inside the network. This is helpful when the
      input have large memory requirements and the replication of state could
      happen after a few layers inside the network.
    training: Whether it is in training mode or inference mode.
    weights: A nested structure of weights w/ the same structure as action.
    name: The name of this policy. All variables in this module will fall
      under that name. Defaults to the class name.
  
  Raises:
    ValueError: If `q_network.action_spec` exists and is not compatible with
      `action_spec`.
docstr_tok: ["Builds","a","CEM-Policy","given","a","network","and","a","sampler",".","Args",":","time_step_spec",":","A","`","TimeStep","`","spec","of","the","expected","time_steps",".","action_spec",":","A","nest","of","BoundedTensorSpec","representing","the","actions",".","q_network",":","An","instance","of","a","`","tf_agents.network.Network","`",",","callable","via","`","network","(","observation",",","step_type",")","-",">","(","output",",","final_state",")","`",".","sampler",":","Samples","the","actions","needed","for","the","CEM",".","init_mean",":","A","list","or","tuple","or","scalar",",","reprenting","initial","mean","for","actions",".","init_var",":","A","list","or","tuple","or","scalar",",","reprenting","initial","var","for","actions",".","actor_policy",":","Optional","actor","policy",".","minimal_var",":","Minimal","variance","to","prevent","CEM","distributon","collapsing",".","info_spec",":","A","policy","info","spec",".","num_samples",":","Number","of","samples","to","sample","each","round",".","num_elites",":","Number","of","best","actions","each","round","to","refit","the","distribution","with",".","num_iterations",":","Number","of","iterations","to","run","the","CEM","loop",".","emit_log_probability",":","Whether","to","emit","log-probs","in","info","of","`","PolicyStep","`",".","preprocess_state_action",":","The","shape","of","state","is","(","B",",","...",")","and","the","shape","of","action","is","(","B",",","N",",","A",")",".","When","preprocess_state_action","is","enabled",",","the","state","will","be","tile_batched","to","be","(","BxN",",","...",")","and","the","action","will","be","reshaped","to","be","(","BxN",",","A",")",".","When","preprocess_state_action","is","not","enabled",",","the","same","operation","needs","to","be","done","inside","the","network",".","This","is","helpful","when","the","input","have","large","memory","requirements","and","the","replication","of","state","could","happen","after","a","few","layers","inside","the","network",".","training",":","Whether","it","is","in","training","mode","or","inference","mode",".","weights",":","A","nested","structure","of","weights","w\/","the","same","structure","as","action",".","name",":","The","name","of","this","policy",".","All","variables","in","this","module","will","fall","under","that","name",".","Defaults","to","the","class","name",".","Raises",":","ValueError",":","If","`","q_network.action_spec","`","exists","and","is","not","compatible","with","`","action_spec","`","."]
code: def __init__(
      self,
      time_step_spec: ts.TimeStep,
      action_spec: types.NestedTensorSpec,
      q_network: network.Network,
      sampler: qtopt_cem_actions_sampler.ActionsSampler,
      init_mean: types.NestedArray,
      init_var: types.NestedArray,
      actor_policy: Optional[tf_policy.TFPolicy] = None,
      minimal_var: float = 0.0001,
      info_spec: types.NestedSpecTensorOrArray = (),
      num_samples: int = 32,
      num_elites: int = 4,
      num_iterations: int = 32,
      emit_log_probability: bool = False,
      preprocess_state_action: bool = True,
      training: bool = False,
      weights: types.NestedTensorOrArray = None,
      name: Optional[str] = None,
  ):
      """Builds a CEM-Policy given a network and a sampler.
  
      Args:
        time_step_spec: A `TimeStep` spec of the expected time_steps.
        action_spec: A nest of BoundedTensorSpec representing the actions.
        q_network: An instance of a `tf_agents.network.Network`, callable via
          `network(observation, step_type) -> (output, final_state)`.
        sampler: Samples the actions needed for the CEM.
        init_mean: A list or tuple or scalar, reprenting initial mean for actions.
        init_var: A list or tuple or scalar, reprenting initial var for actions.
        actor_policy: Optional actor policy.
        minimal_var: Minimal variance to prevent CEM distributon collapsing.
        info_spec: A policy info spec.
        num_samples: Number of samples to sample each round.
        num_elites: Number of best actions each round to refit the distribution
          with.
        num_iterations: Number of iterations to run the CEM loop.
        emit_log_probability: Whether to emit log-probs in info of `PolicyStep`.
        preprocess_state_action: The shape of state is (B, ...) and the shape of
          action is (B, N, A). When preprocess_state_action is enabled, the state
          will be tile_batched to be (BxN, ...) and the action will be reshaped to
          be (BxN, A). When preprocess_state_action is not enabled, the same
          operation needs to be done inside the network. This is helpful when the
          input have large memory requirements and the replication of state could
          happen after a few layers inside the network.
        training: Whether it is in training mode or inference mode.
        weights: A nested structure of weights w/ the same structure as action.
        name: The name of this policy. All variables in this module will fall
          under that name. Defaults to the class name.
  
      Raises:
        ValueError: If `q_network.action_spec` exists and is not compatible with
          `action_spec`.
      """
      network_action_spec = getattr(q_network, "action_spec", None)
  
      if network_action_spec is not None:
          if not action_spec.is_compatible_with(network_action_spec):
              raise ValueError(
                  "action_spec must be compatible with q_network.action_spec; "
                  "instead got action_spec=%s, q_network.action_spec=%s"
                  % (action_spec, network_action_spec)
              )
  
      if q_network:
          network_utils.check_single_floating_network_output(
              q_network.create_variables(),
              expected_output_shape=(),
              label=str(q_network),
          )
          policy_state_spec = q_network.state_spec
      else:
          policy_state_spec = ()
  
      self._actor_policy = actor_policy
      self._q_network = q_network
      self._init_mean = init_mean
      self._init_var = init_var
      self._minimal_var = minimal_var
      self._num_samples = num_samples  # N
      self._num_elites = num_elites  # M
      self._num_iterations = num_iterations
      self._actions_sampler = sampler
      self._observation_spec = time_step_spec.observation
      self._training = training
      self._preprocess_state_action = preprocess_state_action
      self._weights = weights
  
      super(CEMPolicy, self).__init__(
          time_step_spec,
          action_spec,
          info_spec=info_spec,
          policy_state_spec=policy_state_spec,
          clip=False,
          emit_log_probability=emit_log_probability,
          name=name,
      )
  
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 1 to create a new CEMPolicy object from inherited base classes: tf_policy with example: obj = CEMPolicy(time_step_spec, action_spec, q_network, sampler, init_mean, init_var, actor_policy, minimal_var, info_spec, num_samples, num_elites, num_iterations, emit_log_probability, preprocess_state_action, training, weights, name)