Merge branch 'main' into feat/add_sliding_tile_puzzle_environment

.gitignore

@@ -70,6 +70,8 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+# MkDocs documentation
+docs_public/
 
 # PyBuilder
 .pybuilder/

README.md

@@ -1,7 +1,9 @@
 <p align="center">
-    <a href="docs/img/jumanji_logo.png">
-        <img src="docs/img/jumanji_logo.png" alt="Jumanji logo" width="50%"/>
-    </a>
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="docs/img/jumanji_logo_dm.png">
+  <source media="(prefers-color-scheme: light)" srcset="docs/img/jumanji_logo.png">
+  <img alt="Jumanji Logo" src="docs/img/jumanji_logo.png", width="50%">
+</picture>
 </p>
 
 [![Python Versions](https://img.shields.io/pypi/pyversions/jumanji.svg?style=flat-square)](https://www.python.org/doc/versions/)

docs/api/environments/tetris.md

@@ -0,0 +1,8 @@
+::: jumanji.environments.packing.tetris.env.Tetris
+    selection:
+      members:
+        - __init__
+        - reset
+        - step
+        - observation_spec
+        - action_spec

jumanji/environments/routing/robot_warehouse/viewer.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""File adapted from [Rware](https://github.com/semitable/robotic-warehouse). More specifically,
+the rendering code is copied from the original Rware environment and should be ignored from the
+copyright."""
+
 # flake8: noqa: CCR001
 
 from typing import Callable, Optional, Sequence, Tuple

jumanji/specs_test.py

@@ -184,13 +184,13 @@ def test_shape_element_type_error(self) -> None:
 
     def test_dtype_type_error(self) -> None:
         with pytest.raises(TypeError):
-            specs.Array((1, 2, 3), "32")
+            specs.Array((1, 2, 3), "32")  # type: ignore
 
     def test_scalar_shape(self) -> None:
         specs.Array((), jnp.int32)
 
     def test_string_dtype_error(self) -> None:
-        specs.Array((1, 2, 3), "int32")
+        specs.Array((1, 2, 3), "int32")  # type: ignore
 
     def test_dtype(self) -> None:
         specs.Array((1, 2, 3), int)

jumanji/types.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import field
 from typing import TYPE_CHECKING, Dict, Generic, Optional, Sequence, TypeVar, Union
 
 if TYPE_CHECKING:  # https://github.com/python/mypy/issues/6239
@@ -71,14 +72,14 @@ class TimeStep(Generic[Observation]):
         extras: environment metric(s) or information returned by the environment but
             not observed by the agent (hence not in the observation). For example, it
             could be whether an invalid action was taken. In most environments, extras
-            is None.
+            is an empty dictionary.
     """
 
     step_type: StepType
     reward: Array
     discount: Array
     observation: Observation
-    extras: Optional[Dict] = None
+    extras: Dict = field(default_factory=dict)
 
     def first(self) -> Array:
         return self.step_type == StepType.FIRST
@@ -110,6 +111,7 @@ def restart(
     Returns:
         TimeStep identified as a reset.
     """
+    extras = extras or {}
     return TimeStep(
         step_type=StepType.FIRST,
         reward=jnp.zeros(shape, dtype=float),
@@ -144,6 +146,7 @@ def transition(
         TimeStep identified as a transition.
     """
     discount = discount if discount is not None else jnp.ones(shape, dtype=float)
+    extras = extras or {}
     return TimeStep(
         step_type=StepType.MID,
         reward=reward,
@@ -175,6 +178,7 @@ def termination(
     Returns:
         TimeStep identified as the termination of an episode.
     """
+    extras = extras or {}
     return TimeStep(
         step_type=StepType.LAST,
         reward=reward,
@@ -208,6 +212,7 @@ def truncation(
         TimeStep identified as the truncation of an episode.
     """
     discount = discount if discount is not None else jnp.ones(shape, dtype=float)
+    extras = extras or {}
     return TimeStep(
         step_type=StepType.LAST,
         reward=reward,
@@ -228,4 +233,4 @@ def get_valid_dtype(dtype: Union[jnp.dtype, type]) -> jnp.dtype:
     Returns:
         dtype converted to the correct type precision.
     """
-    return jnp.empty((), dtype).dtype
+    return jnp.empty((), dtype).dtype  # type: ignore

jumanji/wrappers.py

@@ -96,6 +96,14 @@ def action_spec(self) -> specs.Spec:
         """Returns the action spec."""
         return self._env.action_spec()
 
+    def reward_spec(self) -> specs.Array:
+        """Returns the reward spec."""
+        return self._env.reward_spec()
+
+    def discount_spec(self) -> specs.BoundedArray:
+        """Returns the discount spec."""
+        return self._env.discount_spec()
+
     def render(self, state: State) -> Any:
         """Compute render frames during initialisation of the environment.
 

requirements/requirements.txt

@@ -2,7 +2,7 @@ chex>=0.1.3
 dm-env>=1.5
 gym>=0.22.0
 jax>=0.2.26
-matplotlib>=3.3.4
+matplotlib~=3.7.4
 numpy>=1.19.5
 Pillow>=9.0.0
 typing-extensions>=4.0.0