Merge pull request #14 from DanielAvdar/init

Init

Author: DanielAvdar

docs/source/api.rst

@@ -1,10 +1,10 @@
 API Reference
-============
+=============
 
 This section provides detailed information about the prob-spaces API.
 
 Probability Spaces
------------------
+------------------
 
 .. automodule:: prob_spaces
    :members:
@@ -20,7 +20,7 @@ Converter
    :show-inheritance:
 
 Distributions
-------------
+-------------
 
 .. automodule:: prob_spaces.dists
    :members:

docs/source/index.rst

@@ -19,18 +19,40 @@ environments.
    :target: https://opensource.org/licenses/MIT
    :alt: License
 
+.. image:: https://img.shields.io/badge/ubuntu-blue?logo=ubuntu
+   :alt: Ubuntu
+
+.. image:: https://img.shields.io/badge/ubuntu-blue?logo=windows
+   :alt: Windows
+
+.. image:: https://img.shields.io/badge/ubuntu-blue?logo=apple
+   :alt: MacOS
+
+
+
 .. toctree::
-   :maxdepth: 2
-   :caption: Contents:
-
-   introduction
-   installation
-   api
-   modules/discrete
-   modules/multi_discrete
-   modules/box
-   modules/dict
-   modules/converter
+   :maxdepth: 1
+   :caption: Introduction:
+
+   ./introduction
+   ./installation
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Usage:
+
+   ./modules/discrete
+   ./modules/multi_discrete
+   ./modules/box
+   ./modules/dict
+   ./modules/converter
+
+.. toctree::
+   :maxdepth: 1
+   :caption: API Reference:
+
+   ./api
+
 
 Indices and tables
 ==================

docs/source/installation.rst

@@ -1,3 +1,5 @@
+.. _installation:
+
 Installation
 ===========
 

docs/source/introduction.rst

@@ -1,3 +1,5 @@
+.. _introduction:
+
 Introduction
 ===========
 

docs/source/modules/box.rst

@@ -1,5 +1,7 @@
+.. _box:
+
 Box Space
-========
+=========
 
 The ``BoxDist`` class extends the Gymnasium Box space to create continuous probability distributions.
 
@@ -14,7 +16,7 @@ API Reference
 -------------
 
 .. autoclass:: prob_spaces.box.BoxDist
-   :members:
+   :members: __call__
    :undoc-members:
    :show-inheritance:
 

docs/source/modules/converter.rst

@@ -1,3 +1,5 @@
+.. _converter:
+
 Space Converter
 ===============
 

docs/source/modules/dict.rst

@@ -1,5 +1,7 @@
+.. _dict:
+
 Dict Space
-=========
+==========
 
 The ``DictDist`` class extends the Gymnasium Dict space to create nested distributions.
 
@@ -13,7 +15,7 @@ API Reference
 ------------
 
 .. autoclass:: prob_spaces.dict.DictDist
-   :members:
+   :members: __call__
    :undoc-members:
    :show-inheritance:
 

docs/source/modules/discrete.rst

@@ -1,5 +1,7 @@
+.. _discrete:
+
 Discrete Space
-=============
+==============
 
 The ``DiscreteDist`` class extends the Gymnasium Discrete space to create categorical distributions.
 
@@ -13,7 +15,7 @@ API Reference
 ------------
 
 .. autoclass:: prob_spaces.discrete.DiscreteDist
-   :members:
+   :members: __call__
    :undoc-members:
    :show-inheritance:
 

docs/source/modules/multi_discrete.rst

@@ -1,5 +1,7 @@
+.. _multi_discrete:
+
 MultiDiscrete Space
-==================
+===================
 
 The ``MultiDiscreteDist`` class extends the Gymnasium MultiDiscrete space to create categorical distributions
 for multiple discrete variables.
@@ -14,7 +16,7 @@ API Reference
 ------------
 
 .. autoclass:: prob_spaces.multi_discrete.MultiDiscreteDist
-   :members:
+   :members: __call__
    :undoc-members:
    :show-inheritance:
 

prob_spaces/__init__.py

@@ -0,0 +1,13 @@
+"""Probability distribution classes for various Gymnasium spaces."""
+
+from .converter import convert_to_prob_space
+from .dict import DictDist
+from .discrete import DiscreteDist
+from .multi_discrete import MultiDiscreteDist
+
+__all__ = [
+    "MultiDiscreteDist",
+    "DiscreteDist",
+    "DictDist",
+    "convert_to_prob_space",
+]

prob_spaces/box.py

@@ -33,6 +33,17 @@ def transforms(self, device: th.device) -> list:
         return transforms
 
     def __call__(self, loc: th.Tensor, scale: th.Tensor) -> th.distributions.Distribution:
+        """
+        Generates a transformed probability distribution based on the input location and scale
+        parameters. The method constructs a base distribution, applies a sequence of
+        transformations to it, and returns the resulting transformed distribution. This
+        allows for creating flexible and expressive probability distributions.
+
+        :param loc: A tensor specifying the location parameters for the base distribution.
+        :param scale: A tensor specifying the scale parameters for the base distribution.
+        :return: A transformed distribution object derived from the specified base distribution
+            and transformations.
+        """
         dist = self.base_dist(loc, scale, validate_args=True)  # type: ignore
         transforms = self.transforms(loc.device)
         transformed_dist = TransformedDistribution(dist, transforms, validate_args=True)

prob_spaces/discrete.py

@@ -6,6 +6,17 @@
 
 class DiscreteDist(spaces.Discrete):
     def __call__(self, prob: th.Tensor, mask: th.Tensor = None) -> MaskedCategorical:
+        """
+        Compute and return a masked categorical distribution based on the given probability
+        tensor and an optional mask. The distribution incorporates specific probabilities
+        and constraints defined by the provided input.
+
+        :param prob: A tensor representing the probabilities for each category.
+        :param mask: A tensor specifying a mask to limit the valid categories.
+                     Defaults to a tensor of ones if not provided.
+        :return: A MaskedCategorical distribution constructed with given probabilities,
+                 mask, and starting values.
+        """
         probs = prob.reshape(self.n)  # type: ignore
         start = self.start
         mask = mask if mask is not None else th.ones_like(probs, dtype=th.bool, device=probs.device)

prob_spaces/multi_discrete.py

@@ -37,6 +37,23 @@ def _internal_mask(self) -> NDArray[np.bool_]:
         return mask
 
     def __call__(self, prob: th.Tensor, mask: th.Tensor = None) -> MaskedCategorical:
+        """
+        Applies a transformation to the input probability tensor and optional mask, creating
+        a `MaskedCategorical` distribution. The method reshapes the input probabilities
+        to match the specified `nvec` dimensions, applies an optional mask for masking
+        specific probabilities, and combines these with an internal mask. The result
+        is used to create a `MaskedCategorical` distribution.
+
+        :param prob: A tensor containing probabilities to be reshaped and used in
+            constructing the distribution.
+        :type prob: th.Tensor
+        :param mask: An optional boolean tensor for masking specific probabilities
+            before creating the distribution. Defaults to None.
+        :type mask: th.Tensor, optional
+        :return: A `MaskedCategorical` distribution object created with reshaped
+            probabilities and combined masking information.
+        :rtype: MaskedCategorical
+        """
         probs = prob.reshape(*self.nvec.shape, self.prob_last_dim)
         start = self.start
         mask = mask if mask is not None else th.ones_like(probs, dtype=th.bool, device=probs.device)

pyproject.toml

@@ -7,7 +7,7 @@ authors = [
 ]
 license = { text = "MIT" }
 readme = "README.md"
-requires-python = ">=3.10,<4"
+requires-python = ">=3.10"
 
 keywords = [
     "python"