From 6513099806d3bd9661d667925d6b4b7a78bf8ba5 Mon Sep 17 00:00:00 2001
From: pseudo-rnd-thoughts <mark.m.towers@gmail.com>
Date: Mon, 11 Mar 2024 12:10:41 +0000
Subject: [PATCH] Update the documentation

---
 docs/api/spaces.md           |  3 ++-
 docs/api/spaces/composite.md |  5 +++++
 gymnasium/spaces/oneof.py    | 13 +++++++++++--
 3 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/docs/api/spaces.md b/docs/api/spaces.md
index 9f749a66a..014daec15 100644
--- a/docs/api/spaces.md
+++ b/docs/api/spaces.md
@@ -66,7 +66,8 @@ Often environment spaces require joining fundamental spaces together for vectori
 * :class:`Dict` - Supports a dictionary of keys and subspaces, used for a fixed number of unordered spaces
 * :class:`Tuple` - Supports a tuple of subspaces, used for multiple for a fixed number of ordered spaces
 * :class:`Sequence` - Supports a variable number of instances of a single subspace, used for entities spaces or selecting a variable number of actions
-* :py:class:`Graph` - Supports graph based actions or observations with discrete or continuous nodes and edge values.
+* :class:`Graph` - Supports graph based actions or observations with discrete or continuous nodes and edge values
+* :class:`OneOf` - Supports optional action spaces such that an action can be one of N possible subspaces
 ```
 
 ## Utility functions
diff --git a/docs/api/spaces/composite.md b/docs/api/spaces/composite.md
index b43a5b0f5..72d8aee27 100644
--- a/docs/api/spaces/composite.md
+++ b/docs/api/spaces/composite.md
@@ -21,4 +21,9 @@
 
     .. automethod:: gymnasium.spaces.Graph.sample
     .. automethod:: gymnasium.spaces.Graph.seed
+
+.. autoclass:: gymnasium.spaces.OneOf
+
+    .. automethod:: gymnasium.spaces.OneOf.sample
+    .. automethod:: gymnasium.spaces.OneOf.seed
 ```
diff --git a/gymnasium/spaces/oneof.py b/gymnasium/spaces/oneof.py
index 097372518..d88f0b130 100644
--- a/gymnasium/spaces/oneof.py
+++ b/gymnasium/spaces/oneof.py
@@ -18,8 +18,16 @@ class OneOf(Space[Any]):
     Example:
         >>> from gymnasium.spaces import OneOf, Box, Discrete
         >>> observation_space = OneOf((Discrete(2), Box(-1, 1, shape=(2,))), seed=42)
-        >>> observation_space.sample()
+        >>> observation_space.sample()  # the first element is the space index (Box in this case) and the second element is the sample from Box
         (1, array([-0.3991573 ,  0.21649833], dtype=float32))
+        >>> observation_space.sample()  # this time the Discrete space was sampled as index=0
+        (0, 0)
+        >>> observation_space[0]
+        Discrete(2)
+        >>> observation_space[1]
+        Box(-1.0, 1.0, (2,), float32)
+        >>> len(observation_space)
+        2
     """
 
     def __init__(
@@ -27,7 +35,7 @@ def __init__(
         spaces: Iterable[Space[Any]],
         seed: int | typing.Sequence[int] | np.random.Generator | None = None,
     ):
-        r"""Constructor of :class:`Tuple` space.
+        r"""Constructor of :class:`OneOf` space.
 
         The generated instance will represent the cartesian product :math:`\text{spaces}[0] \times ... \times \text{spaces}[-1]`.
 
@@ -36,6 +44,7 @@ def __init__(
             seed: Optionally, you can use this argument to seed the RNGs of the ``spaces`` to ensure reproducible sampling.
         """
         self.spaces = tuple(spaces)
+        assert len(self.spaces) > 0, "Empty `OneOf` spaces are not supported."
         for space in self.spaces:
             assert isinstance(
                 space, Space