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: Refits distribution according to actions with index of ind.
  
  Args:
    target_sample_indices: A [B, M] sized tensor indicating the index
    samples: A dict corresponding to action_spec. Each action is a [B, N, A]
      sized tensor.
  
  Returns:
    mean: A dict containing [B, A] sized tensors where each row
      is the refitted mean.
    var: A dict containing [B, A] sized tensors where each row
      is the refitted var.
docstr_tok: ["Refits","distribution","according","to","actions","with","index","of","ind",".","Args",":","target_sample_indices",":","A","[","B",",","M","]","sized","tensor","indicating","the","index","samples",":","A","dict","corresponding","to","action_spec",".","Each","action","is","a","[","B",",","N",",","A","]","sized","tensor",".","Returns",":","mean",":","A","dict","containing","[","B",",","A","]","sized","tensors","where","each","row","is","the","refitted","mean",".","var",":","A","dict","containing","[","B",",","A","]","sized","tensors","where","each","row","is","the","refitted","var","."]
code: def refit_distribution_to(self, target_sample_indices, samples):
      """Refits distribution according to actions with index of ind.
  
      Args:
        target_sample_indices: A [B, M] sized tensor indicating the index
        samples: A dict corresponding to action_spec. Each action is a [B, N, A]
          sized tensor.
  
      Returns:
        mean: A dict containing [B, A] sized tensors where each row
          is the refitted mean.
        var: A dict containing [B, A] sized tensors where each row
          is the refitted var.
      """
  
      def get_mean(
          best_samples, spec, index_range_min, index_range_max
      ):
          if spec.dtype.is_integer:
              return tf.zeros(
                  [tf.shape(target_sample_indices)[0], spec.shape[0]]
              )
          else:
              # In the following we use a customized way to calculate mean and var
              # from best_samples. The reason why we don't use standard tf.nn.moment
              # is because:
              # 1. We only want to calculate mean and var for continuous samples
              # 2. M elites may contain both continuous and categorical samples
              # 3. In a batch (B) of data, numner of continuous elite samples may be
              #    different
              # Also because the value of samples of categorical actions are all 0.0
              # We calculate mean and var in the following way.
              # mean = sum_elites_continuou / num_elites_continuous_expanded
              # var = sum((best_samples - mean)^2 - (mean)^2 * num_elites_categorical)
              # / num_elites_continuous
  
              sum_elites_continuous = tf.reduce_sum(
                  best_samples, axis=1
              )  # [B, A]
  
              # num_elites_continuous: [B]
              num_elites_continuous = tf.reduce_sum(
                  tf.cast(
                      tf.logical_and(
                          tf.greater_equal(
                              target_sample_indices, index_range_min
                          ),
                          tf.less(
                              target_sample_indices, index_range_max
                          ),
                      ),
                      tf.float32,
                  ),
                  axis=1,
              )
  
              # num_elites_continuous_expanded: [B, A]
              num_elites_continuous_expanded = tf.tile(
                  tf.expand_dims(num_elites_continuous, 1),
                  [1, spec.shape.as_list()[0]],
              )
  
              # mean: [B, A]
              mean = tf.math.divide_no_nan(
                  sum_elites_continuous, num_elites_continuous_expanded
              )
  
              return mean
  
      def get_var(
          best_samples, mean, spec, index_range_min, index_range_max
      ):
          if spec.dtype.is_integer:
              return tf.zeros(
                  [tf.shape(target_sample_indices)[0], spec.shape[0]]
              )
          else:
              # num_elites_continuous: [B]
              num_elites_continuous = tf.reduce_sum(
                  tf.cast(
                      tf.logical_and(
                          tf.greater_equal(
                              target_sample_indices, index_range_min
                          ),
                          tf.less(
                              target_sample_indices, index_range_max
                          ),
                      ),
                      tf.float32,
                  ),
                  axis=1,
              )
  
              # num_elites_continuous_expanded: [B, A]
              num_elites_continuous_expanded = tf.tile(
                  tf.expand_dims(num_elites_continuous, 1),
                  [1, spec.shape.as_list()[0]],
              )
  
              num_elites = tf.cast(
                  tf.shape(target_sample_indices)[1], tf.float32
              )
              # num_elites_categorical_expanded: [B, A]
              num_elites_categorical_expanded = (
                  num_elites - num_elites_continuous_expanded
              )
  
              # mean_expanded: [M, B, A]
              mean_expanded = mean * tf.ones(
                  [tf.shape(target_sample_indices)[1], 1, 1]
              )
              # mean_expanded: [B, M, A]
              mean_expanded = tf.transpose(mean_expanded, [1, 0, 2])
  
              var = tf.math.divide_no_nan(
                  tf.reduce_sum(
                      tf.square(best_samples - mean_expanded), axis=1
                  )
                  - tf.multiply(
                      tf.square(mean), num_elites_categorical_expanded
                  ),
                  num_elites_continuous_expanded,
              )
              return var
  
      best_samples = tf.nest.map_structure(
          lambda s: tf.gather(s, target_sample_indices, batch_dims=1),
          samples,
      )
  
      if not self._index_range_min or not self._index_range_max:
          raise ValueError(
              "sample_batch_and_clip must be called before refit_distribution_to!"
          )
  
      mean = tf.nest.map_structure(
          get_mean,
          best_samples,
          self._action_spec,
          self._index_range_min,
          self._index_range_max,
      )
      var = tf.nest.map_structure(
          get_var,
          best_samples,
          mean,
          self._action_spec,
          self._index_range_min,
          self._index_range_max,
      )
  
      return mean, var
  
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 2 for calling the GaussianActionsSampler obj's underlying member method code with example usage: obj.refit_distribution_to(target_sample_indices, samples) and returns: mean,  var, tf, mean, tf, var