gfn.preprocessors
=================

.. py:module:: gfn.preprocessors


Attributes
----------

.. autoapisummary::

   gfn.preprocessors.GetStatesIndicesFn


Classes
-------

.. autoapisummary::

   gfn.preprocessors.EnumPreprocessor
   gfn.preprocessors.IdentityPreprocessor
   gfn.preprocessors.KHotPreprocessor
   gfn.preprocessors.OneHotPreprocessor
   gfn.preprocessors.Preprocessor


Module Contents
---------------

.. py:class:: EnumPreprocessor(get_states_indices)

   Bases: :py:obj:`Preprocessor`


   Preprocessor for environments with discrete, enumerable states.

   This preprocessor converts discrete states to their unique integer indices,
   making them suitable for neural network processing. It is designed for
   environments with a finite number of states where each state can be uniquely
   identified by an index.

   .. attribute:: output_dim

      Always 1, as states are represented by single indices.

   .. attribute:: get_states_indices

      Function that returns unique indices for states.


   .. py:attribute:: get_states_indices


   .. py:method:: preprocess(states)

      Preprocesses the states by returning their unique indices.

      :param states: The discrete states to preprocess.

      :returns: A tensor of shape (*batch_shape, 1) containing the unique indices of the
                states.



.. py:data:: GetStatesIndicesFn

.. py:class:: IdentityPreprocessor(output_dim, target_dtype = None)

   Bases: :py:obj:`Preprocessor`


   Simple preprocessor that returns states without modification.

   This preprocessor serves as the default preprocessor. It can handle both graph and
   tensor-based states by returning them as-is.

   .. attribute:: output_dim

      The dimensionality of the input states.


   .. py:method:: preprocess(states)

      Returns the states without any preprocessing.

      :param states: The states to preprocess.

      :returns: Tensor or GeometricBatch representing the states.



.. py:class:: KHotPreprocessor(height, ndim)

   Bases: :py:obj:`Preprocessor`


   Preprocessor for grid-structured discrete states with multi-dimensional encoding.

   This preprocessor is designed for environments with grid-like state spaces where
   each dimension can take on a finite number of values. It creates a k-hot encoding
   where each dimension is one-hot encoded and then concatenated.

   .. attribute:: output_dim

      The total output dimension (height * ndim).

   .. attribute:: height

      Number of unique values per dimension.

   .. attribute:: ndim

      Number of dimensions in the state space.


   .. py:attribute:: height


   .. py:attribute:: ndim


   .. py:attribute:: output_dim
      :type:  int


   .. py:method:: preprocess(states)

      Preprocesses the states by creating k-hot encoded vectors.

      Each dimension of the state is one-hot encoded and then concatenated into
      a single vector.

      :param states: The discrete states to preprocess.

      :returns: A tensor of shape (*batch_shape, height * ndim) containing k-hot encoded states.

      .. note:: This preprocessor only works for integer state tensors.



.. py:class:: OneHotPreprocessor(n_states, get_states_indices)

   Bases: :py:obj:`Preprocessor`


   Preprocessor that converts discrete states to one-hot encoded vectors.

   This preprocessor is designed for environments with enumerable states where each
   state is represented as a one-hot vector. The output dimension equals the total
   number of possible states.

   .. attribute:: output_dim

      The total number of states in the environment.

   .. attribute:: get_states_indices

      Function that returns unique indices for states.


   .. py:attribute:: get_states_indices


   .. py:attribute:: output_dim
      :type:  int


   .. py:method:: preprocess(states)

      Preprocesses the states by converting them to one-hot encoded vectors.

      :param states: The discrete states to preprocess.

      :returns: A tensor of shape (*batch_shape, n_states) containing one-hot encoded states.



.. py:class:: Preprocessor(output_dim, target_dtype = None)

   Bases: :py:obj:`abc.ABC`


   Base class for state preprocessors.

   Preprocessors transform raw state tensors into formats suitable for neural network
   inputs. They handle the conversion from environment-specific state representations
   to standardized tensor formats that can be processed by neural networks.

   .. attribute:: output_dim

      The dimensionality of the preprocessed output tensor, which is
      compatible with the neural network that will be used. If None, the output
      dimension will not be checked.


   .. py:method:: __call__(states)

      Calls the preprocess method and validates the output shape.

      :param states: The states to preprocess.

      :returns: The preprocessed states as a tensor or GeometricBatch.



   .. py:method:: __repr__()

      Returns a string representation of the Preprocessor.

      :returns: A string summary of the Preprocessor.



   .. py:attribute:: output_dim


   .. py:method:: preprocess(states)
      :abstractmethod:


      Transforms the states to the input format for neural networks.

      :param states: The states to preprocess.

      :returns: A tensor of shape (*batch_shape, output_dim) containing the preprocessed
                states.



   .. py:attribute:: target_dtype
      :value: None