class_docstr: Action sampler that samples continuous actions using Gaussian distribution and one_hot actions.
  
  Supports dict action_spec with arbitrary 1d continuous actions and 1 one_hot
  action.
  
  Given a batch of distribution params(including mean and var) for K (K >= 1)
  mutually exclusive continuous actions, sample [B, N-S+K, A] continuous actions
  and [B, S-K, S] one_hot actions,  where 'N' means num_samples,
  'B' means batch_size, 'A' means action_size_continuous,
  'S' means num_mutually_exclusive_actions. S-K is the possible number of
  categorical actions. If K > 1, the number of samples: N-S+K will be divided
  evenly among these mutually exclusive continuous actions.
docstr: Samples and clips a batch of actions [B, N, A] with mean and var.
  
  Args:
    num_samples: Number of actions to sample each round.
    mean: A dict containing [B, A] shaped tensor representing the mean of the
      actions to be sampled.
    var: A dict containing [B, A] shaped tensor representing the variance of
      the actions to be sampled.
    state: A dict of state tensors constructed according to oberservation_spec
      of the task.
  
  Returns:
    actions:  A dict containing tensor of sampled actions with
      shape [B, N, A]
docstr_tok: ["Samples","and","clips","a","batch","of","actions","[","B",",","N",",","A","]","with","mean","and","var",".","Args",":","num_samples",":","Number","of","actions","to","sample","each","round",".","mean",":","A","dict","containing","[","B",",","A","]","shaped","tensor","representing","the","mean","of","the","actions","to","be","sampled",".","var",":","A","dict","containing","[","B",",","A","]","shaped","tensor","representing","the","variance","of","the","actions","to","be","sampled",".","state",":","A","dict","of","state","tensors","constructed","according","to","oberservation_spec","of","the","task",".","Returns",":","actions",":","A","dict","containing","tensor","of","sampled","actions","with","shape","[","B",",","N",",","A","]"]
code: def sample_batch_and_clip(self, num_samples, mean, var, state=None):
      """Samples and clips a batch of actions [B, N, A] with mean and var.
  
      Args:
        num_samples: Number of actions to sample each round.
        mean: A dict containing [B, A] shaped tensor representing the mean of the
          actions to be sampled.
        var: A dict containing [B, A] shaped tensor representing the variance of
          the actions to be sampled.
        state: A dict of state tensors constructed according to oberservation_spec
          of the task.
  
      Returns:
        actions:  A dict containing tensor of sampled actions with
          shape [B, N, A]
      """
      # At least one sample for each kind of one hot action is generated.
      assert num_samples >= self._num_mutually_exclusive_actions
      num_samples_continuous_total = (
          num_samples
          - self._num_mutually_exclusive_actions
          + self._num_sub_actions
          - 1
      )
      num_samples_continuous_each = (
          num_samples_continuous_total
          // self._num_sub_continuous_actions
      )
  
      # When sampling N samples, we use min_index and max_index to cut N samples
      # into several segments for each sub_actions.
      min_index = 0
      max_index = 0
      self._number_samples_all = []
      for i in range(self._num_sub_actions):
          min_index = max_index
          if i == self._categorical_index:
              max_index += self._num_sub_categorical_actions
          elif i == self._num_sub_actions - 1:
              max_index = num_samples
          else:
              max_index += num_samples_continuous_each
  
          for k in self._action_spec.keys():
              if k in self._sub_actions_fields[i]:
                  self._index_range_min[k] = min_index
                  self._index_range_max[k] = max_index
  
          self._number_samples_all.append(max_index - min_index)
  
      samples_all = []
  
      one_hot_index = 0
      for i in range(self._num_sub_actions):
          if i == self._categorical_index:
              # Samples one_hot actions.
              def sample_one_hot(mean, spec):
                  if spec.dtype.is_integer:
                      full_one_hot = tf.eye(
                          self._num_mutually_exclusive_actions,
                          dtype=tf.int32,
                      )  # [S, S]
  
                      categorical_one_hot = tf.gather(
                          full_one_hot,
                          tf.range(
                              one_hot_index,
                              one_hot_index
                              + self._num_sub_categorical_actions,
                          ),
                      )
  
                      return tf.broadcast_to(
                          categorical_one_hot,
                          [
                              tf.shape(mean)[0],
                              self._num_sub_categorical_actions,
                              spec.shape[0],
                          ],
                      )
                  else:
                      return tf.zeros(
                          [
                              tf.shape(mean)[0],
                              self._num_sub_categorical_actions,
                              spec.shape[0],
                          ]
                      )
  
              samples_one_hot = tf.nest.map_structure(
                  sample_one_hot, mean, self._action_spec
              )
  
              samples_one_hot = tf.nest.map_structure(
                  common.clip_to_spec,
                  samples_one_hot,
                  self._action_spec,
              )
  
              samples_all.append(samples_one_hot)
  
              one_hot_index += self._num_sub_categorical_actions
          else:
              samples_continuous = (
                  self._sample_continuous_and_transpose(
                      mean, var, state, i, one_hot_index
                  )
              )
              samples_all.append(samples_continuous)
  
              one_hot_index += 1
  
      samples_all = tf.nest.map_structure(
          lambda *tensors: tf.concat(tensors, axis=1), *samples_all
      )
  
      return samples_all
  
file_path: agents/tf_agents/policies/samplers/qtopt_cem_actions_sampler_continuous_and_one_hot.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.policies.samplers import qtopt_cem_actions_sampler
  from tf_agents.utils import common
label_desc: Use image node_id 4 for calling the GaussianActionsSampler obj's underlying member method code with example usage: obj.sample_batch_and_clip(num_samples, mean, var, state) and returns: samples_all, tf, tf