class_docstr: An environment based on an arbitrary classification problem.
class_docstr_tok: ["An","environment","based","on","an","arbitrary","classification","problem","."]
docstr: Initialize `ClassificationBanditEnvironment`.
  
  Args:
    dataset: a `tf.data.Dataset` consisting of two `Tensor`s, [inputs, labels]
      where inputs can be of any shape, while labels are integer class labels.
      The label tensor can be of any rank as long as it has 1 element.
    reward_distribution: a `tfd.Distribution` with event_shape `[num_classes,
      num_actions]`. Entry `[i, j]` is the reward for taking action `j` for an
      instance of class `i`.
    batch_size: if `dataset` is batched, this is the size of the batches.
    label_dtype_cast: if not None, casts dataset labels to this dtype.
    shuffle_buffer_size: If None, do not shuffle.  Otherwise, a shuffle buffer
      of the specified size is used in the environment's `dataset`.
    repeat_dataset: Makes the environment iterate on the `dataset` once
      avoiding `OutOfRangeError:  End of sequence` errors when the environment
      is stepped past the end of the `dataset`.
    prefetch_size: If None, do not prefetch.  Otherwise, a prefetch buffer of
      the specified size is used in the environment's `dataset`.
    seed: Used to make results deterministic.
    name: The name of this environment instance.
  
  Raises:
    ValueError: if `reward_distribution` does not have an event shape with
      rank 2.
docstr_tok: ["Initialize","`","ClassificationBanditEnvironment","`",".","Args",":","dataset",":","a","`","tf.data.Dataset","`","consisting","of","two","`","Tensor","`","s",",","[","inputs",",","labels","]","where","inputs","can","be","of","any","shape",",","while","labels","are","integer","class","labels",".","The","label","tensor","can","be","of","any","rank","as","long","as","it","has","1","element",".","reward_distribution",":","a","`","tfd.Distribution","`","with","event_shape","`","[","num_classes",",","num_actions","]","`",".","Entry","`","[","i",",","j","]","`","is","the","reward","for","taking","action","`","j","`","for","an","instance","of","class","`","i","`",".","batch_size",":","if","`","dataset","`","is","batched",",","this","is","the","size","of","the","batches",".","label_dtype_cast",":","if","not","None",",","casts","dataset","labels","to","this","dtype",".","shuffle_buffer_size",":","If","None",",","do","not","shuffle",".","Otherwise",",","a","shuffle","buffer","of","the","specified","size","is","used","in","the","environment","'s","`","dataset","`",".","repeat_dataset",":","Makes","the","environment","iterate","on","the","`","dataset","`","once","avoiding","`","OutOfRangeError",":","End","of","sequence","`","errors","when","the","environment","is","stepped","past","the","end","of","the","`","dataset","`",".","prefetch_size",":","If","None",",","do","not","prefetch",".","Otherwise",",","a","prefetch","buffer","of","the","specified","size","is","used","in","the","environment","'s","`","dataset","`",".","seed",":","Used","to","make","results","deterministic",".","name",":","The","name","of","this","environment","instance",".","Raises",":","ValueError",":","if","`","reward_distribution","`","does","not","have","an","event","shape","with","rank","2","."]
code: def __init__(
      self,
      dataset: tf.data.Dataset,
      reward_distribution: types.Distribution,
      batch_size: types.Int,
      label_dtype_cast: Optional[tf.DType] = None,
      shuffle_buffer_size: Optional[types.Int] = None,
      repeat_dataset: Optional[bool] = True,
      prefetch_size: Optional[types.Int] = None,
      seed: Optional[types.Int] = None,
      name: Optional[Text] = "classification",
  ):
      """Initialize `ClassificationBanditEnvironment`.
  
      Args:
        dataset: a `tf.data.Dataset` consisting of two `Tensor`s, [inputs, labels]
          where inputs can be of any shape, while labels are integer class labels.
          The label tensor can be of any rank as long as it has 1 element.
        reward_distribution: a `tfd.Distribution` with event_shape `[num_classes,
          num_actions]`. Entry `[i, j]` is the reward for taking action `j` for an
          instance of class `i`.
        batch_size: if `dataset` is batched, this is the size of the batches.
        label_dtype_cast: if not None, casts dataset labels to this dtype.
        shuffle_buffer_size: If None, do not shuffle.  Otherwise, a shuffle buffer
          of the specified size is used in the environment's `dataset`.
        repeat_dataset: Makes the environment iterate on the `dataset` once
          avoiding `OutOfRangeError:  End of sequence` errors when the environment
          is stepped past the end of the `dataset`.
        prefetch_size: If None, do not prefetch.  Otherwise, a prefetch buffer of
          the specified size is used in the environment's `dataset`.
        seed: Used to make results deterministic.
        name: The name of this environment instance.
  
      Raises:
        ValueError: if `reward_distribution` does not have an event shape with
          rank 2.
      """
  
      # Computing `action_spec`.
      event_shape = reward_distribution.event_shape
      if len(event_shape) != 2:
          raise ValueError(
              "reward_distribution must have event shape of rank 2; "
              "got event shape {}".format(event_shape)
          )
      _, num_actions = event_shape
      action_spec = tensor_spec.BoundedTensorSpec(
          shape=(),
          dtype=tf.int32,
          minimum=0,
          maximum=num_actions - 1,
          name="action",
      )
      output_shapes = tf.compat.v1.data.get_output_shapes(dataset)
  
      # Computing `time_step_spec`.
      if len(output_shapes) != 2:
          raise ValueError(
              "Dataset must have exactly two outputs; got {}".format(
                  len(output_shapes)
              )
          )
      context_shape = output_shapes[0]
      context_dtype, lbl_dtype = tf.compat.v1.data.get_output_types(
          dataset
      )
      if label_dtype_cast:
          lbl_dtype = label_dtype_cast
      observation_spec = tensor_spec.TensorSpec(
          shape=context_shape, dtype=context_dtype
      )
      time_step_spec = time_step.time_step_spec(observation_spec)
  
      super(ClassificationBanditEnvironment, self).__init__(
          action_spec=action_spec,
          time_step_spec=time_step_spec,
          batch_size=batch_size,
          name=name,
      )
  
      if shuffle_buffer_size:
          dataset = dataset.shuffle(
              buffer_size=shuffle_buffer_size,
              seed=seed,
              reshuffle_each_iteration=True,
          )
      if repeat_dataset:
          dataset = dataset.repeat()
      dataset = dataset.batch(batch_size, drop_remainder=True)
      if prefetch_size:
          dataset = dataset.prefetch(prefetch_size)
      self._data_iterator = eager_utils.dataset_iterator(dataset)
      self._current_label = tf.compat.v2.Variable(
          tf.zeros(batch_size, dtype=lbl_dtype)
      )
      self._previous_label = tf.compat.v2.Variable(
          tf.zeros(batch_size, dtype=lbl_dtype)
      )
      self._reward_distribution = reward_distribution
      self._label_dtype = lbl_dtype
  
      reward_means = self._reward_distribution.mean()
      self._optimal_action_table = tf.argmax(
          reward_means, axis=1, output_type=self._action_spec.dtype
      )
      self._optimal_reward_table = tf.reduce_max(reward_means, axis=1)
  
file_path: agents/tf_agents/bandits/environments/classification_environment.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Text
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.environments import bandit_tf_environment
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import eager_utils
label_desc: Use image node_id 1 to create a new ClassificationBanditEnvironment object from inherited base classes: bte with example: obj = ClassificationBanditEnvironment(dataset, reward_distribution, batch_size, label_dtype_cast, shuffle_buffer_size, repeat_dataset, prefetch_size, seed, name)