class_docstr: Continuous Gaussian actions sampler.
  
  Supports nested action_spec with 1d continuous actions.
  
  Given a batch of distribution params(including mean and var), sample and
  return [B, N, A] actions, where 'B' means batch_size, 'N' means num_samples,
  'A' means action_size.
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 nested structure containing [B, A] shaped tensor representing the
      mean of the actions to be sampled.
    var: A nested structure containing [B, A] shaped tensor representing the
      variance of the actions to be sampled.
    state: Nested state tensor constructed according to oberservation_spec of
      the task.
  
  Returns:
    actions:  A nested structure 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","nested","structure","containing","[","B",",","A","]","shaped","tensor","representing","the","mean","of","the","actions","to","be","sampled",".","var",":","A","nested","structure","containing","[","B",",","A","]","shaped","tensor","representing","the","variance","of","the","actions","to","be","sampled",".","state",":","Nested","state","tensor","constructed","according","to","oberservation_spec","of","the","task",".","Returns",":","actions",":","A","nested","structure","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 nested structure containing [B, A] shaped tensor representing the
          mean of the actions to be sampled.
        var: A nested structure containing [B, A] shaped tensor representing the
          variance of the actions to be sampled.
        state: Nested state tensor constructed according to oberservation_spec of
          the task.
  
      Returns:
        actions:  A nested structure containing tensor of sampled actions with
          shape [B, N, A]
      """
  
      def sample_and_transpose(mean, var):
          dist = tfp.distributions.Normal(loc=mean, scale=tf.sqrt(var))
          sample = tf.transpose(dist.sample(num_samples), [1, 0, 2])
          return sample
  
      batch_size = tf.shape(tf.nest.flatten(mean)[0])[0]
  
      def sample_fn(mean_sample, var_sample, state_sample):
          # [B, N, A]
          samples_continuous = tf.nest.map_structure(
              sample_and_transpose, mean_sample, var_sample
          )
  
          if self._sample_clippers:
              for sample_clipper in self._sample_clippers:
                  samples_continuous = sample_clipper(
                      samples_continuous, state_sample
                  )
  
          samples_continuous = tf.nest.map_structure(
              common.clip_to_spec, samples_continuous, self._action_spec
          )
          return samples_continuous
  
      @tf.function
      def rejection_sampling(sample_rejector):
          valid_batch_samples = tf.nest.map_structure(
              lambda spec: tf.TensorArray(spec.dtype, size=batch_size),
              self._action_spec,
          )
  
          for b_indx in tf.range(batch_size):
              k = tf.constant(0)
              # pylint: disable=cell-var-from-loop
              valid_samples = tf.nest.map_structure(
                  lambda spec: tf.TensorArray(
                      spec.dtype, size=num_samples
                  ),
                  self._action_spec,
              )
  
              count = tf.constant(0)
              while count < self._max_rejection_iterations:
                  count += 1
                  mean_sample = tf.nest.map_structure(
                      lambda t: tf.expand_dims(
                          tf.gather(t, b_indx), axis=0
                      ),
                      mean,
                  )
                  var_sample = tf.nest.map_structure(
                      lambda t: tf.expand_dims(
                          tf.gather(t, b_indx), axis=0
                      ),
                      var,
                  )
                  if state is not None:
                      state_sample = tf.nest.map_structure(
                          lambda t: tf.expand_dims(
                              tf.gather(t, b_indx), axis=0
                          ),
                          state,
                      )
                  else:
                      state_sample = None
  
                  samples = sample_fn(
                      mean_sample, var_sample, state_sample
                  )  # n, a
  
                  mask = sample_rejector(samples, state_sample)
  
                  mask = mask[0, ...]
                  mask_index = tf.where(mask)[:, 0]
  
                  num_mask = tf.shape(mask_index)[0]
                  if num_mask == 0:
                      continue
  
                  good_samples = tf.nest.map_structure(
                      lambda t: tf.gather(t, mask_index, axis=1)[
                          0, ...
                      ],
                      samples,
                  )
  
                  for sample_idx in range(num_mask):
                      if k >= num_samples:
                          break
                      valid_samples = tf.nest.map_structure(
                          lambda gs, vs: vs.write(
                              k, gs[sample_idx : sample_idx + 1, ...]
                          ),
                          good_samples,
                          valid_samples,
                      )
                      k += 1
  
              if k < num_samples:
                  zero_samples = tensor_spec.zero_spec_nest(
                      self._action_spec, outer_dims=(num_samples - k,)
                  )
                  for sample_idx in range(num_samples - k):
                      valid_samples = tf.nest.map_structure(
                          lambda gs, vs: vs.write(
                              k, gs[sample_idx : sample_idx + 1, ...]
                          ),
                          zero_samples,
                          valid_samples,
                      )
  
              valid_samples = tf.nest.map_structure(
                  lambda vs: vs.concat(), valid_samples
              )
  
              valid_batch_samples = tf.nest.map_structure(
                  lambda vbs, vs: vbs.write(b_indx, vs),
                  valid_batch_samples,
                  valid_samples,
              )
  
          samples_continuous = tf.nest.map_structure(
              lambda a: a.stack(), valid_batch_samples
          )
          return samples_continuous
  
      if self._sample_rejecters:
          samples_continuous = rejection_sampling(
              self._sample_rejecters
          )
  
          def set_b_n_shape(t):
              t.set_shape(
                  tf.TensorShape([None, num_samples] + t.shape[2:].dims)
              )
  
          tf.nest.map_structure(set_b_n_shape, samples_continuous)
      else:
          samples_continuous = sample_fn(mean, var, state)
  
      if self._support_integer:
          samples_continuous = tf.nest.map_structure(
              lambda t, s: tf.cast(t, s.dtype),
              samples_continuous,
              self._action_spec,
          )
      return samples_continuous
  
file_path: agents/tf_agents/policies/samplers/qtopt_cem_actions_sampler_continuous.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.specs import tensor_spec
  from tf_agents.utils import common
label_desc: Use image node_id 3 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_continuous, sample, samples_continuous, samples_continuous