Merge pull request #140 from Toni-SM/develop

Develop
Toni-SM · Feb 13, 2024 · 52a75d3 · 52a75d3
2 parents 7b090a1 + b350717
commit 52a75d3
Show file tree

Hide file tree

Showing 116 changed files with 2,638 additions and 378 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,16 @@
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
+## [1.1.0] - 2024-02-12
+### Added
+- MultiCategorical mixin to operate MultiDiscrete action spaces
+
+### Changed (breaking changes)
+- Rename the `ManualTrainer` to `StepTrainer`
+- Output training/evaluation progress messages to system's stdout
+- Get single observation/action spaces for vectorized environments
+- Update Isaac Orbit environment wrapper
+
 ## [1.0.0] - 2023-08-16
 
 Transition from pre-release versions (`1.0.0-rc.1` and`1.0.0-rc.2`) to a stable version.

diff --git a/docs/source/_static/imgs/model_categorical_cnn-dark.svg b/docs/source/_static/imgs/model_categorical_cnn-dark.svg
diff --git a/docs/source/_static/imgs/model_categorical_cnn-light.svg b/docs/source/_static/imgs/model_categorical_cnn-light.svg
diff --git a/docs/source/_static/imgs/model_deterministic_cnn-dark.svg b/docs/source/_static/imgs/model_deterministic_cnn-dark.svg
diff --git a/docs/source/_static/imgs/model_deterministic_cnn-light.svg b/docs/source/_static/imgs/model_deterministic_cnn-light.svg
diff --git a/docs/source/_static/imgs/model_gaussian_cnn-dark.svg b/docs/source/_static/imgs/model_gaussian_cnn-dark.svg
diff --git a/docs/source/_static/imgs/model_gaussian_cnn-light.svg b/docs/source/_static/imgs/model_gaussian_cnn-light.svg
diff --git a/docs/source/_static/imgs/model_multicategorical-dark.svg b/docs/source/_static/imgs/model_multicategorical-dark.svg
diff --git a/docs/source/_static/imgs/model_multicategorical-light.svg b/docs/source/_static/imgs/model_multicategorical-light.svg
diff --git a/docs/source/api/agents/a2c.rst b/docs/source/api/agents/a2c.rst
@@ -164,6 +164,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\blacksquare`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
@@ -194,7 +197,10 @@ The implementation uses 1 stochastic (discrete or continuous) and 1 deterministi
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Categorical <models_categorical>` / :ref:`Gaussian <models_gaussian>` / :ref:`MultivariateGaussian <models_multivariate_gaussian>`
+      - :ref:`Categorical <models_categorical>` /
+        |br| :ref:`Multi-Categorical <models_multicategorical>` /
+        |br| :ref:`Gaussian <models_gaussian>` /
+        |br| :ref:`MultivariateGaussian <models_multivariate_gaussian>`
     * - :math:`V_\phi(s)`
       - Value
       - :literal:`"value"`

diff --git a/docs/source/api/agents/amp.rst b/docs/source/api/agents/amp.rst
@@ -162,6 +162,10 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
@@ -194,7 +198,8 @@ The implementation uses 1 stochastic (continuous) and 2 deterministic function a
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Gaussian <models_gaussian>` / :ref:`MultivariateGaussian <models_multivariate_gaussian>`
+      - :ref:`Gaussian <models_gaussian>` /
+        |br| :ref:`MultivariateGaussian <models_multivariate_gaussian>`
     * - :math:`V_\phi(s)`
       - Value
       - :literal:`"value"`

diff --git a/docs/source/api/agents/cem.rst b/docs/source/api/agents/cem.rst
@@ -119,6 +119,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\blacksquare`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\square`
@@ -149,7 +152,8 @@ The implementation uses 1 discrete function approximator. This function approxim
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Categorical <models_categorical>`
+      - :ref:`Categorical <models_categorical>` /
+        |br| :ref:`Multi-Categorical <models_multicategorical>`
 
 .. raw:: html
 

diff --git a/docs/source/api/agents/ddpg.rst b/docs/source/api/agents/ddpg.rst
@@ -159,6 +159,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`

diff --git a/docs/source/api/agents/ddqn.rst b/docs/source/api/agents/ddqn.rst
@@ -119,6 +119,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\square`

diff --git a/docs/source/api/agents/dqn.rst b/docs/source/api/agents/dqn.rst
@@ -119,6 +119,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\square`

diff --git a/docs/source/api/agents/ppo.rst b/docs/source/api/agents/ppo.rst
@@ -180,6 +180,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\blacksquare`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
@@ -210,7 +213,10 @@ The implementation uses 1 stochastic (discrete or continuous) and 1 deterministi
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Categorical <models_categorical>` / :ref:`Gaussian <models_gaussian>` / :ref:`MultivariateGaussian <models_multivariate_gaussian>`
+      - :ref:`Categorical <models_categorical>` /
+        |br| :ref:`Multi-Categorical <models_multicategorical>` /
+        |br| :ref:`Gaussian <models_gaussian>` /
+        |br| :ref:`MultivariateGaussian <models_multivariate_gaussian>`
     * - :math:`V_\phi(s)`
       - Value
       - :literal:`"value"`

diff --git a/docs/source/api/agents/q_learning.rst b/docs/source/api/agents/q_learning.rst
@@ -99,6 +99,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`

diff --git a/docs/source/api/agents/rpo.rst b/docs/source/api/agents/rpo.rst
@@ -219,6 +219,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
@@ -249,7 +252,8 @@ The implementation uses 1 continuous stochastic and 1 deterministic function app
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Gaussian <models_gaussian>` / :ref:`MultivariateGaussian <models_multivariate_gaussian>`
+      - :ref:`Gaussian <models_gaussian>` /
+        |br| :ref:`MultivariateGaussian <models_multivariate_gaussian>`
     * - :math:`V_\phi(s)`
       - Value
       - :literal:`"value"`

diff --git a/docs/source/api/agents/sac.rst b/docs/source/api/agents/sac.rst
@@ -160,6 +160,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
@@ -190,7 +193,8 @@ The implementation uses 1 stochastic and 4 deterministic function approximators.
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Gaussian <models_gaussian>` / :ref:`MultivariateGaussian <models_multivariate_gaussian>`
+      - :ref:`Gaussian <models_gaussian>` /
+        |br| :ref:`MultivariateGaussian <models_multivariate_gaussian>`
     * - :math:`Q_{\phi 1}(s, a)`
       - Q1-network (critic 1)
       - :literal:`"critic_1"`

diff --git a/docs/source/api/agents/sarsa.rst b/docs/source/api/agents/sarsa.rst
@@ -99,6 +99,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`

diff --git a/docs/source/api/agents/td3.rst b/docs/source/api/agents/td3.rst
@@ -169,6 +169,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`

diff --git a/docs/source/api/agents/trpo.rst b/docs/source/api/agents/trpo.rst
@@ -216,6 +216,9 @@ The implementation supports the following `Gym spaces <https://www.gymlibrary.de
     * - Discrete
       - .. centered:: :math:`\square`
       - .. centered:: :math:`\square`
+    * - MultiDiscrete
+      - .. centered:: :math:`\square`
+      - .. centered:: :math:`\square`
     * - Box
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
@@ -246,7 +249,8 @@ The implementation uses 1 stochastic and 1 deterministic function approximator.
       - :literal:`"policy"`
       - observation
       - action
-      - :ref:`Gaussian <models_gaussian>` / :ref:`MultivariateGaussian <models_multivariate_gaussian>`
+      - :ref:`Gaussian <models_gaussian>` /
+        |br| :ref:`MultivariateGaussian <models_multivariate_gaussian>`
     * - :math:`V_\phi(s)`
       - Value
       - :literal:`"value"`

diff --git a/docs/source/api/models.rst b/docs/source/api/models.rst
@@ -6,6 +6,7 @@ Models
 
     Tabular <models/tabular>
     Categorical <models/categorical>
+    Multi-Categorical <models/multicategorical>
     Gaussian <models/gaussian>
     Multivariate Gaussian <models/multivariate_gaussian>
     Deterministic <models/deterministic>
@@ -29,6 +30,9 @@ Models (or agent models) refer to a representation of the agent's policy, value
     * - :doc:`Categorical model <models/categorical>` (discrete domain)
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`
+    * - :doc:`Multi-Categorical model <models/multicategorical>` (discrete domain)
+      - .. centered:: :math:`\blacksquare`
+      - .. centered:: :math:`\square`
     * - :doc:`Gaussian model <models/gaussian>` (continuous domain)
       - .. centered:: :math:`\blacksquare`
       - .. centered:: :math:`\blacksquare`

diff --git a/docs/source/api/models/categorical.rst b/docs/source/api/models/categorical.rst
@@ -163,6 +163,24 @@ Usage
                             :start-after: [start-cnn-functional-torch]
                             :end-before: [end-cnn-functional-torch]
 
+            .. group-tab:: |_4| |jax| |_4|
+
+                .. tabs::
+
+                    .. group-tab:: setup-style
+
+                        .. literalinclude:: ../../snippets/categorical_model.py
+                            :language: python
+                            :start-after: [start-cnn-setup-jax]
+                            :end-before: [end-cnn-setup-jax]
+
+                    .. group-tab:: compact-style
+
+                        .. literalinclude:: ../../snippets/categorical_model.py
+                            :language: python
+                            :start-after: [start-cnn-compact-jax]
+                            :end-before: [end-cnn-compact-jax]
+
     .. tab:: RNN
 
         .. image:: ../../_static/imgs/model_categorical_rnn-light.svg

diff --git a/docs/source/api/models/deterministic.rst b/docs/source/api/models/deterministic.rst
@@ -163,6 +163,24 @@ Usage
                             :start-after: [start-cnn-functional-torch]
                             :end-before: [end-cnn-functional-torch]
 
+            .. group-tab:: |_4| |jax| |_4|
+
+                .. tabs::
+
+                    .. group-tab:: setup-style
+
+                        .. literalinclude:: ../../snippets/deterministic_model.py
+                            :language: python
+                            :start-after: [start-cnn-setup-jax]
+                            :end-before: [end-cnn-setup-jax]
+
+                    .. group-tab:: compact-style
+
+                        .. literalinclude:: ../../snippets/deterministic_model.py
+                            :language: python
+                            :start-after: [start-cnn-compact-jax]
+                            :end-before: [end-cnn-compact-jax]
+
     .. tab:: RNN
 
         .. image:: ../../_static/imgs/model_deterministic_rnn-light.svg

diff --git a/docs/source/api/models/gaussian.rst b/docs/source/api/models/gaussian.rst
@@ -163,6 +163,24 @@ Usage
                             :start-after: [start-cnn-functional-torch]
                             :end-before: [end-cnn-functional-torch]
 
+            .. group-tab:: |_4| |jax| |_4|
+
+                .. tabs::
+
+                    .. group-tab:: setup-style
+
+                        .. literalinclude:: ../../snippets/gaussian_model.py
+                            :language: python
+                            :start-after: [start-cnn-setup-jax]
+                            :end-before: [end-cnn-setup-jax]
+
+                    .. group-tab:: compact-style
+
+                        .. literalinclude:: ../../snippets/gaussian_model.py
+                            :language: python
+                            :start-after: [start-cnn-compact-jax]
+                            :end-before: [end-cnn-compact-jax]
+
     .. tab:: RNN
 
         .. image:: ../../_static/imgs/model_gaussian_rnn-light.svg