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.
code: def _sample_continuous_and_transpose(
      self, mean, var, state, i, one_hot_index
  ):
      num_samples = self._number_samples_all[i]
  
      def sample_and_transpose(mean, var, spec, mask):
          if spec.dtype.is_integer:
              sample = tf.one_hot(
                  one_hot_index, self._num_mutually_exclusive_actions
              )
              sample = tf.broadcast_to(
                  sample,
                  [
                      tf.shape(mean)[0],
                      tf.constant(
                          num_samples
                      ),  # pylint: disable=cell-var-from-loop
                      tf.shape(mean)[1],
                  ],
              )
          else:
              dist = tfp.distributions.Normal(
                  loc=mean, scale=tf.sqrt(var)
              )
              # Transpose to [B, N, A]
              sample = tf.transpose(
                  dist.sample(num_samples), [1, 0, 2]
              )  # pylint: disable=cell-var-from-loop
              sample = sample * mask
          return tf.cast(sample, spec.dtype)
  
      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,
              self._action_spec,
              self._masks[i],
          )
  
          if self._sample_clippers[i]:
              for sample_clipper in self._sample_clippers[i]:
                  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:
  
                  def sample_zero_and_one_hot(spec):
                      if spec.dtype.is_integer:
                          sample = tf.one_hot(
                              one_hot_index,
                              self._num_mutually_exclusive_actions,
                          )
                      else:
                          sample = tf.zeros(spec.shape, spec.dtype)
  
                      sample = tf.broadcast_to(
                          sample,
                          tf.TensorShape(
                              [num_samples] + sample.shape.dims
                          ),
                      )
                      return tf.cast(sample, spec.dtype)
  
                  zero_samples = tf.nest.map_structure(
                      sample_zero_and_one_hot, self._action_spec
                  )
                  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[i]:
          samples_continuous = rejection_sampling(
              self._sample_rejecters[i]
          )
  
          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)
  
      return samples_continuous
  
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 3 for calling the GaussianActionsSampler obj's underlying member method code with example usage: obj._sample_continuous_and_transpose(mean, var, state, i, one_hot_index) and returns: samples_continuous, tf, samples_continuous, samples_continuous, tf