SAC-pytorch 0.0.2__tar.gz

sac_pytorch-0.0.2/.github/workflows/python-publish.yml

@@ -0,0 +1,38 @@
+
+  
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

sac_pytorch-0.0.2/.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

sac_pytorch-0.0.2/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Phil Wang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sac_pytorch-0.0.2/PKG-INFO

@@ -0,0 +1,138 @@
+Metadata-Version: 2.3
+Name: SAC-pytorch
+Version: 0.0.2
+Summary: Soft Actor Critic - Pytorch
+Project-URL: Homepage, https://pypi.org/project/SAC-pytorch/
+Project-URL: Repository, https://github.com/lucidrains/SAC-pytorch
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2024 Phil Wang
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: artificial intelligence,deep learning,reinforcement learning,soft actor critic
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: adam-atan2-pytorch>=0.0.15
+Requires-Dist: beartype
+Requires-Dist: einops>=0.8.0
+Requires-Dist: einx[torch]>=0.1.3
+Requires-Dist: ema-pytorch>=0.6.2
+Requires-Dist: jaxtyping
+Requires-Dist: pytorch-custom-utils>=0.0.18
+Requires-Dist: soft-moe-pytorch>=0.1.6
+Requires-Dist: torch>=2.0
+Provides-Extra: examples
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Description-Content-Type: text/markdown
+
+## SAC (Soft Actor Critic) - Pytorch (wip)
+
+Implementation of Soft Actor Critic and some of its improvements in Pytorch. Interest comes from watching <a href="https://www.youtube.com/watch?v=17NrtKHdPDw">this lecture</a>
+
+## Citations
+
+```bibtex
+@article{Haarnoja2018SoftAA,
+    title   = {Soft Actor-Critic Algorithms and Applications},
+    author  = {Tuomas Haarnoja and Aurick Zhou and Kristian Hartikainen and G. Tucker and Sehoon Ha and Jie Tan and Vikash Kumar and Henry Zhu and Abhishek Gupta and P. Abbeel and Sergey Levine},
+    journal = {ArXiv},
+    year    = {2018},
+    volume  = {abs/1812.05905},
+    url     = {https://api.semanticscholar.org/CorpusID:55703664}
+}
+```
+
+```bibtex
+@article{Hiraoka2021DropoutQF,
+    title   = {Dropout Q-Functions for Doubly Efficient Reinforcement Learning},
+    author  = {Takuya Hiraoka and Takahisa Imagawa and Taisei Hashimoto and Takashi Onishi and Yoshimasa Tsuruoka},
+    journal = {ArXiv},
+    year    = {2021},
+    volume  = {abs/2110.02034},
+    url     = {https://api.semanticscholar.org/CorpusID:238353966}
+}
+```
+
+```bibtex
+@inproceedings{ObandoCeron2024MixturesOE,
+    title   = {Mixtures of Experts Unlock Parameter Scaling for Deep RL},
+    author  = {Johan S. Obando-Ceron and Ghada Sokar and Timon Willi and Clare Lyle and Jesse Farebrother and Jakob Foerster and Gintare Karolina Dziugaite and Doina Precup and Pablo Samuel Castro},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:267637059}
+}
+```
+
+```bibtex
+@inproceedings{Kumar2023MaintainingPI,
+    title   = {Maintaining Plasticity in Continual Learning via Regenerative Regularization},
+    author  = {Saurabh Kumar and Henrik Marklund and Benjamin Van Roy},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:261076021}
+}
+```
+
+```bibtex
+@inproceedings{Kuznetsov2020ControllingOB,
+    title   = {Controlling Overestimation Bias with Truncated Mixture of Continuous Distributional Quantile Critics},
+    author  = {Arsenii Kuznetsov and Pavel Shvechikov and Alexander Grishin and Dmitry P. Vetrov},
+    booktitle = {International Conference on Machine Learning},
+    year    = {2020},
+    url     = {https://api.semanticscholar.org/CorpusID:218581840}
+}
+```
+
+```bibtex
+@article{Zagoruyko2017DiracNetsTV,
+    title   = {DiracNets: Training Very Deep Neural Networks Without Skip-Connections},
+    author={Sergey Zagoruyko and Nikos Komodakis},
+    journal = {ArXiv},
+    year    = {2017},
+    volume  = {abs/1706.00388},
+    url     = {https://api.semanticscholar.org/CorpusID:1086822}
+}
+```
+
+```bibtex
+@article{Abbas2023LossOP,
+    title  = {Loss of Plasticity in Continual Deep Reinforcement Learning},
+    author = {Zaheer Abbas and Rosie Zhao and Joseph Modayil and Adam White and Marlos C. Machado},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2303.07507},
+    url     = {https://api.semanticscholar.org/CorpusID:257504763}
+}
+```
+
+```bibtex
+@article{Nauman2024BiggerRO,
+    title   = {Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous control},
+    author  = {Michal Nauman and Mateusz Ostaszewski and Krzysztof Jankowski and Piotr Milo's and Marek Cygan},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2405.16158},
+    url     = {https://api.semanticscholar.org/CorpusID:270063045}
+}
+```

sac_pytorch-0.0.2/README.md

@@ -0,0 +1,88 @@
+## SAC (Soft Actor Critic) - Pytorch (wip)
+
+Implementation of Soft Actor Critic and some of its improvements in Pytorch. Interest comes from watching <a href="https://www.youtube.com/watch?v=17NrtKHdPDw">this lecture</a>
+
+## Citations
+
+```bibtex
+@article{Haarnoja2018SoftAA,
+    title   = {Soft Actor-Critic Algorithms and Applications},
+    author  = {Tuomas Haarnoja and Aurick Zhou and Kristian Hartikainen and G. Tucker and Sehoon Ha and Jie Tan and Vikash Kumar and Henry Zhu and Abhishek Gupta and P. Abbeel and Sergey Levine},
+    journal = {ArXiv},
+    year    = {2018},
+    volume  = {abs/1812.05905},
+    url     = {https://api.semanticscholar.org/CorpusID:55703664}
+}
+```
+
+```bibtex
+@article{Hiraoka2021DropoutQF,
+    title   = {Dropout Q-Functions for Doubly Efficient Reinforcement Learning},
+    author  = {Takuya Hiraoka and Takahisa Imagawa and Taisei Hashimoto and Takashi Onishi and Yoshimasa Tsuruoka},
+    journal = {ArXiv},
+    year    = {2021},
+    volume  = {abs/2110.02034},
+    url     = {https://api.semanticscholar.org/CorpusID:238353966}
+}
+```
+
+```bibtex
+@inproceedings{ObandoCeron2024MixturesOE,
+    title   = {Mixtures of Experts Unlock Parameter Scaling for Deep RL},
+    author  = {Johan S. Obando-Ceron and Ghada Sokar and Timon Willi and Clare Lyle and Jesse Farebrother and Jakob Foerster and Gintare Karolina Dziugaite and Doina Precup and Pablo Samuel Castro},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:267637059}
+}
+```
+
+```bibtex
+@inproceedings{Kumar2023MaintainingPI,
+    title   = {Maintaining Plasticity in Continual Learning via Regenerative Regularization},
+    author  = {Saurabh Kumar and Henrik Marklund and Benjamin Van Roy},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:261076021}
+}
+```
+
+```bibtex
+@inproceedings{Kuznetsov2020ControllingOB,
+    title   = {Controlling Overestimation Bias with Truncated Mixture of Continuous Distributional Quantile Critics},
+    author  = {Arsenii Kuznetsov and Pavel Shvechikov and Alexander Grishin and Dmitry P. Vetrov},
+    booktitle = {International Conference on Machine Learning},
+    year    = {2020},
+    url     = {https://api.semanticscholar.org/CorpusID:218581840}
+}
+```
+
+```bibtex
+@article{Zagoruyko2017DiracNetsTV,
+    title   = {DiracNets: Training Very Deep Neural Networks Without Skip-Connections},
+    author={Sergey Zagoruyko and Nikos Komodakis},
+    journal = {ArXiv},
+    year    = {2017},
+    volume  = {abs/1706.00388},
+    url     = {https://api.semanticscholar.org/CorpusID:1086822}
+}
+```
+
+```bibtex
+@article{Abbas2023LossOP,
+    title  = {Loss of Plasticity in Continual Deep Reinforcement Learning},
+    author = {Zaheer Abbas and Rosie Zhao and Joseph Modayil and Adam White and Marlos C. Machado},
+    journal = {ArXiv},
+    year    = {2023},
+    volume  = {abs/2303.07507},
+    url     = {https://api.semanticscholar.org/CorpusID:257504763}
+}
+```
+
+```bibtex
+@article{Nauman2024BiggerRO,
+    title   = {Bigger, Regularized, Optimistic: scaling for compute and sample-efficient continuous control},
+    author  = {Michal Nauman and Mateusz Ostaszewski and Krzysztof Jankowski and Piotr Milo's and Marek Cygan},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2405.16158},
+    url     = {https://api.semanticscholar.org/CorpusID:270063045}
+}
+```

sac_pytorch-0.0.2/SAC_pytorch/SAC.py

@@ -0,0 +1,805 @@
+from __future__ import annotations
+
+import math
+from collections import namedtuple
+
+import torch
+import torch.nn.functional as F
+from torch.distributions import Normal
+from torch import nn, einsum, Tensor, tensor
+from torch.nn import Module, ModuleList, Sequential
+
+from adam_atan2_pytorch.adam_atan2_with_wasserstein_reg import AdamAtan2 as Adam
+
+# tensor typing
+
+import jaxtyping
+from beartype import beartype
+from beartype.door import is_bearable
+
+class TorchTyping:
+    def __init__(self, abstract_dtype):
+        self.abstract_dtype = abstract_dtype
+
+    def __getitem__(self, shapes: str):
+        return self.abstract_dtype[Tensor, shapes]
+
+Float = TorchTyping(jaxtyping.Float)
+Int   = TorchTyping(jaxtyping.Int)
+Bool  = TorchTyping(jaxtyping.Bool)
+
+# ein notations
+# b - batch
+# n - number of actions
+# nc - number of continuous actions
+# nd - number of discrete actions
+# c - critics
+# q - quantiles
+
+from einx import get_at
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+# EMA for target networks
+
+from ema_pytorch import EMA
+
+# constants
+
+ContinuousOutput = namedtuple('ContinuousOutput', [
+    'mu',
+    'sigma'
+])
+
+SoftActorOutput = namedtuple('SoftActorOutput', [
+    'continuous',
+    'discrete'
+])
+
+SampledSoftActorOutput = namedtuple('SampledSoftActorOutput', [
+    'continuous',
+    'continuous_log_prob',
+    'discrete',
+    'discrete_action_logits',
+])
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def compact(arr):
+    return [*filter(exists, arr)]
+
+def identity(t):
+    return t
+
+def cast_tuple(t, length = 1):
+    return t if isinstance(t, tuple) else ((t,) * length)
+
+# tensor helpers
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def entropy(t, eps = 1e-20):
+    prob = t.softmax(dim = -1)
+    return (-prob * log(prob, eps = eps)).sum(dim = -1)
+
+def gumbel_noise(t):
+    noise = torch.rand_like(t)
+    return -log(-log(noise))
+
+def gumbel_sample(t, temperature = 1., dim = -1):
+    assert temperature > 0.
+    return ((t / temperature) + gumbel_noise(t)).argmax(dim = dim)
+
+def init_(m):
+    if not isinstance(m, nn.Linear):
+        return
+
+    gain = torch.nn.init.calculate_gain('relu')
+    torch.nn.init.orthogonal_(m.weight, gain)
+
+    if not exists(m.bias):
+        return
+
+    torch.nn.init.zeros_(m.bias)
+
+# helper classes
+
+def Sequential(*modules):
+    return nn.Sequential(*filter(exists, modules))
+
+class Residual(Module):
+    @beartype
+    def __init__(
+        self,
+        fn: Module,
+        dim_in,
+        dim_out
+    ):
+        super().__init__()
+        self.fn = fn
+        self.residual_proj = nn.Linear(dim_in, dim_out) if dim_in != dim_out else identity
+
+    def forward(self, x, **kwargs):
+        return self.fn(x, **kwargs) + self.residual_proj(x)
+
+# "bro" mlp
+
+class BroMLP(Module):
+    @beartype
+    def __init__(
+        self,
+        dim,
+        dim_out,
+        dim_hidden = None,
+        depth = 3,
+        dropout = 0.,
+        activation = nn.ReLU,
+        expansion_factor = 2,
+        final_norm = False
+    ):
+        super().__init__()
+        """
+        following the design of BroNet https://arxiv.org/abs/2405.16158v1
+        """
+
+        dim_hidden = default(dim_hidden, dim * 2)
+
+        layers = []
+
+        self.proj_in = nn.Sequential(
+            nn.Linear(dim, dim_hidden),
+            activation()
+        )
+
+        dim_inner = dim_hidden * expansion_factor
+
+        for _ in range(depth):
+
+            layer = Sequential(
+                nn.Linear(dim_hidden, dim_inner),
+                nn.Dropout(dropout),
+                nn.LayerNorm(dim_inner, bias = False),
+                activation(),
+                nn.Linear(dim_inner, dim_hidden),
+                nn.LayerNorm(dim_hidden, bias = False),
+            )
+
+            nn.init.constant_(layer[-1].weight, 1e-5)
+            layers.append(layer)
+
+        # final layer out
+
+        self.layers = ModuleList(layers)
+
+        self.final_norm = nn.LayerNorm(dim_hidden) if final_norm else nn.Identity()
+
+        self.proj_out = nn.Linear(dim_hidden, dim_out)
+
+        # init
+
+        self.apply(init_)
+
+    def forward(self, x):
+
+        x = self.proj_in(x)
+
+        for layer in self.layers:
+            x = layer(x) + x
+
+        x = self.final_norm(x)
+        return self.proj_out(x)
+
+# main modules
+
+class Actor(Module):
+    @beartype
+    def __init__(
+        self,
+        *,
+        dim_state,
+        num_cont_actions,
+        mlp_depth = 3,
+        num_discrete_actions: tuple[int, ...] = (),
+        dim_hidden = None,
+        eps = 1e-5
+    ):
+        super().__init__()
+        self.eps = eps
+
+        discrete_action_dims = sum(num_discrete_actions)
+        cont_action_dims = num_cont_actions * 2
+
+        self.num_cont_actions = num_cont_actions
+        self.num_discrete_actions = num_discrete_actions
+        self.split_dims = (discrete_action_dims, cont_action_dims)
+
+        self.to_actions = BroMLP(
+            dim_state,
+            depth = mlp_depth,
+            dim_hidden = dim_hidden,
+            dim_out = discrete_action_dims + cont_action_dims
+        )
+
+
+    def forward(
+        self,
+        state: Float['b ...'],
+        sample = False,
+        cont_reparametrize = False,
+        discrete_sample_temperature = 1.,
+        discrete_sample_deterministic = False,
+    ) -> (
+        SoftActorOutput |
+        SampledSoftActorOutput
+    ):
+
+        action_dims = self.to_actions(state)
+
+        discrete_actions, cont_actions = action_dims.split(self.split_dims, dim = -1)
+        discrete_action_logits = discrete_actions.split(self.num_discrete_actions, dim = -1)
+
+        mu, sigma = rearrange(cont_actions, '... (n mu_sigma) -> mu_sigma ... n', mu_sigma = 2)
+        sigma = sigma.sigmoid().clamp(min = self.eps)
+
+        cont_output = ContinuousOutput(mu, sigma)
+        discrete_output = discrete_action_logits
+
+        if not sample:
+            return SoftActorOutput(cont_output, discrete_output)
+
+        # handle continuous
+
+        if cont_reparametrize:
+            sampled_cont_actions = mu + sigma * torch.randn_like(sigma)
+        else:
+            sampled_cont_actions = torch.normal(mu, sigma)
+
+        squashed_cont_actions = sampled_cont_actions.tanh() # tanh squashing
+
+        cont_log_prob = Normal(mu, sigma).log_prob(sampled_cont_actions)
+        cont_log_prob = cont_log_prob - log(1. - squashed_cont_actions ** 2, eps = self.eps)
+
+        scaled_squashed_cont_actions = squashed_cont_actions * self.num_cont_actions
+
+        # handle discrete
+
+        sampled_discrete_actions = []
+        discrete_log_probs = []
+
+        for logits in discrete_action_logits:
+
+            if discrete_sample_deterministic:
+                sampled_action = logits.argmax(dim = -1)
+            else:
+                sampled_action = gumbel_sample(logits, temperature = discrete_sample_temperature)
+
+            sampled_discrete_actions.append(sampled_action)
+
+            log_probs = logits.log_softmax(dim = -1)
+            discrete_log_prob = get_at('... [d], ... -> ...', log_probs, sampled_action)
+            discrete_log_probs.append(discrete_log_prob)
+
+        # return all sampled continuous and discrete actions with their associated log prob
+
+        return SampledSoftActorOutput(
+            scaled_squashed_cont_actions,
+            torch.stack(sampled_discrete_actions, dim = -1),
+            cont_log_prob,
+            discrete_action_logits
+        )
+
+class Critic(Module):
+    @beartype
+    def __init__(
+        self,
+        *,
+        dim_state,
+        num_cont_actions,
+        mlp_depth = 3,
+        num_discrete_actions: tuple[int, ...] = (),
+        dim_hidden = None,
+        layernorm = False,
+        dropout = 0.,
+        num_quantiles: int | None = None,
+        quantiles: tuple[float, ...] | None = None
+    ):
+        super().__init__()
+        assert not exists(num_quantiles) or num_quantiles > 0
+
+        use_quantiles = exists(num_quantiles)
+        self.returning_quantiles = use_quantiles
+
+        num_actions_split = tensor((num_cont_actions, *num_discrete_actions))
+
+        # determine the output dimension of the critic
+        # which is the sum of all the actions (continous and discrete), multiplied by the number of quantiles
+
+        critic_dim_out = num_actions_split
+        if use_quantiles:
+            critic_dim_out = critic_dim_out * num_quantiles
+
+        self.to_values = BroMLP(
+            dim_state + num_cont_actions,
+            depth = mlp_depth,
+            dim_out = critic_dim_out.sum().item(),
+            dim_hidden = dim_hidden,
+            dropout = dropout
+        )
+
+        # save the number of quantiles and the number of actions, for splitting out the output of the critic correctly
+
+        self.num_quantiles = num_quantiles
+        self.num_actions_split = num_actions_split.tolist()
+
+        # for tensor typing
+
+        self._n = num_cont_actions
+        self._q = num_quantiles
+
+        # quantiles
+
+        if exists(quantiles):
+            quantiles = tensor(quantiles)
+        elif exists(num_quantiles):
+            quantiles = torch.linspace(0., 1., num_quantiles + 2)[1:-1] # excluding 0 and 1 - say 3 quantiles will be 0.25, 0.5, 0.75
+
+        self.register_buffer('quantiles', quantiles)
+
+    def forward(
+        self,
+        state: Float['b ...'],
+        cont_actions: Float['b {self._n}'] | None = None
+    ) -> tuple[Float['b ...'], ...]:
+
+        pack_input = compact([state, cont_actions])
+
+        mlp_input, _ = pack(pack_input, 'b *')
+
+        values = self.to_values(mlp_input)
+
+        if self.returning_quantiles:
+            values = rearrange(values, '... (n q) -> ... n q', q = self.num_quantiles)
+
+        split_dim = -2 if self.returning_quantiles else -1
+
+        values = values.split(self.num_actions_split, dim = split_dim)
+
+        return values
+
+class MultipleCritics(Module):
+    @beartype
+    def __init__(
+        self,
+        *critics: Critic,
+        use_softmin = False
+    ):
+        super().__init__()
+        assert len(critics) > 0
+        assert all([not critic.returning_quantiles for critic in critics]), 'this wrapper only allows for non-quantile critics'
+
+        self.num_critics = len(critics)
+        self.critics = ModuleList(critics)
+
+        self.use_softmin = use_softmin
+
+    def forward(
+        self,
+        states: Float['b ...'],
+        cont_actions: Float['b nc'] | None = None,
+        discrete_actions: Int['b nd'] | None = None,
+        target_values: Float['b n'] | None = None,
+        return_breakdown = False,
+        **kwargs
+    ):
+        critics_values = [critic(states, cont_actions) for critic in self.critics]
+
+        critics_values = [torch.stack(one_value_for_critics) for one_value_for_critics in zip(*critics_values)]
+
+        if not exists(target_values):
+
+            min_critic_values = []
+
+            for critic_values in critics_values:
+                if self.use_softmin:
+                    values = torch.stack(values, dim = -1)
+                    softmin = (-values).softmax(dim = -1)
+
+                    min_critic_value = (softmin * critic_values).sum(dim = -1)
+                else:
+                    min_critic_value = torch.minimum(*critic_values)
+
+                min_critic_values.append(min_critic_value)
+
+            if not return_breakdown:
+                return min_critic_values
+
+            return min_critic_values, values
+
+        cont_critics_values, *discrete_critics_values = critics_values
+
+        discrete_critics_values = [get_at('c b [l], b -> c b', discrete_critics_value, discrete_action) for discrete_critics_value, discrete_action in zip(discrete_critics_values, discrete_actions.unbind(dim = -1))]
+
+        values, _ = pack([cont_critics_values, *discrete_critics_values], 'c b *')
+
+        target_values = repeat(target_values, '... -> c ...', c = self.num_critics)
+        losses = F.mse_loss(values, target_values, reduction = 'none')
+
+        reduced_losses = reduce(losses, 'c b n -> b', 'sum').mean() # sum losses per action per critic, and average over batch
+
+        if not return_breakdown:
+            return reduced_losses
+
+        return reduced_losses, losses
+
+class MultipleQuantileCritics(Module):
+    @beartype
+    def __init__(
+        self,
+        *critics: Critic,
+        frac_atom_keep = 0.75 # will truncate 25% of the top values
+    ):
+        super().__init__()
+        assert len(critics) > 0
+        assert all([critic.returning_quantiles for critic in critics]), 'all critics must be returning quantiles'
+        assert len(set([critic.num_quantiles for critic in critics])) == 1, 'all critics must have same number of quantiles'
+
+        self.num_critics = len(critics)
+
+        self.critics = ModuleList(critics)
+        self.num_atom_keep = int(frac_atom_keep * self.num_critics * self.num_quantiles)
+
+    @property
+    def quantiles(self):
+        return self.critics[0].quantiles
+
+    @property
+    def num_quantiles(self):
+        return self.critics[0].num_quantiles
+
+    def forward(
+        self,
+        states: Float['b ...'],
+        cont_actions: Float['b nc'] | None = None,
+        discrete_actions: Int['b nd'] | None = None,
+        target_values: Float['b n q'] | None = None,
+        truncate_quantiles_across_critics = False,
+        return_breakdown = False
+    ):
+        critics_quantile_atoms = [critic(states, cont_actions) for critic in self.critics]
+
+        critics_quantile_atoms = [torch.stack(one_value_for_critics) for one_value_for_critics in zip(*critics_quantile_atoms)]
+
+        if not exists(target_values):
+
+            if truncate_quantiles_across_critics:
+                truncated_means_per_action = []
+
+                for critic_quantile_atoms in critics_quantile_atoms:
+                    quantile_atoms = rearrange(critic_quantile_atoms, 'c b ... q -> b ... (q c)')
+
+                    # mean of the truncated distribution
+
+                    truncated_quantiles = quantile_atoms.topk(self.num_atom_keep, largest = False).values
+                    truncated_mean = truncated_quantiles.mean(dim = -1)
+
+                    truncated_means_per_action.append(truncated_mean)
+
+                return_value = truncated_means_per_action
+            else:
+
+                min_critic_value_per_action = []
+
+                for critic_quantile_atoms in critics_quantile_atoms:
+                    min_critic_value = torch.minimum(*critic_quantile_atoms)
+                    min_critic_value_per_action.append(min_critic_value)
+
+                return_value = min_critic_value_per_action
+
+            if not return_breakdown:
+                return return_value
+
+            return return_value, critics_quantile_atoms
+
+        cont_quantile_atoms, *discrete_quantile_atoms = critics_quantile_atoms
+        discrete_quantile_atoms = [get_at('c b [l] q, b -> c b q', discrete_quantile_atom, discrete_action) for discrete_quantile_atom, discrete_action in zip(discrete_quantile_atoms, discrete_actions.unbind(dim = -1))]
+
+        quantile_atoms, _ = pack([cont_quantile_atoms, *discrete_quantile_atoms], 'c b * q')
+
+        # quantile regression if training
+
+        quantiles = self.quantiles
+
+        # quantile regression loss
+
+        error = target_values - quantile_atoms
+        losses = torch.maximum(error * quantiles, error * (quantiles - 1.))
+
+        reduced_losses = reduce(losses, 'c b n q -> b', 'sum').mean()
+
+        if not return_breakdown:
+            return reduced_losses
+
+        return reduced_losses, losses
+
+# automatic entropy temperature adjustment
+# will account for both continuous and discrete
+
+class LearnedEntropyTemperature(Module):
+    def __init__(
+        self,
+        num_discrete_actions = 0,
+        num_cont_actions = 0
+    ):
+        super().__init__()
+
+        self.log_alpha = nn.Parameter(tensor(0.))
+
+        self.has_discrete = len(num_discrete_actions) > 0
+        self.has_continuous = num_cont_actions > 0
+
+        self.discrete_entropy_targets = [0.98 * math.log(one_num_discrete_actions)for one_num_discrete_actions in num_discrete_actions] 
+        self.continuous_entropy_target = num_cont_actions
+
+    @property
+    def alpha(self):
+        return self.log_alpha.exp()
+
+    def forward(
+        self,
+        cont_log_prob: Float['b nc'] | None = None,
+        discrete_logits: tuple[Float['b _'], ...] | None = None,
+        return_breakdown = False
+    ):
+        assert exists(cont_log_prob) or exists(discrete_log_prob)
+
+        alpha = self.alpha
+
+        losses = []
+
+        if exists(discrete_logits):
+
+            for one_discrete_logits, discrete_entropy_target in zip(discrete_logits, self.discrete_entropy_targets):
+                discrete_log_prob = one_discrete_logits.log_softmax(dim = -1)
+                discrete_entropy_temp_loss = -alpha * (discrete_log_prob + discrete_entropy_target).detach()
+
+                losses.append(discrete_entropy_temp_loss.mean())
+
+        if exists(cont_log_prob):
+            cont_entropy_temp_loss = -alpha * (cont_log_prob + self.continuous_entropy_target).detach()
+
+            cont_entropy_temp_loss = cont_entropy_temp_loss.mean(dim = 0)
+            losses.append(cont_entropy_temp_loss)
+
+        reduced_losses, _ = pack(losses, '*')
+        reduced_losses = reduced_losses.sum()
+
+        if not return_breakdown:
+            return reduced_losses
+
+        return reduced_losses, losses
+
+# main class
+
+class SAC(Module):
+    @beartype
+    def __init__(
+        self,
+        actor: dict | Actor,
+        critics: (
+            list[dict] |
+            list[Critic] |
+            MultipleCritics |
+            MultipleQuantileCritics
+        ),
+        quantiled_critics = False,
+        reward_discount_rate = 0.99,
+        reward_scale = 1.,
+        actor_learning_rate = 3e-4,
+        actor_regen_reg_rate = 1e-4,
+        critic_target_ema_decay = 0.99,
+        critics_learning_rate = 3e-4,
+        critics_regen_reg_rate = 1e-4,
+        temperature_learning_rate = 3e-4,
+        ema_kwargs: dict = dict()
+    ):
+        super().__init__()
+
+        # set actor
+
+        if isinstance(actor, dict):
+            actor = Actor(**actor)
+
+        self.actor = actor
+
+        self.actor_optimizer = Adam(
+            actor.parameters(),
+            lr = actor_learning_rate,
+            regen_reg_rate = actor_regen_reg_rate
+        )
+        # based on the actor hyperparameters, init the leraned temperature container
+
+        self.learned_entropy_temperature = LearnedEntropyTemperature(
+            num_cont_actions = actor.num_cont_actions,
+            num_discrete_actions = actor.num_discrete_actions
+        )
+
+        self.temperature_optimizer = Adam(
+            self.learned_entropy_temperature.parameters(),
+            lr = temperature_learning_rate,
+        )
+
+        # set critics
+        # allow for usual double+ critic trick or truncated quantiled critics
+
+        if is_bearable(critics, list[dict]):
+            critics = [Critic(**critic) for critic in critics]
+
+        multiple_critic_klass = MultipleCritics if not quantiled_critics else MultipleQuantileCritics
+
+        if is_bearable(critics, list[Critic]):
+            critics = multiple_critic_klass(*critics)
+
+        assert isinstance(critics, multiple_critic_klass), f'expected {multiple_critic_klass.__name__} but received critics wrapped with {type(critics).__name__}'
+
+        self.critics = critics
+
+        self.quantiled_critics = quantiled_critics
+
+        # critic optimizers
+
+        self.critics_optimizer = Adam(
+            critics.parameters(),
+            lr = critics_learning_rate,
+            regen_reg_rate = critics_regen_reg_rate
+        )
+
+        # target critic network
+
+        self.critics_target = EMA(
+            critics,
+            beta = critic_target_ema_decay,
+            include_online_model = False,
+            **ema_kwargs
+        )
+
+        # reward related
+
+        self.reward_scale = reward_scale
+        self.reward_discount_rate = reward_discount_rate
+
+    def forward(
+        self,
+        states: Float['b ...'],
+        cont_actions: Float['b nc'],
+        discrete_actions: Int['b nd'],
+        rewards: Float['b'],
+        done: Bool['b'],
+        next_states: Float['b ...']
+    ):
+
+        rewards = rewards * self.reward_scale
+
+        # bellman equation
+        # todo: setup n-step
+
+        entropy_temp = self.learned_entropy_temperature.alpha.detach()
+
+        γ = self.reward_discount_rate
+        not_terminal = (~done).float()
+
+        with torch.no_grad():
+            self.critics_target.eval()
+
+            next_cont_q_value, *next_discrete_q_values = self.critics_target(next_states, cont_actions = cont_actions)
+
+            # outputs from actor
+
+            actor_output = self.actor(states, sample = True)
+
+            cont_log_prob = actor_output.continuous_log_prob
+            discrete_logits = actor_output.discrete_action_logits
+
+            # learned temperature
+
+            learned_entropy_weight = self.learned_entropy_temperature.alpha
+
+            next_soft_state_values = []
+
+            # first handle continuous soft state value
+
+            if exists(cont_log_prob):
+                if self.quantiled_critics:
+                    cont_log_prob = rearrange(cont_log_prob, '... -> ... 1')
+    
+                cont_soft_state_value = next_cont_q_value - learned_entropy_weight * cont_log_prob
+                next_soft_state_values.append(cont_soft_state_value)
+
+            # then handle discrete contribution
+
+            if len(next_discrete_q_values) > 0:
+
+                for next_discrete_q_value, discrete_logit in zip(next_discrete_q_values, discrete_logits):
+                    discrete_prob = discrete_logit.softmax(dim = -1)
+                    discrete_log_prob = log(discrete_prob)
+
+                    if self.quantiled_critics:
+                        discrete_prob = rearrange(discrete_prob, '... -> ... 1')
+                        discrete_log_prob = rearrange(discrete_log_prob, '... -> ... 1')
+
+                    discrete_soft_state_value = (discrete_prob * (next_discrete_q_value - learned_entropy_weight * discrete_log_prob)).sum(dim = 1)
+                    next_soft_state_values.append(discrete_soft_state_value)
+
+        if self.quantiled_critics:
+            next_soft_state_values, _ = pack(next_soft_state_values, 'b * q')
+        else:
+            next_soft_state_values, _ = pack(next_soft_state_values, 'b *')
+
+        target_q_values = rewards + not_terminal * γ * next_soft_state_values
+
+        critics_losses = self.critics(
+            states,
+            cont_actions = cont_actions,
+            discrete_actions = discrete_actions,
+            target_values = target_q_values
+        )
+
+        # update the critics
+
+        critics_losses.backward()
+        self.critics_optimizer.step()
+        self.critics_optimizer.zero_grad()
+
+        # update the actor
+
+        actor_output = self.actor(states, sample = True, cont_reparametrize = True)
+
+        cont_log_prob = actor_output.continuous_log_prob
+        discrete_logits = actor_output.discrete_action_logits
+
+        cont_q_values, *discrete_q_values = self.critics(
+            states,
+            cont_actions = actor_output.continuous,
+            discrete_actions = actor_output.discrete,
+            truncate_quantiles_across_critics = True
+        )
+
+        actor_action_losses = []
+
+        if exists(cont_log_prob):
+            cont_action_loss = (entropy_temp * cont_log_prob - cont_q_values).sum(dim = -1).mean()
+
+            actor_action_losses.append(cont_action_loss)
+
+        for discrete_logit, one_discrete_q_value in zip(discrete_logits, discrete_q_values):
+            discrete_prob = discrete_logit.softmax(dim = -1)
+            discrete_log_prob = log(discrete_prob)
+
+            one_discrete_actor_loss = (discrete_prob * (entropy_temp * discrete_log_prob - one_discrete_q_value)).mean()
+
+            actor_action_losses.append(one_discrete_actor_loss)
+
+        sum(actor_action_losses).backward()
+        self.actor_optimizer.step()
+        self.actor_optimizer.zero_grad()
+
+        # update the learned entropy temperature
+
+        temperature_loss = self.learned_entropy_temperature(
+            cont_log_prob = cont_log_prob,
+            discrete_logits = discrete_logits
+        )
+
+        temperature_loss.backward()
+        self.temperature_optimizer.step()
+        self.temperature_optimizer.zero_grad()
+
+        # update ema of all critics
+
+        self.critics_target.update()

sac_pytorch-0.0.2/SAC_pytorch/__init__.py

@@ -0,0 +1,4 @@
+from SAC_pytorch.SAC import (
+    SAC,
+    BroMLP
+)

sac_pytorch-0.0.2/pyproject.toml

@@ -0,0 +1,65 @@
+[project]
+name = "SAC-pytorch"
+version = "0.0.2"
+description = "Soft Actor Critic - Pytorch"
+authors = [
+    { name = "Phil Wang", email = "lucidrains@gmail.com" }
+]
+readme = "README.md"
+requires-python = ">= 3.9"
+license = { file = "LICENSE" }
+keywords = [
+    'artificial intelligence',
+    'deep learning',
+    'reinforcement learning',
+    'soft actor critic'
+]
+
+classifiers=[
+    'Development Status :: 4 - Beta',
+    'Intended Audience :: Developers',
+    'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    'License :: OSI Approved :: MIT License',
+    'Programming Language :: Python :: 3.9',
+]
+
+dependencies = [
+    'adam-atan2-pytorch>=0.0.15',
+    'beartype',
+    'einops>=0.8.0',
+    'einx[torch]>=0.1.3',
+    'ema-pytorch>=0.6.2',
+    'jaxtyping',
+    'pytorch-custom-utils>=0.0.18',
+    'soft-moe-pytorch>=0.1.6',
+    'torch>=2.0'
+]
+
+[project.urls]
+Homepage = "https://pypi.org/project/SAC-pytorch/"
+Repository = "https://github.com/lucidrains/SAC-pytorch"
+
+[project.optional-dependencies]
+examples = []
+test = [
+    "pytest"
+]
+
+[tool.pytest.ini_options]
+pythonpath = [
+  "."
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.rye]
+managed = true
+dev-dependencies = []
+
+[tool.hatch.metadata]
+allow-direct-references = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["SAC_pytorch"]