[README]

README.md

@@ -6,3 +6,31 @@
 An open-source non-official community implementation of the model from the paper: Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks: https://surgical-robot-transformer.github.io/
 
 
+## Installation
+
+```bash
+pip3 install srt-torch
+```
+
+
+## Usage
+
+```
+
+```
+
+
+## Citation
+
+
+```bibtex
+@misc{kim2024surgicalrobottransformersrt,
+    title={Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks}, 
+    author={Ji Woong Kim and Tony Z. Zhao and Samuel Schmidgall and Anton Deguet and Marin Kobilarov and Chelsea Finn and Axel Krieger},
+    year={2024},
+    eprint={2407.12998},
+    archivePrefix={arXiv},
+    primaryClass={cs.RO},
+    url={https://arxiv.org/abs/2407.12998}, 
+}
+```

example.py

@@ -4,7 +4,7 @@
 
 import torch
 from loguru import logger
-from srt.main import (
+from srt_torch.main import (
     SurgicalRobotTransformer,
     ModelConfig,
     RobotObservation,

pyproject.toml

@@ -3,15 +3,15 @@ requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
-name = "paper"
+name = "srt-torch"
 version = "0.0.1"
-description = "Paper - Pytorch"
+description = "Surgical Robot Transformer (SRT): Imitation Learning for Surgical Tasks"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]
-homepage = "https://github.com/kyegomez/paper"
-documentation = "https://github.com/kyegomez/paper"  # Add this if you have documentation.
+homepage = "https://github.com/Agora-Lab-AI/SRT"
+documentation = "https://github.com/Agora-Lab-AI/SRT"  # Add this if you have documentation.
 readme = "README.md"  # Assuming you have a README.md
-repository = "https://github.com/kyegomez/paper"
+repository = "https://github.com/Agora-Lab-AI/SRT"
 keywords = ["artificial intelligence", "deep learning", "optimizers", "Prompt Engineering"]
 classifiers = [
     "Development Status :: 4 - Beta",
@@ -23,12 +23,8 @@ classifiers = [
 
 [tool.poetry.dependencies]
 python = "^3.10"
-swarms = "*"
-zetascale = "*"
-
-[tool.poetry.dev-dependencies]
-# Add development dependencies here
-
+loguru = "*"
+torch = "*"
 
 [tool.poetry.group.lint.dependencies]
 ruff = "^0.1.6"

srt_torch/__init__.py

@@ -0,0 +1,21 @@
+from srt_torch.main import (
+    SurgicalRobotTransformer,
+    ModelConfig,
+    RobotObservation,
+    RobotAction,
+    ImageEncoder,
+    TransformerBlock,
+    ActionPredictor,
+    normalize_rotation_6d,
+)
+
+__all__ = [
+    "SurgicalRobotTransformer",
+    "ModelConfig",
+    "RobotObservation",
+    "RobotAction",
+    "ImageEncoder",
+    "TransformerBlock",
+    "ActionPredictor",
+    "normalize_rotation_6d",
+]