SRT-H 0.0.1__tar.gz

srt_h-0.0.1/.github/workflows/python-publish.yml

@@ -0,0 +1,36 @@
+# 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 }}

srt_h-0.0.1/.github/workflows/test.yml

@@ -0,0 +1,21 @@
+name: Pytest
+on: [push, pull_request]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.10"
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install -e .[test]
+    - name: Test with pytest
+      run: |
+        python -m pytest tests/

srt_h-0.0.1/.gitignore

@@ -0,0 +1,207 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$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
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.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
+#poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+#pdm.lock
+#pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+#pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi
+
+# 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
+.envrc
+.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/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/

srt_h-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 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.

srt_h-0.0.1/PKG-INFO

@@ -0,0 +1,103 @@
+Metadata-Version: 2.4
+Name: SRT-H
+Version: 0.0.1
+Summary: MMDiT
+Project-URL: Homepage, https://pypi.org/project/SRT-H/
+Project-URL: Repository, https://github.com/lucidrains/SRT-H
+Author-email: Phil Wang <lucidrains@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 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,attention mechanism,deep learning,surgical-robotics,transformers
+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: bidirectional-cross-attention>=0.1.0
+Requires-Dist: einops>=0.8.0
+Requires-Dist: einx>=0.3.0
+Requires-Dist: requests
+Requires-Dist: torch>=2.0
+Requires-Dist: validators
+Requires-Dist: vit-pytorch>=1.11.4
+Requires-Dist: x-transformers>=2.5.6
+Provides-Extra: examples
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Description-Content-Type: text/markdown
+
+<img src="./fig2.png" width="450px"></img>
+
+## SRT-H (wip)
+
+Implementation of the model architecture for [SRT-H](https://h-surgical-robot-transformer.github.io/) out of Johns Hopkins
+
+## Citations
+
+```bibtex
+@misc{kim2025srthhierarchicalframeworkautonomous,
+    title   = {SRT-H: A Hierarchical Framework for Autonomous Surgery via Language Conditioned Imitation Learning}, 
+    author  = {Ji Woong Kim and Juo-Tung Chen and Pascal Hansen and Lucy X. Shi and Antony Goldenberg and Samuel Schmidgall and Paul Maria Scheikl and Anton Deguet and Brandon M. White and De Ru Tsai and Richard Cha and Jeffrey Jopling and Chelsea Finn and Axel Krieger},
+    year    = {2025},
+    eprint  = {2505.10251},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2505.10251}, 
+}
+```
+
+```bibtex
+@article{Kim2024SurgicalRT,
+    title   = {Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks},
+    author  = {Ji Woong Kim and Tony Zhao and Samuel Schmidgall and Anton Deguet and Marin Kobilarov and Chelsea Finn and Axel Krieger},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2407.12998},
+    url     = {https://api.semanticscholar.org/CorpusID:271270134}
+}
+```
+
+```bibtex
+@misc{zhao2023learningfinegrainedbimanualmanipulation,
+    title   = {Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware}, 
+    author  = {Tony Z. Zhao and Vikash Kumar and Sergey Levine and Chelsea Finn},
+    year    = {2023},
+    eprint  = {2304.13705},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2304.13705}, 
+}
+```
+
+```bibtex
+@misc{heng2025vitacformerlearningcrossmodalrepresentation,
+    title   = {ViTacFormer: Learning Cross-Modal Representation for Visuo-Tactile Dexterous Manipulation}, 
+    author  = {Liang Heng and Haoran Geng and Kaifeng Zhang and Pieter Abbeel and Jitendra Malik},
+    year    = {2025},
+    eprint  = {2506.15953},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2506.15953}, 
+}
+```

srt_h-0.0.1/README.md

@@ -0,0 +1,54 @@
+<img src="./fig2.png" width="450px"></img>
+
+## SRT-H (wip)
+
+Implementation of the model architecture for [SRT-H](https://h-surgical-robot-transformer.github.io/) out of Johns Hopkins
+
+## Citations
+
+```bibtex
+@misc{kim2025srthhierarchicalframeworkautonomous,
+    title   = {SRT-H: A Hierarchical Framework for Autonomous Surgery via Language Conditioned Imitation Learning}, 
+    author  = {Ji Woong Kim and Juo-Tung Chen and Pascal Hansen and Lucy X. Shi and Antony Goldenberg and Samuel Schmidgall and Paul Maria Scheikl and Anton Deguet and Brandon M. White and De Ru Tsai and Richard Cha and Jeffrey Jopling and Chelsea Finn and Axel Krieger},
+    year    = {2025},
+    eprint  = {2505.10251},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2505.10251}, 
+}
+```
+
+```bibtex
+@article{Kim2024SurgicalRT,
+    title   = {Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks},
+    author  = {Ji Woong Kim and Tony Zhao and Samuel Schmidgall and Anton Deguet and Marin Kobilarov and Chelsea Finn and Axel Krieger},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2407.12998},
+    url     = {https://api.semanticscholar.org/CorpusID:271270134}
+}
+```
+
+```bibtex
+@misc{zhao2023learningfinegrainedbimanualmanipulation,
+    title   = {Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware}, 
+    author  = {Tony Z. Zhao and Vikash Kumar and Sergey Levine and Chelsea Finn},
+    year    = {2023},
+    eprint  = {2304.13705},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2304.13705}, 
+}
+```
+
+```bibtex
+@misc{heng2025vitacformerlearningcrossmodalrepresentation,
+    title   = {ViTacFormer: Learning Cross-Modal Representation for Visuo-Tactile Dexterous Manipulation}, 
+    author  = {Liang Heng and Haoran Geng and Kaifeng Zhang and Pieter Abbeel and Jitendra Malik},
+    year    = {2025},
+    eprint  = {2506.15953},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2506.15953}, 
+}
+```

srt_h-0.0.1/SRT_H/SRT_H.py

@@ -0,0 +1,354 @@
+from __future__ import annotations
+from collections import namedtuple
+
+import torch
+from torch import nn, cat, stack, Tensor, is_tensor
+from torch.nn import Module, ModuleList, Parameter, Identity, Linear, Sequential
+
+from x_transformers import Encoder, Attention, AttentionPool
+
+from einops import rearrange, repeat
+from einops.layers.torch import Rearrange
+
+from vit_pytorch.accept_video_wrapper import AcceptVideoWrapper
+
+from bidirectional_cross_attention import BidirectionalCrossAttentionTransformer as BiCrossAttnTransformer
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# pretrained model related
+# they successfully apply
+
+# 1. efficient net for low level vision
+# 2. swin t for high level vision
+# 3. distilbert for clinician language feedback
+
+class AcceptVideoSwin(Module):
+    def __init__(
+        self,
+        hub_url = 'SharanSMenon/swin-transformer-hub',
+        model_name = 'swin_tiny_patch4_window7_224',
+        dim_model = 768,
+        max_time_seq_len = 8 # say 8 frames
+    ):
+        super().__init__()
+        swin = torch.hub.load(hub_url, model_name, pretrained = True)
+        swin.avgpool = Identity()
+        swin.head = Rearrange('b (d n) -> b n d', d = dim_model)
+
+        self.model = AcceptVideoWrapper(
+            swin,
+            add_time_pos_emb = True,
+            time_seq_len = max_time_seq_len,
+            dim_emb = dim_model
+        )
+
+    def forward(
+        self,
+        video
+    ):
+        embeds = self.model(video)
+        return rearrange(embeds, 'b t n d -> b (t n) d')
+
+# ACT - Action Chunking Transformer - Zhou et al.
+
+Losses = namedtuple('Losses', ('action_recon', 'vae_kl_div'))
+
+class ACT(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        dim_joint_state,
+        action_chunk_len,
+        dim_action = 20,
+        dim_head = 64,
+        dim_style_vector = None,
+        dim_lang_condition = None,
+        heads = 8,
+        vae_encoder_depth = 3,
+        encoder_depth = 6,
+        decoder_depth = 6,
+        vae_encoder_kwargs: dict = dict(),
+        vae_encoder_attn_pool_depth = 2,
+        encoder_kwargs: dict = dict(),
+        decoder: dict = dict(),
+        image_model: Module | None = None,
+        image_model_dim_emb = None,
+        dim_tactile_input = None,
+        tactile_self_attn_depth = 2,
+        tactile_image_fusion_cross_attn_depth = 2, # ViTacFormer
+        max_num_image_frames = 32,
+        vae_kl_loss_weight = 1.,
+        action_loss_fn = nn.L1Loss()
+    ):
+        super().__init__()
+
+        self.dim = dim
+
+        # style vector dimension related
+
+        dim_style_vector = default(dim_style_vector, dim)
+        need_style_proj = dim_style_vector != dim
+
+        self.dim_style_vector = dim_style_vector
+
+        # projections
+
+        self.joint_to_token = nn.Linear(dim_joint_state, dim)
+        self.action_to_vae_tokens = nn.Linear(dim_action, dim)
+
+        # for the cvae and style vector
+
+        self.vae_encoder = Encoder(
+            dim = dim,
+            depth = vae_encoder_depth,
+            heads = heads,
+            attn_dim_head = dim_head,
+            use_rmsnorm = True
+        )
+
+        self.attn_pooler = AttentionPool(dim = dim, depth = vae_encoder_attn_pool_depth, heads = heads, dim_head = dim_head)
+
+        self.to_style_vector_mean_log_variance = Sequential(
+            Linear(dim, dim_style_vector * 2, bias = False),
+            Rearrange('... (d mean_log_var) -> mean_log_var ... d', mean_log_var = 2)
+        )
+
+        self.style_vector_to_token = nn.Linear(dim_style_vector, dim) if need_style_proj else nn.Identity()
+
+        # detr like
+
+        self.encoder = Encoder(
+            dim = dim,
+            depth = vae_encoder_depth,
+            heads = heads,
+            attn_dim_head = dim_head,
+            use_rmsnorm = True
+        )
+
+        self.action_queries = Parameter(torch.randn(action_chunk_len, dim) * 1e-2)
+
+        self.decoder = Encoder(
+            dim = dim,
+            depth = vae_encoder_depth,
+            heads = heads,
+            attn_dim_head = dim_head,
+            cross_attend = True,
+            use_rmsnorm = True,
+            rotary_pos_emb = True
+        )
+
+        self.decoder_embed_to_actions = nn.Linear(dim, dim_action)
+
+        # image model
+
+        image_model_dim_emb = default(image_model_dim_emb, dim)
+        need_image_to_state_proj = image_model_dim_emb != dim
+
+        # they used efficient net in the paper, but allow for others
+
+        if not exists(image_model):
+            efficientnet = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_efficientnet_b0', pretrained = True)
+            utils = torch.hub.load('NVIDIA/DeepLearningExamples:torchhub', 'nvidia_convnets_processing_utils')
+
+            image_model_dim_emb = 1280
+            efficientnet.classifier = Rearrange('b d h w -> b (h w) d') # replace the classifier layer in efficient net
+
+            image_model = efficientnet
+
+        # set the image model and the projection to image tokens (state tokens)
+
+        self.image_model = image_model
+        self.to_state_tokens = nn.Linear(image_model_dim_emb, dim) if exists(image_model) and need_image_to_state_proj else nn.Identity()
+
+        if exists(image_model):
+            self.accept_video_wrapper = AcceptVideoWrapper(image_model, add_time_pos_emb = True, time_seq_len = max_num_image_frames, dim_emb = image_model_dim_emb)
+
+        # tactile
+
+        self.to_tactile_tokens = nn.Linear(dim_tactile_input, dim) if exists(dim_tactile_input) else None
+
+        self.tactile_self_attn = Encoder(
+            dim = dim,
+            depth = tactile_self_attn_depth,
+            heads = heads,
+            attn_dim_head = dim_head,
+            pre_norm_has_final_norm = False
+        )
+
+        self.tactile_fuse = BiCrossAttnTransformer(
+            dim = dim,
+            context_dim = dim,
+            heads = heads,
+            depth = tactile_image_fusion_cross_attn_depth
+        )
+
+        # take care of clinician feedback which is conditioning the state tokens with FiLM
+
+        self.to_film_scale_offset = None
+
+        if exists(dim_lang_condition):
+            self.to_film_scale_offset = nn.Linear(dim_lang_condition, dim * 2, bias = False)
+            nn.init.zeros_(self.film.weight)
+
+        # loss related
+
+        self.action_loss_fn = action_loss_fn
+        self.vae_kl_loss_weight = vae_kl_loss_weight
+
+    def forward(
+        self,
+        *,
+        joint_state,                 # (d)
+        video = None,                # (b c t h w)
+        state_tokens = None,         # (b n d)
+        tactile_input = None,        # (b nt dt)
+        tactile_tokens = None,       # (b nt d)
+        actions = None,              # (b na da)
+        style_vector = None,         # (d) | (b d)
+        lang_condition = None,       # (b d)
+        return_loss_breakdown = False
+    ):
+
+        # take care of video -> image tokens
+
+        assert exists(state_tokens) or exists(video), '`video` or its encoded `state_tokens` must be passed in'
+        assert not (exists(video) and not exists(self.image_model)), '`video` cannot be passed in if `image_model` is not set'
+
+        if exists(video):
+            assert video.ndim == 5
+
+            images_embeds = self.accept_video_wrapper(video, eval_with_no_grad = True)
+            state_tokens = self.to_state_tokens(images_embeds)
+
+            state_tokens = rearrange(state_tokens, 'b t n d -> b (t n) d')
+
+        # if tactile tokens are presented, fuse it with cross attention, as proposed by ViTacFormer - force feedback is becoming a thing
+
+        if exists(tactile_input):
+            assert not exists(tactile_tokens) and exists(self.to_tactile_tokens)
+
+            tactile_tokens = self.to_tactile_tokens(tactile_input)
+
+        if exists(tactile_tokens):
+            tactile_tokens = self.tactile_self_attn(tactile_tokens)
+
+            state_tokens, tactile_tokens = self.tactile_fuse(state_tokens, tactile_tokens)
+
+        # maybe condition state tokens
+
+        if exists(lang_condition):
+            assert exists(self.to_film_scale_offset), f'`dim_lang_condition` must be set if doing further conditioning (clinician feedback in this paper)'
+
+            scale, offset = self.to_film_scale_offset(lang_condition).chunk(2, dim = -1)
+
+            scale, offset = tuple(rearrange(t, 'b d -> b 1 d'), (scale, offset))
+
+            state_tokens = state_tokens * (scale + 1.) + offset
+
+        # variables
+
+        batch, device = state_tokens.shape[0], state_tokens.device
+
+        is_training = exists(actions)
+        is_sampling = not is_training
+
+        assert not (is_training and exists(style_vector)), 'style vector z cannot be set during training'
+
+        # joint token
+
+        joint_tokens = self.joint_to_token(joint_state)
+        joint_tokens = rearrange(joint_tokens, 'b d -> b 1 d')
+
+        # take care of the needed style token during training
+
+        if is_training:
+            action_vae_tokens = self.action_to_vae_tokens(actions)
+
+            vae_input = cat((action_vae_tokens, joint_tokens), dim = 1)
+
+            vae_encoder_embed = self.vae_encoder(vae_input)
+
+            # cross attention pool
+
+            pooled_vae_embed = self.attn_pooler(vae_encoder_embed)
+
+            style_mean, style_log_variance = self.to_style_vector_mean_log_variance(pooled_vae_embed)
+
+            # reparam
+
+            style_std = (0.5 * style_log_variance).exp()
+
+            noise = torch.randn_like(style_mean)
+
+            style_vector = style_mean + style_std * noise
+
+        elif exists(style_vector) and style_vector.ndim == 1:
+
+            style_vector = repeat(style_vector, 'd -> b 1 d', b = batch)
+
+        else:
+            # or just zeros during inference, as in the paper
+
+            style_vector = torch.zeros((batch, 1, self.dim_style_vector), device = device)
+
+        style_token = self.style_vector_to_token(style_vector)
+
+        # detr like encoder / decoder
+
+        encoder_input = cat((style_token, state_tokens, joint_tokens), dim = 1)
+
+        encoded = self.encoder(encoder_input)
+
+        decoder_input = repeat(self.action_queries, 'na d -> b na d', b = batch)
+
+        decoded = self.decoder(decoder_input, context = encoded)
+
+        pred_actions = self.decoder_embed_to_actions(decoded)
+
+        if not is_training:
+            return pred_actions
+
+        # take care of training loss
+
+        action_recon_loss = self.action_loss_fn(pred_actions, actions)
+
+        vae_kl_loss = (0.5 * (
+            style_log_variance.exp()
+            + style_mean.square()
+            - style_log_variance
+            - 1.
+        )).sum(dim = -1).mean()
+
+        loss_breakdown = Losses(action_recon_loss, vae_kl_loss)
+
+        total_loss = (
+            action_recon_loss +
+            vae_kl_loss * self.vae_kl_loss_weight
+        )
+
+        if not return_loss_breakdown:
+            return total_loss
+
+        return total_loss, loss_breakdown
+
+# classes
+
+class SRT(Module):
+    def __init__(
+        self
+    ):
+        super().__init__()
+
+    def forward(
+        self,
+        state
+    ):
+        raise NotImplementedError

srt_h-0.0.1/SRT_H/__init__.py

@@ -0,0 +1,4 @@
+from SRT_H.SRT_H import (
+    ACT,
+    SRT
+)

srt_h-0.0.1/pyproject.toml

@@ -0,0 +1,65 @@
+[project]
+name = "SRT-H"
+version = "0.0.1"
+description = "MMDiT"
+authors = [
+    { name = "Phil Wang", email = "lucidrains@gmail.com" }
+]
+readme = "README.md"
+requires-python = ">= 3.9"
+license = { file = "LICENSE" }
+keywords = [
+    'artificial intelligence',
+    'deep learning',
+    'surgical-robotics',
+    'transformers',
+    'attention mechanism'
+]
+
+classifiers=[
+    'Development Status :: 4 - Beta',
+    'Intended Audience :: Developers',
+    'Topic :: Scientific/Engineering :: Artificial Intelligence',
+    'License :: OSI Approved :: MIT License',
+    'Programming Language :: Python :: 3.9',
+]
+
+dependencies = [
+    "bidirectional-cross-attention>=0.1.0",
+    "einx>=0.3.0",
+    "einops>=0.8.0",
+    "requests",
+    "torch>=2.0",
+    "validators",
+    "vit-pytorch>=1.11.4",
+    "x-transformers>=2.5.6",
+]
+
+[project.urls]
+Homepage = "https://pypi.org/project/SRT-H/"
+Repository = "https://github.com/lucidrains/SRT-H"
+
+[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 = ["SRT_H"]

srt_h-0.0.1/tests/test_model.py

@@ -0,0 +1,94 @@
+import torch
+
+import pytest
+param = pytest.mark.parametrize
+
+@param('pass_custom_style', (False, True))
+def test_act(
+    pass_custom_style
+):
+    from SRT_H.SRT_H import ACT
+
+    act = ACT(
+        dim = 512,
+        dim_joint_state = 17,
+        action_chunk_len = 16
+    )
+
+    states = torch.randn(3, 512, 512)
+    joint_state = torch.randn(3, 17)
+
+    actions = torch.randn(3, 16, 20)
+
+    loss = act(
+        state_tokens = states,
+        joint_state = joint_state,
+        actions = actions
+    )
+
+    loss.backward()
+
+    # after a lot of data and training ...
+
+    style_vector = torch.ones(512) if pass_custom_style else None
+
+    sampled_actions = act(state_tokens = states, joint_state = joint_state, style_vector = style_vector) # (3, 16, 20)
+
+    assert sampled_actions.shape == (3, 16, 20)
+
+@param('tactile', (False, True))
+@param('efficient_net', (False, True))
+def test_act_with_image_model(
+    tactile,
+    efficient_net
+):
+
+    from SRT_H.SRT_H import ACT
+
+    from vit_pytorch import ViT
+    from vit_pytorch.extractor import Extractor
+
+    v = ViT(
+        image_size = 256,
+        patch_size = 32,
+        num_classes = 1000,
+        dim = 1024,
+        depth = 6,
+        heads = 16,
+        mlp_dim = 2048,
+        dropout = 0.1,
+        emb_dropout = 0.1
+    )
+
+    v = Extractor(v, return_embeddings_only = True)
+
+    act = ACT(
+        image_model = v if not efficient_net else None,
+        image_model_dim_emb = 1024,
+        dim = 512,
+        dim_joint_state = 17,
+        action_chunk_len = 16,
+        dim_tactile_input = 37
+    )
+
+    states = torch.randn(3, 512, 512)
+    joint_state = torch.randn(3, 17)
+
+    tactile_input = torch.randn(3, 16, 37) if tactile else None
+
+    actions = torch.randn(3, 16, 20)
+
+    video = torch.randn(3, 3, 2, 224, 224)
+
+    loss = act(
+        video = video,
+        joint_state = joint_state,
+        tactile_input = tactile_input,
+        actions = actions
+    )
+
+    loss.backward()
+
+    # after a lot of data and training ...
+
+    sampled_actions = act(state_tokens = states, joint_state = joint_state) # (3, 16, 20)