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: Scores the sample actions internally as part of CEM.
  
  Args:
    observation: A batch of state tensors or NamedTuples, whatever the q_func
      will handle. CEM is agnostic to it.
    sample_actions: A [BxT, N, A] sized tensor, where batch is the batch size,
      N is the sample size for the CEM, a is the size of the action space.
    step_type: A `StepType` enum value.
    policy_state: A Tensor, or a nested dict, list or tuple of Tensors
      representing the previous policy_state.
    seq_size: Size for the time dimension.
  
  Returns:
    a tensor of shape [BxT, N] representing the scores for the actions.
docstr_tok: ["Scores","the","sample","actions","internally","as","part","of","CEM",".","Args",":","observation",":","A","batch","of","state","tensors","or","NamedTuples",",","whatever","the","q_func","will","handle",".","CEM","is","agnostic","to","it",".","sample_actions",":","A","[","BxT",",","N",",","A","]","sized","tensor",",","where","batch","is","the","batch","size",",","N","is","the","sample","size","for","the","CEM",",","a","is","the","size","of","the","action","space",".","step_type",":","A","`","StepType","`","enum","value",".","policy_state",":","A","Tensor",",","or","a","nested","dict",",","list","or","tuple","of","Tensors","representing","the","previous","policy_state",".","seq_size",":","Size","for","the","time","dimension",".","Returns",":","a","tensor","of","shape","[","BxT",",","N","]","representing","the","scores","for","the","actions","."]
code: def _score_with_time(
      self,
      observation: types.NestedTensorOrArray,
      sample_actions: types.NestedTensorOrArray,
      step_type: Optional[tf.Tensor],
      policy_state: Sequence[tf.Tensor],
      seq_size: tf.Tensor,
  ) -> Tuple[tf.Tensor, Sequence[tf.Tensor]]:
      """Scores the sample actions internally as part of CEM.
  
      Args:
        observation: A batch of state tensors or NamedTuples, whatever the q_func
          will handle. CEM is agnostic to it.
        sample_actions: A [BxT, N, A] sized tensor, where batch is the batch size,
          N is the sample size for the CEM, a is the size of the action space.
        step_type: A `StepType` enum value.
        policy_state: A Tensor, or a nested dict, list or tuple of Tensors
          representing the previous policy_state.
        seq_size: Size for the time dimension.
  
      Returns:
        a tensor of shape [BxT, N] representing the scores for the actions.
      """
  
      # Reshape input to be compatible with network: [BxN, T, A]
  
      def expand_to_megabatch(feature):
          # Collapse second dimension of megabatch.
          dim = tf.shape(feature)[2]
          dim_sample = tf.shape(feature)[1]
          feature = tf.reshape(feature, [-1, seq_size, dim_sample, dim])
          feature = tf.transpose(feature, [0, 2, 1, 3])
          return tf.reshape(feature, [-1, seq_size, dim])
  
      def decouple_batch_time(feature):
          dim = tf.shape(feature)[2]
          dim_sample = tf.shape(feature)[1]
          return tf.reshape(feature, [-1, seq_size, dim_sample, dim])
  
      if self._preprocess_state_action:
          # [BxT, N, A] -> [BxN, T, A]
          sample_actions = tf.nest.map_structure(
              expand_to_megabatch, sample_actions
          )
          # TODO(b/138331671) Move tf.contrib.seq2seq.tile_batch to utils.common
          # [B, T, ...] -> [BxN, T, ...]
          observation = nest_utils.tile_batch(
              observation, self._num_samples
          )
          step_type = nest_utils.tile_batch(
              step_type, self._num_samples
          )
          policy_state = nest_utils.tile_batch(
              policy_state, self._num_samples
          )
      else:
          # [BxT, N, A] -> [B, T, N, A]
          sample_actions = tf.nest.map_structure(
              decouple_batch_time, sample_actions
          )
  
      scores, next_policy_state = self.compute_target_q(
          observation, sample_actions, step_type, policy_state
      )  # [BxN, T]
  
      if self._preprocess_state_action:
          next_policy_state = tf.nest.map_structure(
              lambda x: tf.reshape(
                  x,
                  [-1, self._num_samples]
                  + x.shape.as_list()[
                      1:
                  ],  # pylint:disable=g-long-lambda
              )[:, 0, ...],
              next_policy_state,
          )
  
      scores = tf.reshape(
          scores, [-1, self._num_samples, seq_size]
      )  # [B, N, T]
      scores = tf.transpose(scores, [0, 2, 1])  # [B, T, N]
      scores = tf.reshape(scores, [-1, self._num_samples])  # [BxT, N]
  
      return scores, 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 6 for calling the CEMPolicy obj's underlying member method code with example usage: obj._score_with_time(observation, sample_actions, step_type, policy_state, seq_size) and returns: scores,  next_policy_state, tf, tf