diff --git a/.gitattributes b/.gitattributes index df7d2d5..4135de8 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,2 @@ *.memmap filter=lfs diff=lfs merge=lfs -text +*.stl filter=lfs diff=lfs merge=lfs -text diff --git a/.github/poetry/cpu/poetry.lock b/.github/poetry/cpu/poetry.lock index d558d50..ec2b558 100644 --- a/.github/poetry/cpu/poetry.lock +++ b/.github/poetry/cpu/poetry.lock @@ -338,73 +338,6 @@ files = [ {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -[[package]] -name = "cython" -version = "3.0.9" -description = "The Cython compiler for writing C extensions in the Python language." -optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" -files = [ - {file = "Cython-3.0.9-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:296bd30d4445ac61b66c9d766567f6e81a6e262835d261e903c60c891a6729d3"}, - {file = "Cython-3.0.9-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f496b52845cb45568a69d6359a2c335135233003e708ea02155c10ce3548aa89"}, - {file = "Cython-3.0.9-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:858c3766b9aa3ab8a413392c72bbab1c144a9766b7c7bfdef64e2e414363fa0c"}, - {file = "Cython-3.0.9-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c0eb1e6ef036028a52525fd9a012a556f6dd4788a0e8755fe864ba0e70cde2ff"}, - {file = "Cython-3.0.9-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:c8191941073ea5896321de3c8c958fd66e5f304b0cd1f22c59edd0b86c4dd90d"}, - {file = "Cython-3.0.9-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:e32b016030bc72a8a22a1f21f470a2f57573761a4f00fbfe8347263f4fbdb9f1"}, - {file = "Cython-3.0.9-cp310-cp310-win32.whl", hash = "sha256:d6f3ff1cd6123973fe03e0fb8ee936622f976c0c41138969975824d08886572b"}, - {file = "Cython-3.0.9-cp310-cp310-win_amd64.whl", hash = "sha256:56f3b643dbe14449248bbeb9a63fe3878a24256664bc8c8ef6efd45d102596d8"}, - {file = "Cython-3.0.9-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:35e6665a20d6b8a152d72b7fd87dbb2af6bb6b18a235b71add68122d594dbd41"}, - {file = "Cython-3.0.9-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92f4960c40ad027bd8c364c50db11104eadc59ffeb9e5b7f605ca2f05946e20"}, - {file = "Cython-3.0.9-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:38df37d0e732fbd9a2fef898788492e82b770c33d1e4ed12444bbc8a3b3f89c0"}, - {file = "Cython-3.0.9-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ad7fd88ebaeaf2e76fd729a8919fae80dab3d6ac0005e28494261d52ff347a8f"}, - {file = "Cython-3.0.9-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:1365d5f76bf4d19df3d19ce932584c9bb76e9fb096185168918ef9b36e06bfa4"}, - {file = "Cython-3.0.9-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:c232e7f279388ac9625c3e5a5a9f0078a9334959c5d6458052c65bbbba895e1e"}, - {file = "Cython-3.0.9-cp311-cp311-win32.whl", hash = "sha256:357e2fad46a25030b0c0496487e01a9dc0fdd0c09df0897f554d8ba3c1bc4872"}, - {file = "Cython-3.0.9-cp311-cp311-win_amd64.whl", hash = "sha256:1315aee506506e8d69cf6631d8769e6b10131fdcc0eb66df2698f2a3ddaeeff2"}, - {file = "Cython-3.0.9-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:157973807c2796addbed5fbc4d9c882ab34bbc60dc297ca729504901479d5df7"}, - {file = "Cython-3.0.9-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:00b105b5d050645dd59e6767bc0f18b48a4aa11c85f42ec7dd8181606f4059e3"}, - {file = "Cython-3.0.9-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ac5536d09bef240cae0416d5a703d298b74c7bbc397da803ac9d344e732d4369"}, - {file = "Cython-3.0.9-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:09c44501d476d16aaa4cbc29c87f8c0f54fc20e69b650d59cbfa4863426fc70c"}, - {file = "Cython-3.0.9-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:cc9c3b9f20d8e298618e5ccd32083ca386e785b08f9893fbec4c50b6b85be772"}, - {file = "Cython-3.0.9-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:a30d96938c633e3ec37000ac3796525da71254ef109e66bdfd78f29891af6454"}, - {file = "Cython-3.0.9-cp312-cp312-win32.whl", hash = "sha256:757ca93bdd80702546df4d610d2494ef2e74249cac4d5ba9464589fb464bd8a3"}, - {file = "Cython-3.0.9-cp312-cp312-win_amd64.whl", hash = "sha256:1dc320a9905ab95414013f6de805efbff9e17bb5fb3b90bbac533f017bec8136"}, - {file = "Cython-3.0.9-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:4ae349960ebe0da0d33724eaa7f1eb866688fe5434cc67ce4dbc06d6a719fbfc"}, - {file = "Cython-3.0.9-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:63d2537bf688247f76ded6dee28ebd26274f019309aef1eb4f2f9c5c482fde2d"}, - {file = "Cython-3.0.9-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:36f5a2dfc724bea1f710b649f02d802d80fc18320c8e6396684ba4a48412445a"}, - {file = "Cython-3.0.9-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:deaf4197d4b0bcd5714a497158ea96a2bd6d0f9636095437448f7e06453cc83d"}, - {file = "Cython-3.0.9-cp36-cp36m-musllinux_1_1_aarch64.whl", hash = "sha256:000af6deb7412eb7ac0c635ff5e637fb8725dd0a7b88cc58dfc2b3de14e701c4"}, - {file = "Cython-3.0.9-cp36-cp36m-musllinux_1_1_x86_64.whl", hash = "sha256:15c7f5c2d35bed9aa5f2a51eaac0df23ae72f2dbacf62fc672dd6bfaa75d2d6f"}, - {file = "Cython-3.0.9-cp36-cp36m-win32.whl", hash = "sha256:f49aa4970cd3bec66ac22e701def16dca2a49c59cceba519898dd7526e0be2c0"}, - {file = "Cython-3.0.9-cp36-cp36m-win_amd64.whl", hash = "sha256:4558814fa025b193058d42eeee498a53d6b04b2980d01339fc2444b23fd98e58"}, - {file = "Cython-3.0.9-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:539cd1d74fd61f6cfc310fa6bbbad5adc144627f2b7486a07075d4e002fd6aad"}, - {file = "Cython-3.0.9-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c3232926cd406ee02eabb732206f6e882c3aed9d58f0fea764013d9240405bcf"}, - {file = "Cython-3.0.9-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33b6ac376538a7fc8c567b85d3c71504308a9318702ec0485dd66c059f3165cb"}, - {file = "Cython-3.0.9-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2cc92504b5d22ac66031ffb827bd3a967fc75a5f0f76ab48bce62df19be6fdfd"}, - {file = "Cython-3.0.9-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:22b8fae756c5c0d8968691bed520876de452f216c28ec896a00739a12dba3bd9"}, - {file = "Cython-3.0.9-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:9cda0d92a09f3520f29bd91009f1194ba9600777c02c30c6d2d4ac65fb63e40d"}, - {file = "Cython-3.0.9-cp37-cp37m-win32.whl", hash = "sha256:ec612418490941ed16c50c8d3784c7bdc4c4b2a10c361259871790b02ec8c1db"}, - {file = "Cython-3.0.9-cp37-cp37m-win_amd64.whl", hash = "sha256:976c8d2bedc91ff6493fc973d38b2dc01020324039e2af0e049704a8e1b22936"}, - {file = "Cython-3.0.9-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5055988b007c92256b6e9896441c3055556038c3497fcbf8c921a6c1fce90719"}, - {file = "Cython-3.0.9-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d9360606d964c2d0492a866464efcf9d0a92715644eede3f6a2aa696de54a137"}, - {file = "Cython-3.0.9-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:02c6e809f060bed073dc7cba1648077fe3b68208863d517c8b39f3920eecf9dd"}, - {file = "Cython-3.0.9-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:95ed792c966f969cea7489c32ff90150b415c1f3567db8d5a9d489c7c1602dac"}, - {file = "Cython-3.0.9-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:8edd59d22950b400b03ca78d27dc694d2836a92ef0cac4f64cb4b2ff902f7e25"}, - {file = "Cython-3.0.9-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:4cf0ed273bf60e97922fcbbdd380c39693922a597760160b4b4355e6078ca188"}, - {file = "Cython-3.0.9-cp38-cp38-win32.whl", hash = "sha256:5eb9bd4ae12ebb2bc79a193d95aacf090fbd8d7013e11ed5412711650cb34934"}, - {file = "Cython-3.0.9-cp38-cp38-win_amd64.whl", hash = "sha256:44457279da56e0f829bb1fc5a5dc0836e5d498dbcf9b2324f32f7cc9d2ec6569"}, - {file = "Cython-3.0.9-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:c4b419a1adc2af43f4660e2f6eaf1e4fac2dbac59490771eb8ac3d6063f22356"}, - {file = "Cython-3.0.9-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f836192140f033b2319a0128936367c295c2b32e23df05b03b672a6015757ea"}, - {file = "Cython-3.0.9-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2fd198c1a7f8e9382904d622cc0efa3c184605881fd5262c64cbb7168c4c1ec5"}, - {file = "Cython-3.0.9-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a274fe9ca5c53fafbcf5c8f262f8ad6896206a466f0eeb40aaf36a7951e957c0"}, - {file = "Cython-3.0.9-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:158c38360bbc5063341b1e78d3737f1251050f89f58a3df0d10fb171c44262be"}, - {file = "Cython-3.0.9-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:8bf30b045f7deda0014b042c1b41c1d272facc762ab657529e3b05505888e878"}, - {file = "Cython-3.0.9-cp39-cp39-win32.whl", hash = "sha256:9a001fd95c140c94d934078544ff60a3c46aca2dc86e75a76e4121d3cd1f4b33"}, - {file = "Cython-3.0.9-cp39-cp39-win_amd64.whl", hash = "sha256:530c01c4aebba709c0ec9c7ecefe07177d0b9fd7ffee29450a118d92192ccbdf"}, - {file = "Cython-3.0.9-py2.py3-none-any.whl", hash = "sha256:bf96417714353c5454c2e3238fca9338599330cf51625cdc1ca698684465646f"}, - {file = "Cython-3.0.9.tar.gz", hash = "sha256:a2d354f059d1f055d34cfaa62c5b68bc78ac2ceab6407148d47fb508cf3ba4f3"}, -] - [[package]] name = "debugpy" version = "1.8.1" @@ -500,13 +433,13 @@ files = [ [[package]] name = "dm-control" -version = "1.0.16" +version = "1.0.14" description = "Continuous control environments and MuJoCo Python bindings." optional = false python-versions = ">=3.8" files = [ - {file = "dm_control-1.0.16-py3-none-any.whl", hash = "sha256:341f582e26d88556ac0dd1963e11c25f93ef6649b07c75d04570a7de3edebd1b"}, - {file = "dm_control-1.0.16.tar.gz", hash = "sha256:bbdd6dc54b4dbc33eef6c43294edb0a927c3a5f35621e698480f88f55f48a1fd"}, + {file = "dm_control-1.0.14-py3-none-any.whl", hash = "sha256:883c63244a7ebf598700a97564ed19fffd3479ca79efd090aed881609cdb9fc6"}, + {file = "dm_control-1.0.14.tar.gz", hash = "sha256:def1ece747b6f175c581150826b50f1a6134086dab34f8f3fd2d088ea035cf3d"}, ] [package.dependencies] @@ -516,7 +449,7 @@ dm-tree = "!=0.1.2" glfw = "*" labmaze = "*" lxml = "*" -mujoco = ">=3.1.1" +mujoco = ">=2.3.7" numpy = ">=1.9.0" protobuf = ">=3.19.4" pyopengl = ">=3.1.4" @@ -635,43 +568,6 @@ files = [ {file = "einops-0.7.0.tar.gz", hash = "sha256:b2b04ad6081a3b227080c9bf5e3ace7160357ff03043cd66cc5b2319eb7031d1"}, ] -[[package]] -name = "etils" -version = "1.7.0" -description = "Collection of common python utils" -optional = false -python-versions = ">=3.10" -files = [ - {file = "etils-1.7.0-py3-none-any.whl", hash = "sha256:61af8f7c242171de15e22e5da02d527cb9e677d11f8bcafe18fcc3548eee3e60"}, - {file = "etils-1.7.0.tar.gz", hash = "sha256:97b68fd25e185683215286ef3a54e38199b6245f5fe8be6bedc1189be4256350"}, -] - -[package.dependencies] -fsspec = {version = "*", optional = true, markers = "extra == \"epath\""} -importlib_resources = {version = "*", optional = true, markers = "extra == \"epath\""} -typing_extensions = {version = "*", optional = true, markers = "extra == \"epy\""} -zipp = {version = "*", optional = true, markers = "extra == \"epath\""} - -[package.extras] -all = ["etils[array-types]", "etils[eapp]", "etils[ecolab]", "etils[edc]", "etils[enp]", "etils[epath-gcs]", "etils[epath-s3]", "etils[epath]", "etils[epy]", "etils[etqdm]", "etils[etree-dm]", "etils[etree-jax]", "etils[etree-tf]", "etils[etree]"] -array-types = ["etils[enp]"] -dev = ["chex", "dataclass_array", "optree", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-subtests", "pytest-xdist", "torch"] -docs = ["etils[all,dev]", "sphinx-apitree[ext]"] -eapp = ["absl-py", "etils[epy]", "simple_parsing"] -ecolab = ["etils[enp]", "etils[epy]", "etils[etree]", "jupyter", "mediapy", "numpy", "packaging", "protobuf"] -edc = ["etils[epy]"] -enp = ["etils[epy]", "numpy"] -epath = ["etils[epy]", "fsspec", "importlib_resources", "typing_extensions", "zipp"] -epath-gcs = ["etils[epath]", "gcsfs"] -epath-s3 = ["etils[epath]", "s3fs"] -epy = ["typing_extensions"] -etqdm = ["absl-py", "etils[epy]", "tqdm"] -etree = ["etils[array-types]", "etils[enp]", "etils[epy]", "etils[etqdm]"] -etree-dm = ["dm-tree", "etils[etree]"] -etree-jax = ["etils[etree]", "jax[cpu]"] -etree-tf = ["etils[etree]", "tensorflow"] -lazy-imports = ["etils[ecolab]"] - [[package]] name = "exceptiongroup" version = "1.2.0" @@ -686,6 +582,17 @@ files = [ [package.extras] test = ["pytest (>=6)"] +[[package]] +name = "farama-notifications" +version = "0.0.4" +description = "Notifications for all Farama Foundation maintained libraries." +optional = false +python-versions = "*" +files = [ + {file = "Farama-Notifications-0.0.4.tar.gz", hash = "sha256:13fceff2d14314cf80703c8266462ebf3733c7d165336eee998fc58e545efd18"}, + {file = "Farama_Notifications-0.0.4-py3-none-any.whl", hash = "sha256:14de931035a41961f7c056361dc7f980762a143d05791ef5794a751a2caf05ae"}, +] + [[package]] name = "fasteners" version = "0.19" @@ -887,43 +794,58 @@ files = [ protobuf = ["grpcio-tools (>=1.62.1)"] [[package]] -name = "gym" -version = "0.26.2" -description = "Gym: A universal API for reinforcement learning environments" +name = "gymnasium" +version = "0.29.1" +description = "A standard API for reinforcement learning and a diverse set of reference environments (formerly Gym)." optional = false -python-versions = ">=3.6" +python-versions = ">=3.8" files = [ - {file = "gym-0.26.2.tar.gz", hash = "sha256:e0d882f4b54f0c65f203104c24ab8a38b039f1289986803c7d02cdbe214fbcc4"}, + {file = "gymnasium-0.29.1-py3-none-any.whl", hash = "sha256:61c3384b5575985bb7f85e43213bcb40f36fcdff388cae6bc229304c71f2843e"}, + {file = "gymnasium-0.29.1.tar.gz", hash = "sha256:1a532752efcb7590478b1cc7aa04f608eb7a2fdad5570cd217b66b6a35274bb1"}, ] [package.dependencies] cloudpickle = ">=1.2.0" -gym_notices = ">=0.0.4" -numpy = ">=1.18.0" +farama-notifications = ">=0.0.1" +numpy = ">=1.21.0" +typing-extensions = ">=4.3.0" [package.extras] accept-rom-license = ["autorom[accept-rom-license] (>=0.4.2,<0.5.0)"] -all = ["ale-py (>=0.8.0,<0.9.0)", "box2d-py (==2.3.5)", "imageio (>=2.14.1)", "lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "mujoco (==2.2)", "mujoco_py (>=2.1,<2.2)", "opencv-python (>=3.0)", "pygame (==2.1.0)", "pytest (==7.0.1)", "swig (==4.*)"] -atari = ["ale-py (>=0.8.0,<0.9.0)"] -box2d = ["box2d-py (==2.3.5)", "pygame (==2.1.0)", "swig (==4.*)"] -classic-control = ["pygame (==2.1.0)"] -mujoco = ["imageio (>=2.14.1)", "mujoco (==2.2)"] -mujoco-py = ["mujoco_py (>=2.1,<2.2)"] -other = ["lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "opencv-python (>=3.0)"] -testing = ["box2d-py (==2.3.5)", "imageio (>=2.14.1)", "lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "mujoco (==2.2)", "mujoco_py (>=2.1,<2.2)", "opencv-python (>=3.0)", "pygame (==2.1.0)", "pytest (==7.0.1)", "swig (==4.*)"] -toy-text = ["pygame (==2.1.0)"] +all = ["box2d-py (==2.3.5)", "cython (<3)", "imageio (>=2.14.1)", "jax (>=0.4.0)", "jaxlib (>=0.4.0)", "lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "mujoco (>=2.3.3)", "mujoco-py (>=2.1,<2.2)", "opencv-python (>=3.0)", "pygame (>=2.1.3)", "shimmy[atari] (>=0.1.0,<1.0)", "swig (==4.*)", "torch (>=1.0.0)"] +atari = ["shimmy[atari] (>=0.1.0,<1.0)"] +box2d = ["box2d-py (==2.3.5)", "pygame (>=2.1.3)", "swig (==4.*)"] +classic-control = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"] +jax = ["jax (>=0.4.0)", "jaxlib (>=0.4.0)"] +mujoco = ["imageio (>=2.14.1)", "mujoco (>=2.3.3)"] +mujoco-py = ["cython (<3)", "cython (<3)", "mujoco-py (>=2.1,<2.2)", "mujoco-py (>=2.1,<2.2)"] +other = ["lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "opencv-python (>=3.0)", "torch (>=1.0.0)"] +testing = ["pytest (==7.1.3)", "scipy (>=1.7.3)"] +toy-text = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"] [[package]] -name = "gym-notices" -version = "0.0.8" -description = "Notices for gym" +name = "gymnasium-robotics" +version = "1.2.4" +description = "Robotics environments for the Gymnasium repo." optional = false -python-versions = "*" +python-versions = ">=3.8" files = [ - {file = "gym-notices-0.0.8.tar.gz", hash = "sha256:ad25e200487cafa369728625fe064e88ada1346618526102659b4640f2b4b911"}, - {file = "gym_notices-0.0.8-py3-none-any.whl", hash = "sha256:e5f82e00823a166747b4c2a07de63b6560b1acb880638547e0cabf825a01e463"}, + {file = "gymnasium-robotics-1.2.4.tar.gz", hash = "sha256:d304192b066f8b800599dfbe3d9d90bba9b761ee884472bdc4d05968a8bc61cb"}, + {file = "gymnasium_robotics-1.2.4-py3-none-any.whl", hash = "sha256:c2cb23e087ca0280ae6802837eb7b3a6d14e5bd24c00803ab09f015fcff3eef5"}, ] +[package.dependencies] +gymnasium = ">=0.26" +imageio = "*" +Jinja2 = ">=3.0.3" +mujoco = ">=2.3.3,<3.0" +numpy = ">=1.21.0" +PettingZoo = ">=1.23.0" + +[package.extras] +mujoco-py = ["cython (<3)", "mujoco-py (>=2.1,<2.2)"] +testing = ["Jinja2 (>=3.0.3)", "PettingZoo (>=1.23.0)", "cython (<3)", "mujoco-py (>=2.1,<2.2)", "pytest (==7.0.1)"] + [[package]] name = "h5py" version = "3.10.0" @@ -1105,21 +1027,6 @@ docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.link perf = ["ipython"] testing = ["flufl.flake8", "importlib-resources (>=1.3)", "jaraco.test (>=5.4)", "packaging", "pyfakefs", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-mypy", "pytest-perf (>=0.9.2)", "pytest-ruff (>=0.2.1)"] -[[package]] -name = "importlib-resources" -version = "6.3.2" -description = "Read resources from Python packages" -optional = false -python-versions = ">=3.8" -files = [ - {file = "importlib_resources-6.3.2-py3-none-any.whl", hash = "sha256:f41f4098b16cd140a97d256137cfd943d958219007990b2afb00439fc623f580"}, - {file = "importlib_resources-6.3.2.tar.gz", hash = "sha256:963eb79649252b0160c1afcfe5a1d3fe3ad66edd0a8b114beacffb70c0674223"}, -] - -[package.extras] -docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (<7.2.5)", "sphinx (>=3.5)", "sphinx-lint"] -testing = ["jaraco.collections", "pytest (>=6)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-mypy", "pytest-ruff (>=0.2.1)", "zipp (>=3.17)"] - [[package]] name = "iniconfig" version = "2.0.0" @@ -1457,65 +1364,44 @@ tests = ["pytest (>=4.6)"] [[package]] name = "mujoco" -version = "3.1.3" +version = "2.3.7" description = "MuJoCo Physics Simulator" optional = false python-versions = ">=3.8" files = [ - {file = "mujoco-3.1.3-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:1a07e33443ca88c77128336e550502c58721e37b3830af29f0118311c17d826e"}, - {file = "mujoco-3.1.3-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:da442c45fa08cf7f307a6f2484ff382b90714b9f52aaceffd5fcb8536dbdc11c"}, - {file = "mujoco-3.1.3-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ec08dbfddef6e4c6d7b03685b929ed134e8eb9d0dbc788752ff54216b7b3544e"}, - {file = "mujoco-3.1.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f1578279ff581ed1c70893cc16ecf48a048a14568e9e64b446a2d32c22b1154c"}, - {file = "mujoco-3.1.3-cp310-cp310-win_amd64.whl", hash = "sha256:9a359e7787e1d0bbdb9fafeb31df61261a4cdc42d0a5d77c91fbe57c63e4c6fd"}, - {file = "mujoco-3.1.3-cp311-cp311-macosx_10_16_x86_64.whl", hash = "sha256:b070805d65ee6b708ddf1a16a16fc2073ce2d1eea8ea26352b8aee4071de274c"}, - {file = "mujoco-3.1.3-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d789a95150cf1bef21e3a3431c26263730b0437ec3b4794b2eed0f900185746e"}, - {file = "mujoco-3.1.3-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6661aa27c81be338ce0973ba6e83f655ff3cc023ea9d62398f130b46478f708a"}, - {file = "mujoco-3.1.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f79dc6134c90a7274d2663c07bea6d45629ea52ce40bf6722c5d506df909b4b9"}, - {file = "mujoco-3.1.3-cp311-cp311-win_amd64.whl", hash = "sha256:b1df674d9486e1bd2e93fb69009d8db4adcf4b3b7edc92da5c98d1c6a2ea7a28"}, - {file = "mujoco-3.1.3-cp312-cp312-macosx_10_16_x86_64.whl", hash = "sha256:51841750310a1c4b5e7c7f19d28fe5e3deea0e2c7cc60ebab33c2f07360b1700"}, - {file = "mujoco-3.1.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fe4318af5b14ea39bc5b8892c69797a1a9deb02199178814be16abb5611308fb"}, - {file = "mujoco-3.1.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:079f293a56c2b3aa6b4101c3822ee5587b5cc9bf35028afdd1f2128db102ad20"}, - {file = "mujoco-3.1.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d4c120414bf89a11538e3f5eb1de6bcd6c4aeade9775ecad3e4eea27d88e1492"}, - {file = "mujoco-3.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:7fc9d69383dc0f7c4b775b2be829a065fb78dca743a25f9d864d52174c916b2b"}, - {file = "mujoco-3.1.3-cp38-cp38-macosx_10_16_x86_64.whl", hash = "sha256:5a29004079a40d23836228647bae9ea41f77fd7e407e8ad642dc72054e5a099e"}, - {file = "mujoco-3.1.3-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:c1d1e083d8825faf9e2609d4e749cc5629ed7735374ed68eb3dde63dd0e4fe73"}, - {file = "mujoco-3.1.3-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9cc267cf0de8de3c8b317f7c12b2d7a484a7f462263f8ce4c8ae18e9d6817897"}, - {file = "mujoco-3.1.3-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d67a2aae8d5d58f7ac41d0bcffcd745955415887843bc34da8e3d794b46afbae"}, - {file = "mujoco-3.1.3-cp38-cp38-win_amd64.whl", hash = "sha256:8ddb9a07d5ad59c67f2d7e79568cba27ad68cf2284a68370f2054dce2e6e4128"}, - {file = "mujoco-3.1.3-cp39-cp39-macosx_10_16_x86_64.whl", hash = "sha256:acdc761e8fa7d4bfb9f262b8886dbb3dd41a957c3ef7ec126aae3342f68b1293"}, - {file = "mujoco-3.1.3-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:49bad0da02ebf67ab37a6f6fe435dfc6339f0b46b51b452ee79aaffa5b73659b"}, - {file = "mujoco-3.1.3-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:daa6a3ca50a3769ebfd59274651d2edc76b177cd950560022120fb77cd51f607"}, - {file = "mujoco-3.1.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:80f3e99d1f2bb02bf7903ba4fef9c31e05fba439b7292ed751390ec78e1eb890"}, - {file = "mujoco-3.1.3-cp39-cp39-win_amd64.whl", hash = "sha256:2d2fe38b1a7f64e708e8b9a96cf7677027b33fb6e059184163976c6c03fef4cc"}, - {file = "mujoco-3.1.3.tar.gz", hash = "sha256:f700d074031060b46111ddb60432d00425f821eeeaf0ccc76ed95d47861bd4de"}, + {file = "mujoco-2.3.7-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:e8714a5ff6a1561b364b7b4648d4c0c8d13e751874cf7401c309b9d23fa9598b"}, + {file = "mujoco-2.3.7-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a934315f858a4e0c4b90a682fde519471cfdd7baa64435179da8cd20d4ae3f99"}, + {file = "mujoco-2.3.7-cp310-cp310-macosx_11_0_x86_64.whl", hash = "sha256:36513024330f88b5f9a43558efef5692b33599bffd5141029b690a27918ffcbe"}, + {file = "mujoco-2.3.7-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6d4eede8ba8210fbd3d3cd1dbf69e24dd1541aa74c5af5b8adbbbf65504b6dba"}, + {file = "mujoco-2.3.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ab85fafc9d5a091c712947573b7e694512d283876bf7f33ae3f8daad3a20c0db"}, + {file = "mujoco-2.3.7-cp310-cp310-win_amd64.whl", hash = "sha256:f8b7e13fef8c813d91b78f975ed0815157692777907ffa4b4be53a4edb75019b"}, + {file = "mujoco-2.3.7-cp311-cp311-macosx_10_16_x86_64.whl", hash = "sha256:779520216f72a8e370e3f0cdd71b45c3b7384c63331a3189194c930a3e7cff5c"}, + {file = "mujoco-2.3.7-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:9d4018053879016282d27ab7a91e292c72d44efb5a88553feacfe5b843dde103"}, + {file = "mujoco-2.3.7-cp311-cp311-macosx_11_0_x86_64.whl", hash = "sha256:3149b16b8122ee62642474bfd2871064e8edc40235471cf5d84be3569afc0312"}, + {file = "mujoco-2.3.7-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c08660a8d52ef3efde76095f0991e807703a950c1e882d2bcd984b9a846626f7"}, + {file = "mujoco-2.3.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:426af8965f8636d94a0f75740c3024a62b3e585020ee817ef5208ec844a1ad94"}, + {file = "mujoco-2.3.7-cp311-cp311-win_amd64.whl", hash = "sha256:215415a8e98a4b50625beae859079d5e0810b2039e50420f0ba81763c34abb59"}, + {file = "mujoco-2.3.7-cp38-cp38-macosx_10_16_x86_64.whl", hash = "sha256:8b78d14f4c60cea3c58e046bd4de453fb5b9b33aca6a25fc91d39a53f3a5342a"}, + {file = "mujoco-2.3.7-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:5c6f5a51d6f537a4bf294cf73816f3a6384573f8f10a5452b044df2771412a96"}, + {file = "mujoco-2.3.7-cp38-cp38-macosx_11_0_x86_64.whl", hash = "sha256:ea8911e6047f92d7d775701f37e4c093971b6def3160f01d0b6926e29a7e962e"}, + {file = "mujoco-2.3.7-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7473a3de4dd1a8762d569ffb139196b4c5e7eca27d256df97b6cd4c66d2a09b2"}, + {file = "mujoco-2.3.7-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:40e7e2d8f93d2495ec74efec84e5118ecc6e1d85157a844789c73c9ac9a4e28e"}, + {file = "mujoco-2.3.7-cp38-cp38-win_amd64.whl", hash = "sha256:720bc228a2023b3b0ed6af78f5b0f8ea36867be321d473321555c57dbf6e4e5b"}, + {file = "mujoco-2.3.7-cp39-cp39-macosx_10_16_x86_64.whl", hash = "sha256:855e79686366442aa410246043b44f7d842d3900d68fe7e37feb42147db9d707"}, + {file = "mujoco-2.3.7-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:98947f4a742d34d36f3c3f83e9167025bb0414bbaa4bd859b0673bdab9959963"}, + {file = "mujoco-2.3.7-cp39-cp39-macosx_11_0_x86_64.whl", hash = "sha256:d42818f2ee5d1632dbce31d136ed5ff868db54b04e4e9aca0c5a3ac329f8a90f"}, + {file = "mujoco-2.3.7-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9237e1ba14bced9449c31199e6d5be49547f3a4c99bc83b196af7ca45fd73b83"}, + {file = "mujoco-2.3.7-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:39b728ea638245b150e2650c5433e6952e0ed3798c63e47e264574270caea2a3"}, + {file = "mujoco-2.3.7-cp39-cp39-win_amd64.whl", hash = "sha256:9c721a5042b99d948d5f0296a534bcce3f142c777c4d7642f503a539513f3912"}, + {file = "mujoco-2.3.7.tar.gz", hash = "sha256:422041f1ce37c6d151fbced1048df626837e94fe3cd9f813585907046336a7d0"}, ] [package.dependencies] absl-py = "*" -etils = {version = "*", extras = ["epath"]} glfw = "*" numpy = "*" pyopengl = "*" -[[package]] -name = "mujoco-py" -version = "2.1.2.14" -description = "" -optional = false -python-versions = ">=3.6" -files = [ - {file = "mujoco-py-2.1.2.14.tar.gz", hash = "sha256:eb5b14485acf80a3cf8c15f4b080c6a28a9f79e68869aa696d16cbd51ea7706f"}, - {file = "mujoco_py-2.1.2.14-py3-none-any.whl", hash = "sha256:37c0b41bc0153a8a0eb3663103a67c60f65467753f74e4ff6e68b879f3e3a71f"}, -] - -[package.dependencies] -cffi = ">=1.10" -Cython = ">=0.27.2" -fasteners = ">=0.15,<1.0" -glfw = ">=1.4.0" -imageio = ">=2.1.2" -numpy = ">=1.11" - [[package]] name = "networkx" version = "3.2.1" @@ -1790,6 +1676,31 @@ sql-other = ["SQLAlchemy (>=2.0.0)", "adbc-driver-postgresql (>=0.8.0)", "adbc-d test = ["hypothesis (>=6.46.1)", "pytest (>=7.3.2)", "pytest-xdist (>=2.2.0)"] xml = ["lxml (>=4.9.2)"] +[[package]] +name = "pettingzoo" +version = "1.24.3" +description = "Gymnasium for multi-agent reinforcement learning." +optional = false +python-versions = ">=3.8" +files = [ + {file = "pettingzoo-1.24.3-py3-none-any.whl", hash = "sha256:23ed90517d2e8a7098bdaf5e31234b3a7f7b73ca578d70d1ca7b9d0cb0e37982"}, + {file = "pettingzoo-1.24.3.tar.gz", hash = "sha256:91f9094f18e06fb74b98f4099cd22e8ae4396125e51719d50b30c9f1c7ab07e6"}, +] + +[package.dependencies] +gymnasium = ">=0.28.0" +numpy = ">=1.21.0" + +[package.extras] +all = ["box2d-py (==2.3.5)", "chess (==1.9.4)", "multi-agent-ale-py (==0.1.11)", "pillow (>=8.0.1)", "pygame (==2.3.0)", "pymunk (==6.2.0)", "rlcard (==1.0.5)", "scipy (>=1.4.1)", "shimmy[openspiel] (>=1.2.0)"] +atari = ["multi-agent-ale-py (==0.1.11)", "pygame (==2.3.0)"] +butterfly = ["pygame (==2.3.0)", "pymunk (==6.2.0)"] +classic = ["chess (==1.9.4)", "pygame (==2.3.0)", "rlcard (==1.0.5)", "shimmy[openspiel] (>=1.2.0)"] +mpe = ["pygame (==2.3.0)"] +other = ["pillow (>=8.0.1)"] +sisl = ["box2d-py (==2.3.5)", "pygame (==2.3.0)", "pymunk (==6.2.0)", "scipy (>=1.4.1)"] +testing = ["AutoROM", "pre-commit", "pynput", "pytest", "pytest-cov", "pytest-markdown-docs", "pytest-xdist"] + [[package]] name = "pillow" version = "10.2.0" @@ -3305,4 +3216,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "cbd9aedcb3a24417b85124fb94db706dd6ca0a90dfb610b0aebdcd3aa2a0333c" +content-hash = "93c406139c456780b3d309d7ed3d68ea60cc0e8893c1ee717692984e573d3404" diff --git a/.github/poetry/cpu/pyproject.toml b/.github/poetry/cpu/pyproject.toml index 2f5a542..586ef21 100644 --- a/.github/poetry/cpu/pyproject.toml +++ b/.github/poetry/cpu/pyproject.toml @@ -21,7 +21,6 @@ packages = [{include = "lerobot"}] [tool.poetry.dependencies] python = "^3.10" -cython = "^3.0.8" termcolor = "^2.4.0" omegaconf = "^2.3.0" dm-env = "^1.6" @@ -42,17 +41,17 @@ mpmath = "^1.3.0" torch = {version = "^2.2.1", source = "torch-cpu"} tensordict = {git = "https://github.com/pytorch/tensordict"} torchrl = {git = "https://github.com/pytorch/rl", rev = "13bef426dcfa5887c6e5034a6e9697993fa92c37"} -mujoco = "^3.1.2" -mujoco-py = "^2.1.2.14" -gym = "^0.26.2" +mujoco = "^2.3.7" opencv-python = "^4.9.0.80" diffusers = "^0.26.3" torchvision = {version = "^0.17.1", source = "torch-cpu"} h5py = "^3.10.0" dm = "^1.3" -dm-control = "^1.0.16" +dm-control = "1.0.14" robomimic = "0.2.0" huggingface-hub = "^0.21.4" +gymnasium-robotics = "^1.2.4" +gymnasium = "^0.29.1" [tool.poetry.group.dev.dependencies] diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 728a978..4788a17 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -16,12 +16,8 @@ jobs: ${{ github.event_name == 'push' }} runs-on: ubuntu-latest env: - POETRY_VERSION: 1.8.1 + POETRY_VERSION: 1.8.2 DATA_DIR: tests/data - TMPDIR: ~/tmp - TEMP: ~/tmp - TMP: ~/tmp - PYOPENGL_PLATFORM: egl MUJOCO_GL: egl LEROBOT_TESTS_DEVICE: cpu steps: @@ -86,6 +82,10 @@ jobs: - name: Install dependencies if: steps.restore-dependencies-cache.outputs.cache-hit != 'true' + env: + TMPDIR: ~/tmp + TEMP: ~/tmp + TMP: ~/tmp run: | mkdir ~/tmp poetry install --no-interaction --no-root diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2b79434..765b678 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,11 +14,11 @@ repos: - id: end-of-file-fixer - id: trailing-whitespace - repo: https://github.com/asottile/pyupgrade - rev: v3.15.1 + rev: v3.15.2 hooks: - id: pyupgrade - repo: https://github.com/astral-sh/ruff-pre-commit - rev: v0.2.2 + rev: v0.3.4 hooks: - id: ruff args: [--fix] diff --git a/LICENSE b/LICENSE index 26534b4..a603343 100644 --- a/LICENSE +++ b/LICENSE @@ -253,6 +253,31 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +## Some of lerobot's code is derived from simxarm, which is subject to the following copyright notice: + +MIT License + +Copyright (c) 2023 Nicklas Hansen & Yanjie Ze + +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. + + ## Some of lerobot's code is derived from ALOHA, which is subject to the following copyright notice: MIT License diff --git a/lerobot/__version__.py b/lerobot/__version__.py index 0fc9dd1..6232b69 100644 --- a/lerobot/__version__.py +++ b/lerobot/__version__.py @@ -1,4 +1,4 @@ -""" To enable `lerobot.__version__` """ +"""To enable `lerobot.__version__`""" from importlib.metadata import PackageNotFoundError, version diff --git a/lerobot/common/datasets/simxarm.py b/lerobot/common/datasets/simxarm.py index d7e2e18..74bec4c 100644 --- a/lerobot/common/datasets/simxarm.py +++ b/lerobot/common/datasets/simxarm.py @@ -40,7 +40,7 @@ class SimxarmExperienceReplay(AbstractExperienceReplay): def __init__( self, dataset_id: str, - version: str | None = None, + version: str | None = "v1.1", batch_size: int = None, *, shuffle: bool = True, @@ -67,11 +67,11 @@ class SimxarmExperienceReplay(AbstractExperienceReplay): ) def _download_and_preproc_obsolete(self): - assert self.root is not None + # assert self.root is not None # TODO(rcadene): finish download - download() + # download() - dataset_path = self.root / f"{self.dataset_id}_raw" / "buffer.pkl" + dataset_path = self.root / f"{self.dataset_id}" / "buffer.pkl" print(f"Using offline dataset '{dataset_path}'") with open(dataset_path, "rb") as f: dataset_dict = pickle.load(f) @@ -105,15 +105,19 @@ class SimxarmExperienceReplay(AbstractExperienceReplay): "frame_id": torch.arange(0, num_frames, 1), ("next", "observation", "image"): next_image, ("next", "observation", "state"): next_state, - ("next", "observation", "reward"): next_reward, - ("next", "observation", "done"): next_done, + ("next", "reward"): next_reward, + ("next", "done"): next_done, }, batch_size=num_frames, ) if episode_id == 0: # hack to initialize tensordict data structure to store episodes - td_data = episode[0].expand(total_frames).memmap_like(self.root / f"{self.dataset_id}") + td_data = ( + episode[0] + .expand(total_frames) + .memmap_like(self.root / f"{self.dataset_id}" / "replay_buffer") + ) td_data[idx0:idx1] = episode diff --git a/lerobot/common/envs/abstract.py b/lerobot/common/envs/abstract.py index 01250d1..bca0af3 100644 --- a/lerobot/common/envs/abstract.py +++ b/lerobot/common/envs/abstract.py @@ -4,7 +4,7 @@ from typing import Optional from tensordict import TensorDict from torchrl.envs import EnvBase -from lerobot.common.utils import set_seed +from lerobot.common.utils import set_global_seed class AbstractEnv(EnvBase): @@ -67,4 +67,4 @@ class AbstractEnv(EnvBase): raise NotImplementedError("Abstract method") def _set_seed(self, seed: Optional[int]): - set_seed(seed) + set_global_seed(seed) diff --git a/lerobot/common/envs/aloha/assets/tabletop.stl b/lerobot/common/envs/aloha/assets/tabletop.stl index ab35cdf..1c17d3f 100644 Binary files a/lerobot/common/envs/aloha/assets/tabletop.stl and b/lerobot/common/envs/aloha/assets/tabletop.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl b/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl index 534c7af..ef1f3f3 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl and b/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_left.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl b/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl index d6a492c..7eb8aef 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl and b/lerobot/common/envs/aloha/assets/vx300s_10_custom_finger_right.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl b/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl index d6df86b..4c2b3a1 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl and b/lerobot/common/envs/aloha/assets/vx300s_10_gripper_finger.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl b/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl index 193014b..8a30f7c 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl and b/lerobot/common/envs/aloha/assets/vx300s_11_ar_tag.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_1_base.stl b/lerobot/common/envs/aloha/assets/vx300s_1_base.stl index 5a7efda..9198e62 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_1_base.stl and b/lerobot/common/envs/aloha/assets/vx300s_1_base.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl b/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl index dc22aa7..ab3d957 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl and b/lerobot/common/envs/aloha/assets/vx300s_2_shoulder.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl b/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl index 111c586..3d6f663 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl and b/lerobot/common/envs/aloha/assets/vx300s_3_upper_arm.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl b/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl index 8170d21..4eb249e 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl and b/lerobot/common/envs/aloha/assets/vx300s_4_upper_forearm.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl b/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl index 39581f8..34c7622 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl and b/lerobot/common/envs/aloha/assets/vx300s_5_lower_forearm.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl b/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl index ab8423e..232fabf 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl and b/lerobot/common/envs/aloha/assets/vx300s_6_wrist.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl b/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl index 043db9c..946c3c8 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl and b/lerobot/common/envs/aloha/assets/vx300s_7_gripper.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl b/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl index 36099b4..28d5bd7 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl and b/lerobot/common/envs/aloha/assets/vx300s_8_gripper_prop.stl differ diff --git a/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl b/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl index eba3caa..5201d5e 100644 Binary files a/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl and b/lerobot/common/envs/aloha/assets/vx300s_9_gripper_bar.stl differ diff --git a/lerobot/common/envs/aloha/env.py b/lerobot/common/envs/aloha/env.py index a001ca5..d38d7f0 100644 --- a/lerobot/common/envs/aloha/env.py +++ b/lerobot/common/envs/aloha/env.py @@ -29,9 +29,9 @@ from lerobot.common.envs.aloha.tasks.sim_end_effector import ( TransferCubeEndEffectorTask, ) from lerobot.common.envs.aloha.utils import sample_box_pose, sample_insertion_pose -from lerobot.common.utils import set_seed +from lerobot.common.utils import set_global_seed -_has_gym = importlib.util.find_spec("gym") is not None +_has_gym = importlib.util.find_spec("gymnasium") is not None class AlohaEnv(AbstractEnv): @@ -63,7 +63,7 @@ class AlohaEnv(AbstractEnv): def _make_env(self): if not _has_gym: - raise ImportError("Cannot import gym.") + raise ImportError("Cannot import gymnasium.") if not self.from_pixels: raise NotImplementedError() @@ -290,7 +290,7 @@ class AlohaEnv(AbstractEnv): ) def _set_seed(self, seed: Optional[int]): - set_seed(seed) + set_global_seed(seed) # TODO(rcadene): seed the env # self._env.seed(seed) logging.warning("Aloha env is not seeded") diff --git a/lerobot/common/envs/factory.py b/lerobot/common/envs/factory.py index 06c7c43..855e073 100644 --- a/lerobot/common/envs/factory.py +++ b/lerobot/common/envs/factory.py @@ -17,7 +17,7 @@ def make_env(cfg, transform=None): } if cfg.env.name == "simxarm": - from lerobot.common.envs.simxarm import SimxarmEnv + from lerobot.common.envs.simxarm.env import SimxarmEnv kwargs["task"] = cfg.env.task clsfunc = SimxarmEnv diff --git a/lerobot/common/envs/pusht/env.py b/lerobot/common/envs/pusht/env.py index 070c718..ca39bf4 100644 --- a/lerobot/common/envs/pusht/env.py +++ b/lerobot/common/envs/pusht/env.py @@ -16,9 +16,9 @@ from torchrl.data.tensor_specs import ( from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform from lerobot.common.envs.abstract import AbstractEnv -from lerobot.common.utils import set_seed +from lerobot.common.utils import set_global_seed -_has_gym = importlib.util.find_spec("gym") is not None +_has_gym = importlib.util.find_spec("gymnasium") is not None class PushtEnv(AbstractEnv): @@ -50,7 +50,7 @@ class PushtEnv(AbstractEnv): def _make_env(self): if not _has_gym: - raise ImportError("Cannot import gym.") + raise ImportError("Cannot import gymnasium.") # TODO(rcadene) (PushTEnv is similar to PushTImageEnv, but without the image rendering, it's faster to iterate on) # from lerobot.common.envs.pusht.pusht_env import PushTEnv @@ -238,6 +238,6 @@ class PushtEnv(AbstractEnv): def _set_seed(self, seed: Optional[int]): # Set global seed. - set_seed(seed) + set_global_seed(seed) # Set PushTImageEnv seed as it relies on it's own internal _seed attribute. self._env.seed(seed) diff --git a/lerobot/common/envs/pusht/pusht_env.py b/lerobot/common/envs/pusht/pusht_env.py index 186f9e3..6ef70ae 100644 --- a/lerobot/common/envs/pusht/pusht_env.py +++ b/lerobot/common/envs/pusht/pusht_env.py @@ -1,14 +1,14 @@ import collections import cv2 -import gym +import gymnasium as gym import numpy as np import pygame import pymunk import pymunk.pygame_util import shapely.geometry as sg import skimage.transform as st -from gym import spaces +from gymnasium import spaces from pymunk.vec2d import Vec2d from lerobot.common.envs.pusht.pymunk_override import DrawOptions diff --git a/lerobot/common/envs/pusht/pusht_image_env.py b/lerobot/common/envs/pusht/pusht_image_env.py index 4981eb6..6547835 100644 --- a/lerobot/common/envs/pusht/pusht_image_env.py +++ b/lerobot/common/envs/pusht/pusht_image_env.py @@ -1,5 +1,5 @@ import numpy as np -from gym import spaces +from gymnasium import spaces from lerobot.common.envs.pusht.pusht_env import PushTEnv diff --git a/lerobot/common/envs/simxarm.py b/lerobot/common/envs/simxarm/env.py similarity index 92% rename from lerobot/common/envs/simxarm.py rename to lerobot/common/envs/simxarm/env.py index eac3666..f3c358d 100644 --- a/lerobot/common/envs/simxarm.py +++ b/lerobot/common/envs/simxarm/env.py @@ -1,4 +1,5 @@ import importlib +import logging from collections import deque from typing import Optional @@ -15,12 +16,11 @@ from torchrl.data.tensor_specs import ( from torchrl.envs.libs.gym import _gym_to_torchrl_spec_transform from lerobot.common.envs.abstract import AbstractEnv -from lerobot.common.utils import set_seed +from lerobot.common.utils import set_global_seed MAX_NUM_ACTIONS = 4 -_has_gym = importlib.util.find_spec("gym") is not None -_has_simxarm = importlib.util.find_spec("simxarm") is not None and _has_gym +_has_gym = importlib.util.find_spec("gymnasium") is not None class SimxarmEnv(AbstractEnv): @@ -49,13 +49,12 @@ class SimxarmEnv(AbstractEnv): ) def _make_env(self): - if not _has_simxarm: - raise ImportError("Cannot import simxarm.") if not _has_gym: - raise ImportError("Cannot import gym.") + raise ImportError("Cannot import gymnasium.") - import gym - from simxarm import TASKS + import gymnasium + + from lerobot.common.envs.simxarm.simxarm import TASKS if self.task not in TASKS: raise ValueError(f"Unknown task {self.task}. Must be one of {list(TASKS.keys())}") @@ -63,7 +62,7 @@ class SimxarmEnv(AbstractEnv): self._env = TASKS[self.task]["env"]() num_actions = len(TASKS[self.task]["action_space"]) - self._action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(num_actions,)) + self._action_space = gymnasium.spaces.Box(low=-1.0, high=1.0, shape=(num_actions,)) self._action_padding = np.zeros((MAX_NUM_ACTIONS - num_actions), dtype=np.float32) if "w" not in TASKS[self.task]["action_space"]: self._action_padding[-1] = 1.0 @@ -84,7 +83,7 @@ class SimxarmEnv(AbstractEnv): else: obs = {"state": torch.tensor(raw_obs["observation"], dtype=torch.float32)} - obs = TensorDict(obs, batch_size=[]) + # obs = TensorDict(obs, batch_size=[]) return obs def _reset(self, tensordict: Optional[TensorDict] = None): @@ -229,5 +228,7 @@ class SimxarmEnv(AbstractEnv): ) def _set_seed(self, seed: Optional[int]): - set_seed(seed) - self._env.seed(seed) + set_global_seed(seed) + self._seed = seed + # TODO(aliberts): change self._reset so that it takes in a seed value + logging.warning("simxarm env is not properly seeded") diff --git a/lerobot/common/envs/simxarm/simxarm/__init__.py b/lerobot/common/envs/simxarm/simxarm/__init__.py new file mode 100644 index 0000000..903d604 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/__init__.py @@ -0,0 +1,166 @@ +from collections import OrderedDict, deque + +import gymnasium as gym +import numpy as np +from gymnasium.wrappers import TimeLimit + +from lerobot.common.envs.simxarm.simxarm.tasks.base import Base as Base +from lerobot.common.envs.simxarm.simxarm.tasks.lift import Lift +from lerobot.common.envs.simxarm.simxarm.tasks.peg_in_box import PegInBox +from lerobot.common.envs.simxarm.simxarm.tasks.push import Push +from lerobot.common.envs.simxarm.simxarm.tasks.reach import Reach + +TASKS = OrderedDict( + ( + ( + "reach", + { + "env": Reach, + "action_space": "xyz", + "episode_length": 50, + "description": "Reach a target location with the end effector", + }, + ), + ( + "push", + { + "env": Push, + "action_space": "xyz", + "episode_length": 50, + "description": "Push a cube to a target location", + }, + ), + ( + "peg_in_box", + { + "env": PegInBox, + "action_space": "xyz", + "episode_length": 50, + "description": "Insert a peg into a box", + }, + ), + ( + "lift", + { + "env": Lift, + "action_space": "xyzw", + "episode_length": 50, + "description": "Lift a cube above a height threshold", + }, + ), + ) +) + + +class SimXarmWrapper(gym.Wrapper): + """ + A wrapper for the SimXarm environments. This wrapper is used to + convert the action and observation spaces to the correct format. + """ + + def __init__(self, env, task, obs_mode, image_size, action_repeat, frame_stack=1, channel_last=False): + super().__init__(env) + self._env = env + self.obs_mode = obs_mode + self.image_size = image_size + self.action_repeat = action_repeat + self.frame_stack = frame_stack + self._frames = deque([], maxlen=frame_stack) + self.channel_last = channel_last + self._max_episode_steps = task["episode_length"] // action_repeat + + image_shape = ( + (image_size, image_size, 3 * frame_stack) + if channel_last + else (3 * frame_stack, image_size, image_size) + ) + if obs_mode == "state": + self.observation_space = env.observation_space["observation"] + elif obs_mode == "rgb": + self.observation_space = gym.spaces.Box(low=0, high=255, shape=image_shape, dtype=np.uint8) + elif obs_mode == "all": + self.observation_space = gym.spaces.Dict( + state=gym.spaces.Box(low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32), + rgb=gym.spaces.Box(low=0, high=255, shape=image_shape, dtype=np.uint8), + ) + else: + raise ValueError(f"Unknown obs_mode {obs_mode}. Must be one of [rgb, all, state]") + self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(len(task["action_space"]),)) + self.action_padding = np.zeros(4 - len(task["action_space"]), dtype=np.float32) + if "w" not in task["action_space"]: + self.action_padding[-1] = 1.0 + + def _render_obs(self): + obs = self.render(mode="rgb_array", width=self.image_size, height=self.image_size) + if not self.channel_last: + obs = obs.transpose(2, 0, 1) + return obs.copy() + + def _update_frames(self, reset=False): + pixels = self._render_obs() + self._frames.append(pixels) + if reset: + for _ in range(1, self.frame_stack): + self._frames.append(pixels) + assert len(self._frames) == self.frame_stack + + def transform_obs(self, obs, reset=False): + if self.obs_mode == "state": + return obs["observation"] + elif self.obs_mode == "rgb": + self._update_frames(reset=reset) + rgb_obs = np.concatenate(list(self._frames), axis=-1 if self.channel_last else 0) + return rgb_obs + elif self.obs_mode == "all": + self._update_frames(reset=reset) + rgb_obs = np.concatenate(list(self._frames), axis=-1 if self.channel_last else 0) + return OrderedDict((("rgb", rgb_obs), ("state", self.robot_state))) + else: + raise ValueError(f"Unknown obs_mode {self.obs_mode}. Must be one of [rgb, all, state]") + + def reset(self): + return self.transform_obs(self._env.reset(), reset=True) + + def step(self, action): + action = np.concatenate([action, self.action_padding]) + reward = 0.0 + for _ in range(self.action_repeat): + obs, r, done, info = self._env.step(action) + reward += r + return self.transform_obs(obs), reward, done, info + + def render(self, mode="rgb_array", width=384, height=384, **kwargs): + return self._env.render(mode, width=width, height=height) + + @property + def state(self): + return self._env.robot_state + + +def make(task, obs_mode="state", image_size=84, action_repeat=1, frame_stack=1, channel_last=False, seed=0): + """ + Create a new environment. + Args: + task (str): The task to create an environment for. Must be one of: + - 'reach' + - 'push' + - 'peg-in-box' + - 'lift' + obs_mode (str): The observation mode to use. Must be one of: + - 'state': Only state observations + - 'rgb': RGB images + - 'all': RGB images and state observations + image_size (int): The size of the image observations + action_repeat (int): The number of times to repeat the action + seed (int): The random seed to use + Returns: + gym.Env: The environment + """ + if task not in TASKS: + raise ValueError(f"Unknown task {task}. Must be one of {list(TASKS.keys())}") + env = TASKS[task]["env"]() + env = TimeLimit(env, TASKS[task]["episode_length"]) + env = SimXarmWrapper(env, TASKS[task], obs_mode, image_size, action_repeat, frame_stack, channel_last) + env.seed(seed) + + return env diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/__init__.py b/lerobot/common/envs/simxarm/simxarm/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/lift.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/lift.xml new file mode 100644 index 0000000..92231f9 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/lift.xml @@ -0,0 +1,53 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/base_link.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/base_link.stl new file mode 100644 index 0000000..f1f5295 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/base_link.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21fb81ae7fba19e3c6b2d2ca60c8051712ba273357287eb5a397d92d61c7a736 +size 1211434 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner.stl new file mode 100644 index 0000000..6cb8894 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:be68ce180d11630a667a5f37f4dffcc3feebe4217d4bb3912c813b6d9ca3ec66 +size 3284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner2.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner2.stl new file mode 100644 index 0000000..dab55ef --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_inner2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2c6448552bf6b1c4f17334d686a5320ce051bcdfe31431edf69303d8a570d1de +size 3284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_outer.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_outer.stl new file mode 100644 index 0000000..21cf11f --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/block_outer.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:748b9e197e6521914f18d1f6383a36f211136b3f33f2ad2a8c11b9f921c2cf86 +size 6284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_finger.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_finger.stl new file mode 100644 index 0000000..6bf4e50 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_finger.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a44756eb72f9c214cb37e61dc209cd7073fdff3e4271a7423476ef6fd090d2d4 +size 242684 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_inner_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_inner_knuckle.stl new file mode 100644 index 0000000..817c7e1 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_inner_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e8e48692ad26837bb3d6a97582c89784d09948fc09bfe4e5a59017859ff04dac +size 366284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_outer_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_outer_knuckle.stl new file mode 100644 index 0000000..010c0f3 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/left_outer_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:501665812b08d67e764390db781e839adc6896a9540301d60adf606f57648921 +size 22284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link1.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link1.stl new file mode 100644 index 0000000..f2b676f --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link1.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34b541122df84d2ef5fcb91b715eb19659dc15ad8d44a191dde481f780265636 +size 184184 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link2.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link2.stl new file mode 100644 index 0000000..bf93580 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link2.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:61e641cd47c169ecef779683332e00e4914db729bf02dfb61bfbe69351827455 +size 225584 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link3.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link3.stl new file mode 100644 index 0000000..d316d23 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link3.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9e2798e7946dd70046c95455d5ba96392d0b54a6069caba91dc4ca66e1379b42 +size 237084 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link4.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link4.stl new file mode 100644 index 0000000..f6d5fe9 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link4.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c757fee95f873191a0633c355c07a360032960771cabbd7593a6cdb0f1ffb089 +size 243684 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link5.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link5.stl new file mode 100644 index 0000000..e037b8b --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link5.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:715ad5787c5dab57589937fd47289882707b5e1eb997e340d567785b02f4ec90 +size 229084 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link6.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link6.stl new file mode 100644 index 0000000..198c530 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link6.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85b320aa420497827223d16d492bba8de091173374e361396fc7a5dad7bdb0cb +size 399384 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link7.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link7.stl new file mode 100644 index 0000000..ce9a39a --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link7.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:97115d848fbf802cb770cd9be639ae2af993103b9d9bbb0c50c943c738a36f18 +size 231684 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link_base.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link_base.stl new file mode 100644 index 0000000..110b953 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/link_base.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f6fcbc18258090eb56c21cfb17baa5ae43abc98b1958cd366f3a73b9898fc7f0 +size 2106184 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_finger.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_finger.stl new file mode 100644 index 0000000..03f26e9 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_finger.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c5dee87c7f37baf554b8456ebfe0b3e8ed0b22b8938bd1add6505c2ad6d32c7d +size 242684 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_inner_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_inner_knuckle.stl new file mode 100644 index 0000000..8586f34 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_inner_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b41dd2c2c550281bf78d7cc6fa117b14786700e5c453560a0cb5fd6dfa0ffb3e +size 366284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_outer_knuckle.stl b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_outer_knuckle.stl new file mode 100644 index 0000000..ae7afc2 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/mesh/right_outer_knuckle.stl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:75ca1107d0a42a0f03802a9a49cab48419b31851ee8935f8f1ca06be1c1c91e8 +size 22284 diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/peg_in_box.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/peg_in_box.xml new file mode 100644 index 0000000..0f85459 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/peg_in_box.xml @@ -0,0 +1,74 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/push.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/push.xml new file mode 100644 index 0000000..42a78c8 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/push.xml @@ -0,0 +1,54 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/reach.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/reach.xml new file mode 100644 index 0000000..ded6d20 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/reach.xml @@ -0,0 +1,48 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/shared.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/shared.xml new file mode 100644 index 0000000..ee56f8f --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/shared.xml @@ -0,0 +1,51 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/assets/xarm.xml b/lerobot/common/envs/simxarm/simxarm/tasks/assets/xarm.xml new file mode 100644 index 0000000..023474d --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/assets/xarm.xml @@ -0,0 +1,88 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/base.py b/lerobot/common/envs/simxarm/simxarm/tasks/base.py new file mode 100644 index 0000000..b937b29 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/base.py @@ -0,0 +1,145 @@ +import os + +import mujoco +import numpy as np +from gymnasium_robotics.envs import robot_env + +from lerobot.common.envs.simxarm.simxarm.tasks import mocap + + +class Base(robot_env.MujocoRobotEnv): + """ + Superclass for all simxarm environments. + Args: + xml_name (str): name of the xml environment file + gripper_rotation (list): initial rotation of the gripper (given as a quaternion) + """ + + def __init__(self, xml_name, gripper_rotation=None): + if gripper_rotation is None: + gripper_rotation = [0, 1, 0, 0] + self.gripper_rotation = np.array(gripper_rotation, dtype=np.float32) + self.center_of_table = np.array([1.655, 0.3, 0.63625]) + self.max_z = 1.2 + self.min_z = 0.2 + super().__init__( + model_path=os.path.join(os.path.dirname(__file__), "assets", xml_name + ".xml"), + n_substeps=20, + n_actions=4, + initial_qpos={}, + ) + + @property + def dt(self): + return self.n_substeps * self.model.opt.timestep + + @property + def eef(self): + return self._utils.get_site_xpos(self.model, self.data, "grasp") + + @property + def obj(self): + return self._utils.get_site_xpos(self.model, self.data, "object_site") + + @property + def robot_state(self): + gripper_angle = self._utils.get_joint_qpos(self.model, self.data, "right_outer_knuckle_joint") + return np.concatenate([self.eef, gripper_angle]) + + def is_success(self): + return NotImplementedError() + + def get_reward(self): + raise NotImplementedError() + + def _sample_goal(self): + raise NotImplementedError() + + def get_obs(self): + return self._get_obs() + + def _step_callback(self): + self._mujoco.mj_forward(self.model, self.data) + + def _limit_gripper(self, gripper_pos, pos_ctrl): + if gripper_pos[0] > self.center_of_table[0] - 0.105 + 0.15: + pos_ctrl[0] = min(pos_ctrl[0], 0) + if gripper_pos[0] < self.center_of_table[0] - 0.105 - 0.3: + pos_ctrl[0] = max(pos_ctrl[0], 0) + if gripper_pos[1] > self.center_of_table[1] + 0.3: + pos_ctrl[1] = min(pos_ctrl[1], 0) + if gripper_pos[1] < self.center_of_table[1] - 0.3: + pos_ctrl[1] = max(pos_ctrl[1], 0) + if gripper_pos[2] > self.max_z: + pos_ctrl[2] = min(pos_ctrl[2], 0) + if gripper_pos[2] < self.min_z: + pos_ctrl[2] = max(pos_ctrl[2], 0) + return pos_ctrl + + def _apply_action(self, action): + assert action.shape == (4,) + action = action.copy() + pos_ctrl, gripper_ctrl = action[:3], action[3] + pos_ctrl = self._limit_gripper( + self._utils.get_site_xpos(self.model, self.data, "grasp"), pos_ctrl + ) * (1 / self.n_substeps) + gripper_ctrl = np.array([gripper_ctrl, gripper_ctrl]) + mocap.apply_action( + self.model, + self._model_names, + self.data, + np.concatenate([pos_ctrl, self.gripper_rotation, gripper_ctrl]), + ) + + def _render_callback(self): + self._mujoco.mj_forward(self.model, self.data) + + def _reset_sim(self): + self.data.time = self.initial_time + self.data.qpos[:] = np.copy(self.initial_qpos) + self.data.qvel[:] = np.copy(self.initial_qvel) + self._sample_goal() + self._mujoco.mj_step(self.model, self.data, nstep=10) + return True + + def _set_gripper(self, gripper_pos, gripper_rotation): + self._utils.set_mocap_pos(self.model, self.data, "robot0:mocap", gripper_pos) + self._utils.set_mocap_quat(self.model, self.data, "robot0:mocap", gripper_rotation) + self._utils.set_joint_qpos(self.model, self.data, "right_outer_knuckle_joint", 0) + self.data.qpos[10] = 0.0 + self.data.qpos[12] = 0.0 + + def _env_setup(self, initial_qpos): + for name, value in initial_qpos.items(): + self.data.set_joint_qpos(name, value) + mocap.reset(self.model, self.data) + mujoco.mj_forward(self.model, self.data) + self._sample_goal() + mujoco.mj_forward(self.model, self.data) + + def reset(self): + self._reset_sim() + return self._get_obs() + + def step(self, action): + assert action.shape == (4,) + assert self.action_space.contains(action), "{!r} ({}) invalid".format(action, type(action)) + self._apply_action(action) + self._mujoco.mj_step(self.model, self.data, nstep=2) + self._step_callback() + obs = self._get_obs() + reward = self.get_reward() + done = False + info = {"is_success": self.is_success(), "success": self.is_success()} + return obs, reward, done, info + + def render(self, mode="rgb_array", width=384, height=384): + self._render_callback() + # HACK + self.model.vis.global_.offwidth = width + self.model.vis.global_.offheight = height + return self.mujoco_renderer.render(mode) + + def close(self): + if self.mujoco_renderer is not None: + self.mujoco_renderer.close() diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/lift.py b/lerobot/common/envs/simxarm/simxarm/tasks/lift.py new file mode 100644 index 0000000..0b11196 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/lift.py @@ -0,0 +1,100 @@ +import numpy as np + +from lerobot.common.envs.simxarm.simxarm import Base + + +class Lift(Base): + def __init__(self): + self._z_threshold = 0.15 + super().__init__("lift") + + @property + def z_target(self): + return self._init_z + self._z_threshold + + def is_success(self): + return self.obj[2] >= self.z_target + + def get_reward(self): + reach_dist = np.linalg.norm(self.obj - self.eef) + reach_dist_xy = np.linalg.norm(self.obj[:-1] - self.eef[:-1]) + pick_completed = self.obj[2] >= (self.z_target - 0.01) + obj_dropped = (self.obj[2] < (self._init_z + 0.005)) and (reach_dist > 0.02) + + # Reach + if reach_dist < 0.05: + reach_reward = -reach_dist + max(self._action[-1], 0) / 50 + elif reach_dist_xy < 0.05: + reach_reward = -reach_dist + else: + z_bonus = np.linalg.norm(np.linalg.norm(self.obj[-1] - self.eef[-1])) + reach_reward = -reach_dist - 2 * z_bonus + + # Pick + if pick_completed and not obj_dropped: + pick_reward = self.z_target + elif (reach_dist < 0.1) and (self.obj[2] > (self._init_z + 0.005)): + pick_reward = min(self.z_target, self.obj[2]) + else: + pick_reward = 0 + + return reach_reward / 100 + pick_reward + + def _get_obs(self): + eef_velp = self._utils.get_site_xvelp(self.model, self.data, "grasp") * self.dt + gripper_angle = self._utils.get_joint_qpos(self.model, self.data, "right_outer_knuckle_joint") + eef = self.eef - self.center_of_table + + obj = self.obj - self.center_of_table + obj_rot = self._utils.get_joint_qpos(self.model, self.data, "object_joint0")[-4:] + obj_velp = self._utils.get_site_xvelp(self.model, self.data, "object_site") * self.dt + obj_velr = self._utils.get_site_xvelr(self.model, self.data, "object_site") * self.dt + + obs = np.concatenate( + [ + eef, + eef_velp, + obj, + obj_rot, + obj_velp, + obj_velr, + eef - obj, + np.array( + [ + np.linalg.norm(eef - obj), + np.linalg.norm(eef[:-1] - obj[:-1]), + self.z_target, + self.z_target - obj[-1], + self.z_target - eef[-1], + ] + ), + gripper_angle, + ], + axis=0, + ) + return {"observation": obs, "state": eef, "achieved_goal": eef, "desired_goal": eef} + + def _sample_goal(self): + # Gripper + gripper_pos = np.array([1.280, 0.295, 0.735]) + self.np_random.uniform(-0.05, 0.05, size=3) + super()._set_gripper(gripper_pos, self.gripper_rotation) + + # Object + object_pos = self.center_of_table - np.array([0.15, 0.10, 0.07]) + object_pos[0] += self.np_random.uniform(-0.05, 0.05, size=1) + object_pos[1] += self.np_random.uniform(-0.05, 0.05, size=1) + object_qpos = self._utils.get_joint_qpos(self.model, self.data, "object_joint0") + object_qpos[:3] = object_pos + self._utils.set_joint_qpos(self.model, self.data, "object_joint0", object_qpos) + self._init_z = object_pos[2] + + # Goal + return object_pos + np.array([0, 0, self._z_threshold]) + + def reset(self): + self._action = np.zeros(4) + return super().reset() + + def step(self, action): + self._action = action.copy() + return super().step(action) diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/mocap.py b/lerobot/common/envs/simxarm/simxarm/tasks/mocap.py new file mode 100644 index 0000000..4295bf1 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/mocap.py @@ -0,0 +1,67 @@ +# import mujoco_py +import mujoco +import numpy as np + + +def apply_action(model, model_names, data, action): + if model.nmocap > 0: + pos_action, gripper_action = np.split(action, (model.nmocap * 7,)) + if data.ctrl is not None: + for i in range(gripper_action.shape[0]): + data.ctrl[i] = gripper_action[i] + pos_action = pos_action.reshape(model.nmocap, 7) + pos_delta, quat_delta = pos_action[:, :3], pos_action[:, 3:] + reset_mocap2body_xpos(model, model_names, data) + data.mocap_pos[:] = data.mocap_pos + pos_delta + data.mocap_quat[:] = data.mocap_quat + quat_delta + + +def reset(model, data): + if model.nmocap > 0 and model.eq_data is not None: + for i in range(model.eq_data.shape[0]): + # if sim.model.eq_type[i] == mujoco_py.const.EQ_WELD: + if model.eq_type[i] == mujoco.mjtEq.mjEQ_WELD: + # model.eq_data[i, :] = np.array([0., 0., 0., 1., 0., 0., 0.]) + model.eq_data[i, :] = np.array( + [ + 0.0, + 0.0, + 0.0, + 1.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + ] + ) + # sim.forward() + mujoco.mj_forward(model, data) + + +def reset_mocap2body_xpos(model, model_names, data): + if model.eq_type is None or model.eq_obj1id is None or model.eq_obj2id is None: + return + + # For all weld constraints + for eq_type, obj1_id, obj2_id in zip(model.eq_type, model.eq_obj1id, model.eq_obj2id, strict=False): + # if eq_type != mujoco_py.const.EQ_WELD: + if eq_type != mujoco.mjtEq.mjEQ_WELD: + continue + # body2 = model.body_id2name(obj2_id) + body2 = model_names.body_id2name[obj2_id] + if body2 == "B0" or body2 == "B9" or body2 == "B1": + continue + mocap_id = model.body_mocapid[obj1_id] + if mocap_id != -1: + # obj1 is the mocap, obj2 is the welded body + body_idx = obj2_id + else: + # obj2 is the mocap, obj1 is the welded body + mocap_id = model.body_mocapid[obj2_id] + body_idx = obj1_id + assert mocap_id != -1 + data.mocap_pos[mocap_id][:] = data.xpos[body_idx] + data.mocap_quat[mocap_id][:] = data.xquat[body_idx] diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/peg_in_box.py b/lerobot/common/envs/simxarm/simxarm/tasks/peg_in_box.py new file mode 100644 index 0000000..42e4152 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/peg_in_box.py @@ -0,0 +1,86 @@ +import numpy as np + +from lerobot.common.envs.simxarm.simxarm import Base + + +class PegInBox(Base): + def __init__(self): + super().__init__("peg_in_box") + + def _reset_sim(self): + self._act_magnitude = 0 + super()._reset_sim() + for _ in range(10): + self._apply_action(np.array([0, 0, 0, 1], dtype=np.float32)) + self.sim.step() + + @property + def box(self): + return self.sim.data.get_site_xpos("box_site") + + def is_success(self): + return np.linalg.norm(self.obj - self.box) <= 0.05 + + def get_reward(self): + dist_xy = np.linalg.norm(self.obj[:2] - self.box[:2]) + dist_xyz = np.linalg.norm(self.obj - self.box) + return float(dist_xy <= 0.045) * (2 - 6 * dist_xyz) - 0.2 * np.square(self._act_magnitude) - dist_xy + + def _get_obs(self): + eef_velp = self.sim.data.get_site_xvelp("grasp") * self.dt + gripper_angle = self.sim.data.get_joint_qpos("right_outer_knuckle_joint") + eef, box = self.eef - self.center_of_table, self.box - self.center_of_table + + obj = self.obj - self.center_of_table + obj_rot = self.sim.data.get_joint_qpos("object_joint0")[-4:] + obj_velp = self.sim.data.get_site_xvelp("object_site") * self.dt + obj_velr = self.sim.data.get_site_xvelr("object_site") * self.dt + + obs = np.concatenate( + [ + eef, + eef_velp, + box, + obj, + obj_rot, + obj_velp, + obj_velr, + eef - box, + eef - obj, + obj - box, + np.array( + [ + np.linalg.norm(eef - box), + np.linalg.norm(eef - obj), + np.linalg.norm(obj - box), + gripper_angle, + ] + ), + ], + axis=0, + ) + return {"observation": obs, "state": eef, "achieved_goal": eef, "desired_goal": box} + + def _sample_goal(self): + # Gripper + gripper_pos = np.array([1.280, 0.295, 0.9]) + self.np_random.uniform(-0.05, 0.05, size=3) + super()._set_gripper(gripper_pos, self.gripper_rotation) + + # Object + object_pos = gripper_pos - np.array([0, 0, 0.06]) + self.np_random.uniform(-0.005, 0.005, size=3) + object_qpos = self.sim.data.get_joint_qpos("object_joint0") + object_qpos[:3] = object_pos + self.sim.data.set_joint_qpos("object_joint0", object_qpos) + + # Box + box_pos = np.array([1.61, 0.18, 0.58]) + box_pos[:2] += self.np_random.uniform(-0.11, 0.11, size=2) + box_qpos = self.sim.data.get_joint_qpos("box_joint0") + box_qpos[:3] = box_pos + self.sim.data.set_joint_qpos("box_joint0", box_qpos) + + return self.box + + def step(self, action): + self._act_magnitude = np.linalg.norm(action[:3]) + return super().step(action) diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/push.py b/lerobot/common/envs/simxarm/simxarm/tasks/push.py new file mode 100644 index 0000000..36c4a55 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/push.py @@ -0,0 +1,78 @@ +import numpy as np + +from lerobot.common.envs.simxarm.simxarm import Base + + +class Push(Base): + def __init__(self): + super().__init__("push") + + def _reset_sim(self): + self._act_magnitude = 0 + super()._reset_sim() + + def is_success(self): + return np.linalg.norm(self.obj - self.goal) <= 0.05 + + def get_reward(self): + dist = np.linalg.norm(self.obj - self.goal) + penalty = self._act_magnitude**2 + return -(dist + 0.15 * penalty) + + def _get_obs(self): + eef_velp = self.sim.data.get_site_xvelp("grasp") * self.dt + gripper_angle = self.sim.data.get_joint_qpos("right_outer_knuckle_joint") + eef, goal = self.eef - self.center_of_table, self.goal - self.center_of_table + + obj = self.obj - self.center_of_table + obj_rot = self.sim.data.get_joint_qpos("object_joint0")[-4:] + obj_velp = self.sim.data.get_site_xvelp("object_site") * self.dt + obj_velr = self.sim.data.get_site_xvelr("object_site") * self.dt + + obs = np.concatenate( + [ + eef, + eef_velp, + goal, + obj, + obj_rot, + obj_velp, + obj_velr, + eef - goal, + eef - obj, + obj - goal, + np.array( + [ + np.linalg.norm(eef - goal), + np.linalg.norm(eef - obj), + np.linalg.norm(obj - goal), + gripper_angle, + ] + ), + ], + axis=0, + ) + return {"observation": obs, "state": eef, "achieved_goal": eef, "desired_goal": goal} + + def _sample_goal(self): + # Gripper + gripper_pos = np.array([1.280, 0.295, 0.735]) + self.np_random.uniform(-0.05, 0.05, size=3) + super()._set_gripper(gripper_pos, self.gripper_rotation) + + # Object + object_pos = self.center_of_table - np.array([0.25, 0, 0.07]) + object_pos[0] += self.np_random.uniform(-0.08, 0.08, size=1) + object_pos[1] += self.np_random.uniform(-0.08, 0.08, size=1) + object_qpos = self.sim.data.get_joint_qpos("object_joint0") + object_qpos[:3] = object_pos + self.sim.data.set_joint_qpos("object_joint0", object_qpos) + + # Goal + self.goal = np.array([1.600, 0.200, 0.545]) + self.goal[:2] += self.np_random.uniform(-0.1, 0.1, size=2) + self.sim.model.site_pos[self.sim.model.site_name2id("target0")] = self.goal + return self.goal + + def step(self, action): + self._act_magnitude = np.linalg.norm(action[:3]) + return super().step(action) diff --git a/lerobot/common/envs/simxarm/simxarm/tasks/reach.py b/lerobot/common/envs/simxarm/simxarm/tasks/reach.py new file mode 100644 index 0000000..941a586 --- /dev/null +++ b/lerobot/common/envs/simxarm/simxarm/tasks/reach.py @@ -0,0 +1,44 @@ +import numpy as np + +from lerobot.common.envs.simxarm.simxarm import Base + + +class Reach(Base): + def __init__(self): + super().__init__("reach") + + def _reset_sim(self): + self._act_magnitude = 0 + super()._reset_sim() + + def is_success(self): + return np.linalg.norm(self.eef - self.goal) <= 0.05 + + def get_reward(self): + dist = np.linalg.norm(self.eef - self.goal) + penalty = self._act_magnitude**2 + return -(dist + 0.15 * penalty) + + def _get_obs(self): + eef_velp = self.sim.data.get_site_xvelp("grasp") * self.dt + gripper_angle = self.sim.data.get_joint_qpos("right_outer_knuckle_joint") + eef, goal = self.eef - self.center_of_table, self.goal - self.center_of_table + obs = np.concatenate( + [eef, eef_velp, goal, eef - goal, np.array([np.linalg.norm(eef - goal), gripper_angle])], axis=0 + ) + return {"observation": obs, "state": eef, "achieved_goal": eef, "desired_goal": goal} + + def _sample_goal(self): + # Gripper + gripper_pos = np.array([1.280, 0.295, 0.735]) + self.np_random.uniform(-0.05, 0.05, size=3) + super()._set_gripper(gripper_pos, self.gripper_rotation) + + # Goal + self.goal = np.array([1.550, 0.287, 0.580]) + self.goal[:2] += self.np_random.uniform(-0.125, 0.125, size=2) + self.sim.model.site_pos[self.sim.model.site_name2id("target0")] = self.goal + return self.goal + + def step(self, action): + self._act_magnitude = np.linalg.norm(action[:3]) + return super().step(action) diff --git a/lerobot/common/policies/act/position_encoding.py b/lerobot/common/policies/act/position_encoding.py index 94e862f..63bb484 100644 --- a/lerobot/common/policies/act/position_encoding.py +++ b/lerobot/common/policies/act/position_encoding.py @@ -1,6 +1,7 @@ """ Various positional encodings for the transformer. """ + import math import torch diff --git a/lerobot/common/policies/act/transformer.py b/lerobot/common/policies/act/transformer.py index b2bd368..20cfc81 100644 --- a/lerobot/common/policies/act/transformer.py +++ b/lerobot/common/policies/act/transformer.py @@ -6,6 +6,7 @@ Copy-paste from torch.nn.Transformer with modifications: * extra LN at the end of encoder is removed * decoder returns a stack of activations from all decoding layers """ + import copy from typing import Optional diff --git a/lerobot/common/policies/act/utils.py b/lerobot/common/policies/act/utils.py index 2ce9209..0d93583 100644 --- a/lerobot/common/policies/act/utils.py +++ b/lerobot/common/policies/act/utils.py @@ -3,6 +3,7 @@ Misc functions, including distributed helpers. Mostly copy-paste from torchvision references. """ + import datetime import os import pickle diff --git a/lerobot/common/policies/diffusion/model/tensor_utils.py b/lerobot/common/policies/diffusion/model/tensor_utils.py index 0801e29..df9a568 100644 --- a/lerobot/common/policies/diffusion/model/tensor_utils.py +++ b/lerobot/common/policies/diffusion/model/tensor_utils.py @@ -2,6 +2,7 @@ A collection of utilities for working with nested tensor structures consisting of numpy arrays and torch tensors. """ + import collections import numpy as np diff --git a/lerobot/common/utils.py b/lerobot/common/utils.py index a56543b..2af1d96 100644 --- a/lerobot/common/utils.py +++ b/lerobot/common/utils.py @@ -26,7 +26,7 @@ def get_safe_torch_device(cfg_device: str, log: bool = False) -> torch.device: return device -def set_seed(seed): +def set_global_seed(seed): """Set seed for reproducibility.""" random.seed(seed) np.random.seed(seed) diff --git a/lerobot/configs/default.yaml b/lerobot/configs/default.yaml index 2dc313e..01a02a7 100644 --- a/lerobot/configs/default.yaml +++ b/lerobot/configs/default.yaml @@ -27,6 +27,7 @@ fps: ??? offline_prioritized_sampler: true n_action_steps: ??? +n_obs_steps: ??? env: ??? policy: ??? diff --git a/lerobot/configs/policy/tdmpc.yaml b/lerobot/configs/policy/tdmpc.yaml index 16b7018..ff0e6b0 100644 --- a/lerobot/configs/policy/tdmpc.yaml +++ b/lerobot/configs/policy/tdmpc.yaml @@ -1,6 +1,7 @@ # @package _global_ n_action_steps: 1 +n_obs_steps: 1 policy: name: tdmpc diff --git a/lerobot/scripts/eval.py b/lerobot/scripts/eval.py index b3d107a..e30cd9d 100644 --- a/lerobot/scripts/eval.py +++ b/lerobot/scripts/eval.py @@ -50,7 +50,7 @@ from lerobot.common.envs.factory import make_env from lerobot.common.logger import log_output_dir from lerobot.common.policies.abstract import AbstractPolicy from lerobot.common.policies.factory import make_policy -from lerobot.common.utils import get_safe_torch_device, init_logging, set_seed +from lerobot.common.utils import get_safe_torch_device, init_logging, set_global_seed def write_video(video_path, stacked_frames, fps): @@ -188,7 +188,7 @@ def eval(cfg: dict, out_dir=None, stats_path=None): torch.backends.cudnn.benchmark = True torch.backends.cuda.matmul.allow_tf32 = True - set_seed(cfg.seed) + set_global_seed(cfg.seed) log_output_dir(out_dir) diff --git a/lerobot/scripts/train.py b/lerobot/scripts/train.py index cf71ad2..18c3715 100644 --- a/lerobot/scripts/train.py +++ b/lerobot/scripts/train.py @@ -12,7 +12,7 @@ from lerobot.common.datasets.factory import make_offline_buffer from lerobot.common.envs.factory import make_env from lerobot.common.logger import Logger, log_output_dir from lerobot.common.policies.factory import make_policy -from lerobot.common.utils import format_big_number, get_safe_torch_device, init_logging, set_seed +from lerobot.common.utils import format_big_number, get_safe_torch_device, init_logging, set_global_seed from lerobot.scripts.eval import eval_policy @@ -122,7 +122,7 @@ def train(cfg: dict, out_dir=None, job_name=None): torch.backends.cudnn.benchmark = True torch.backends.cuda.matmul.allow_tf32 = True - set_seed(cfg.seed) + set_global_seed(cfg.seed) logging.info("make_offline_buffer") offline_buffer = make_offline_buffer(cfg) @@ -224,7 +224,22 @@ def train(cfg: dict, out_dir=None, job_name=None): policy=td_policy, auto_cast_to_device=True, ) - assert len(rollout) <= cfg.env.episode_length + + assert ( + len(rollout.batch_size) == 2 + ), "2 dimensions expected: number of env in parallel x max number of steps during rollout" + + num_parallel_env = rollout.batch_size[0] + if num_parallel_env != 1: + # TODO(rcadene): when num_parallel_env > 1, rollout["episode"] needs to be properly set and we need to add tests + raise NotImplementedError() + + num_max_steps = rollout.batch_size[1] + assert num_max_steps <= cfg.env.episode_length + + # reshape to have a list of steps to insert into online_buffer + rollout = rollout.reshape(num_parallel_env * num_max_steps) + # set same episode index for all time steps contained in this rollout rollout["episode"] = torch.tensor([env_step] * len(rollout), dtype=torch.int) online_buffer.extend(rollout) diff --git a/poetry.lock b/poetry.lock index d2d39e7..e47b020 100644 --- a/poetry.lock +++ b/poetry.lock @@ -338,73 +338,6 @@ files = [ {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -[[package]] -name = "cython" -version = "3.0.9" -description = "The Cython compiler for writing C extensions in the Python language." -optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" -files = [ - {file = "Cython-3.0.9-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:296bd30d4445ac61b66c9d766567f6e81a6e262835d261e903c60c891a6729d3"}, - {file = "Cython-3.0.9-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f496b52845cb45568a69d6359a2c335135233003e708ea02155c10ce3548aa89"}, - {file = "Cython-3.0.9-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:858c3766b9aa3ab8a413392c72bbab1c144a9766b7c7bfdef64e2e414363fa0c"}, - {file = "Cython-3.0.9-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c0eb1e6ef036028a52525fd9a012a556f6dd4788a0e8755fe864ba0e70cde2ff"}, - {file = "Cython-3.0.9-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:c8191941073ea5896321de3c8c958fd66e5f304b0cd1f22c59edd0b86c4dd90d"}, - {file = "Cython-3.0.9-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:e32b016030bc72a8a22a1f21f470a2f57573761a4f00fbfe8347263f4fbdb9f1"}, - {file = "Cython-3.0.9-cp310-cp310-win32.whl", hash = "sha256:d6f3ff1cd6123973fe03e0fb8ee936622f976c0c41138969975824d08886572b"}, - {file = "Cython-3.0.9-cp310-cp310-win_amd64.whl", hash = "sha256:56f3b643dbe14449248bbeb9a63fe3878a24256664bc8c8ef6efd45d102596d8"}, - {file = "Cython-3.0.9-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:35e6665a20d6b8a152d72b7fd87dbb2af6bb6b18a235b71add68122d594dbd41"}, - {file = "Cython-3.0.9-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f92f4960c40ad027bd8c364c50db11104eadc59ffeb9e5b7f605ca2f05946e20"}, - {file = "Cython-3.0.9-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:38df37d0e732fbd9a2fef898788492e82b770c33d1e4ed12444bbc8a3b3f89c0"}, - {file = "Cython-3.0.9-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ad7fd88ebaeaf2e76fd729a8919fae80dab3d6ac0005e28494261d52ff347a8f"}, - {file = "Cython-3.0.9-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:1365d5f76bf4d19df3d19ce932584c9bb76e9fb096185168918ef9b36e06bfa4"}, - {file = "Cython-3.0.9-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:c232e7f279388ac9625c3e5a5a9f0078a9334959c5d6458052c65bbbba895e1e"}, - {file = "Cython-3.0.9-cp311-cp311-win32.whl", hash = "sha256:357e2fad46a25030b0c0496487e01a9dc0fdd0c09df0897f554d8ba3c1bc4872"}, - {file = "Cython-3.0.9-cp311-cp311-win_amd64.whl", hash = "sha256:1315aee506506e8d69cf6631d8769e6b10131fdcc0eb66df2698f2a3ddaeeff2"}, - {file = "Cython-3.0.9-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:157973807c2796addbed5fbc4d9c882ab34bbc60dc297ca729504901479d5df7"}, - {file = "Cython-3.0.9-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:00b105b5d050645dd59e6767bc0f18b48a4aa11c85f42ec7dd8181606f4059e3"}, - {file = "Cython-3.0.9-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ac5536d09bef240cae0416d5a703d298b74c7bbc397da803ac9d344e732d4369"}, - {file = "Cython-3.0.9-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:09c44501d476d16aaa4cbc29c87f8c0f54fc20e69b650d59cbfa4863426fc70c"}, - {file = "Cython-3.0.9-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:cc9c3b9f20d8e298618e5ccd32083ca386e785b08f9893fbec4c50b6b85be772"}, - {file = "Cython-3.0.9-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:a30d96938c633e3ec37000ac3796525da71254ef109e66bdfd78f29891af6454"}, - {file = "Cython-3.0.9-cp312-cp312-win32.whl", hash = "sha256:757ca93bdd80702546df4d610d2494ef2e74249cac4d5ba9464589fb464bd8a3"}, - {file = "Cython-3.0.9-cp312-cp312-win_amd64.whl", hash = "sha256:1dc320a9905ab95414013f6de805efbff9e17bb5fb3b90bbac533f017bec8136"}, - {file = "Cython-3.0.9-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:4ae349960ebe0da0d33724eaa7f1eb866688fe5434cc67ce4dbc06d6a719fbfc"}, - {file = "Cython-3.0.9-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:63d2537bf688247f76ded6dee28ebd26274f019309aef1eb4f2f9c5c482fde2d"}, - {file = "Cython-3.0.9-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:36f5a2dfc724bea1f710b649f02d802d80fc18320c8e6396684ba4a48412445a"}, - {file = "Cython-3.0.9-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:deaf4197d4b0bcd5714a497158ea96a2bd6d0f9636095437448f7e06453cc83d"}, - {file = "Cython-3.0.9-cp36-cp36m-musllinux_1_1_aarch64.whl", hash = "sha256:000af6deb7412eb7ac0c635ff5e637fb8725dd0a7b88cc58dfc2b3de14e701c4"}, - {file = "Cython-3.0.9-cp36-cp36m-musllinux_1_1_x86_64.whl", hash = "sha256:15c7f5c2d35bed9aa5f2a51eaac0df23ae72f2dbacf62fc672dd6bfaa75d2d6f"}, - {file = "Cython-3.0.9-cp36-cp36m-win32.whl", hash = "sha256:f49aa4970cd3bec66ac22e701def16dca2a49c59cceba519898dd7526e0be2c0"}, - {file = "Cython-3.0.9-cp36-cp36m-win_amd64.whl", hash = "sha256:4558814fa025b193058d42eeee498a53d6b04b2980d01339fc2444b23fd98e58"}, - {file = "Cython-3.0.9-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:539cd1d74fd61f6cfc310fa6bbbad5adc144627f2b7486a07075d4e002fd6aad"}, - {file = "Cython-3.0.9-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c3232926cd406ee02eabb732206f6e882c3aed9d58f0fea764013d9240405bcf"}, - {file = "Cython-3.0.9-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33b6ac376538a7fc8c567b85d3c71504308a9318702ec0485dd66c059f3165cb"}, - {file = "Cython-3.0.9-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2cc92504b5d22ac66031ffb827bd3a967fc75a5f0f76ab48bce62df19be6fdfd"}, - {file = "Cython-3.0.9-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:22b8fae756c5c0d8968691bed520876de452f216c28ec896a00739a12dba3bd9"}, - {file = "Cython-3.0.9-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:9cda0d92a09f3520f29bd91009f1194ba9600777c02c30c6d2d4ac65fb63e40d"}, - {file = "Cython-3.0.9-cp37-cp37m-win32.whl", hash = "sha256:ec612418490941ed16c50c8d3784c7bdc4c4b2a10c361259871790b02ec8c1db"}, - {file = "Cython-3.0.9-cp37-cp37m-win_amd64.whl", hash = "sha256:976c8d2bedc91ff6493fc973d38b2dc01020324039e2af0e049704a8e1b22936"}, - {file = "Cython-3.0.9-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5055988b007c92256b6e9896441c3055556038c3497fcbf8c921a6c1fce90719"}, - {file = "Cython-3.0.9-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d9360606d964c2d0492a866464efcf9d0a92715644eede3f6a2aa696de54a137"}, - {file = "Cython-3.0.9-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:02c6e809f060bed073dc7cba1648077fe3b68208863d517c8b39f3920eecf9dd"}, - {file = "Cython-3.0.9-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:95ed792c966f969cea7489c32ff90150b415c1f3567db8d5a9d489c7c1602dac"}, - {file = "Cython-3.0.9-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:8edd59d22950b400b03ca78d27dc694d2836a92ef0cac4f64cb4b2ff902f7e25"}, - {file = "Cython-3.0.9-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:4cf0ed273bf60e97922fcbbdd380c39693922a597760160b4b4355e6078ca188"}, - {file = "Cython-3.0.9-cp38-cp38-win32.whl", hash = "sha256:5eb9bd4ae12ebb2bc79a193d95aacf090fbd8d7013e11ed5412711650cb34934"}, - {file = "Cython-3.0.9-cp38-cp38-win_amd64.whl", hash = "sha256:44457279da56e0f829bb1fc5a5dc0836e5d498dbcf9b2324f32f7cc9d2ec6569"}, - {file = "Cython-3.0.9-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:c4b419a1adc2af43f4660e2f6eaf1e4fac2dbac59490771eb8ac3d6063f22356"}, - {file = "Cython-3.0.9-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4f836192140f033b2319a0128936367c295c2b32e23df05b03b672a6015757ea"}, - {file = "Cython-3.0.9-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2fd198c1a7f8e9382904d622cc0efa3c184605881fd5262c64cbb7168c4c1ec5"}, - {file = "Cython-3.0.9-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a274fe9ca5c53fafbcf5c8f262f8ad6896206a466f0eeb40aaf36a7951e957c0"}, - {file = "Cython-3.0.9-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:158c38360bbc5063341b1e78d3737f1251050f89f58a3df0d10fb171c44262be"}, - {file = "Cython-3.0.9-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:8bf30b045f7deda0014b042c1b41c1d272facc762ab657529e3b05505888e878"}, - {file = "Cython-3.0.9-cp39-cp39-win32.whl", hash = "sha256:9a001fd95c140c94d934078544ff60a3c46aca2dc86e75a76e4121d3cd1f4b33"}, - {file = "Cython-3.0.9-cp39-cp39-win_amd64.whl", hash = "sha256:530c01c4aebba709c0ec9c7ecefe07177d0b9fd7ffee29450a118d92192ccbdf"}, - {file = "Cython-3.0.9-py2.py3-none-any.whl", hash = "sha256:bf96417714353c5454c2e3238fca9338599330cf51625cdc1ca698684465646f"}, - {file = "Cython-3.0.9.tar.gz", hash = "sha256:a2d354f059d1f055d34cfaa62c5b68bc78ac2ceab6407148d47fb508cf3ba4f3"}, -] - [[package]] name = "debugpy" version = "1.8.1" @@ -639,6 +572,17 @@ files = [ [package.extras] test = ["pytest (>=6)"] +[[package]] +name = "farama-notifications" +version = "0.0.4" +description = "Notifications for all Farama Foundation maintained libraries." +optional = false +python-versions = "*" +files = [ + {file = "Farama-Notifications-0.0.4.tar.gz", hash = "sha256:13fceff2d14314cf80703c8266462ebf3733c7d165336eee998fc58e545efd18"}, + {file = "Farama_Notifications-0.0.4-py3-none-any.whl", hash = "sha256:14de931035a41961f7c056361dc7f980762a143d05791ef5794a751a2caf05ae"}, +] + [[package]] name = "fasteners" version = "0.19" @@ -840,43 +784,58 @@ files = [ protobuf = ["grpcio-tools (>=1.62.1)"] [[package]] -name = "gym" -version = "0.26.2" -description = "Gym: A universal API for reinforcement learning environments" +name = "gymnasium" +version = "0.29.1" +description = "A standard API for reinforcement learning and a diverse set of reference environments (formerly Gym)." optional = false -python-versions = ">=3.6" +python-versions = ">=3.8" files = [ - {file = "gym-0.26.2.tar.gz", hash = "sha256:e0d882f4b54f0c65f203104c24ab8a38b039f1289986803c7d02cdbe214fbcc4"}, + {file = "gymnasium-0.29.1-py3-none-any.whl", hash = "sha256:61c3384b5575985bb7f85e43213bcb40f36fcdff388cae6bc229304c71f2843e"}, + {file = "gymnasium-0.29.1.tar.gz", hash = "sha256:1a532752efcb7590478b1cc7aa04f608eb7a2fdad5570cd217b66b6a35274bb1"}, ] [package.dependencies] cloudpickle = ">=1.2.0" -gym_notices = ">=0.0.4" -numpy = ">=1.18.0" +farama-notifications = ">=0.0.1" +numpy = ">=1.21.0" +typing-extensions = ">=4.3.0" [package.extras] accept-rom-license = ["autorom[accept-rom-license] (>=0.4.2,<0.5.0)"] -all = ["ale-py (>=0.8.0,<0.9.0)", "box2d-py (==2.3.5)", "imageio (>=2.14.1)", "lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "mujoco (==2.2)", "mujoco_py (>=2.1,<2.2)", "opencv-python (>=3.0)", "pygame (==2.1.0)", "pytest (==7.0.1)", "swig (==4.*)"] -atari = ["ale-py (>=0.8.0,<0.9.0)"] -box2d = ["box2d-py (==2.3.5)", "pygame (==2.1.0)", "swig (==4.*)"] -classic-control = ["pygame (==2.1.0)"] -mujoco = ["imageio (>=2.14.1)", "mujoco (==2.2)"] -mujoco-py = ["mujoco_py (>=2.1,<2.2)"] -other = ["lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "opencv-python (>=3.0)"] -testing = ["box2d-py (==2.3.5)", "imageio (>=2.14.1)", "lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "mujoco (==2.2)", "mujoco_py (>=2.1,<2.2)", "opencv-python (>=3.0)", "pygame (==2.1.0)", "pytest (==7.0.1)", "swig (==4.*)"] -toy-text = ["pygame (==2.1.0)"] +all = ["box2d-py (==2.3.5)", "cython (<3)", "imageio (>=2.14.1)", "jax (>=0.4.0)", "jaxlib (>=0.4.0)", "lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "mujoco (>=2.3.3)", "mujoco-py (>=2.1,<2.2)", "opencv-python (>=3.0)", "pygame (>=2.1.3)", "shimmy[atari] (>=0.1.0,<1.0)", "swig (==4.*)", "torch (>=1.0.0)"] +atari = ["shimmy[atari] (>=0.1.0,<1.0)"] +box2d = ["box2d-py (==2.3.5)", "pygame (>=2.1.3)", "swig (==4.*)"] +classic-control = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"] +jax = ["jax (>=0.4.0)", "jaxlib (>=0.4.0)"] +mujoco = ["imageio (>=2.14.1)", "mujoco (>=2.3.3)"] +mujoco-py = ["cython (<3)", "cython (<3)", "mujoco-py (>=2.1,<2.2)", "mujoco-py (>=2.1,<2.2)"] +other = ["lz4 (>=3.1.0)", "matplotlib (>=3.0)", "moviepy (>=1.0.0)", "opencv-python (>=3.0)", "torch (>=1.0.0)"] +testing = ["pytest (==7.1.3)", "scipy (>=1.7.3)"] +toy-text = ["pygame (>=2.1.3)", "pygame (>=2.1.3)"] [[package]] -name = "gym-notices" -version = "0.0.8" -description = "Notices for gym" +name = "gymnasium-robotics" +version = "1.2.4" +description = "Robotics environments for the Gymnasium repo." optional = false -python-versions = "*" +python-versions = ">=3.8" files = [ - {file = "gym-notices-0.0.8.tar.gz", hash = "sha256:ad25e200487cafa369728625fe064e88ada1346618526102659b4640f2b4b911"}, - {file = "gym_notices-0.0.8-py3-none-any.whl", hash = "sha256:e5f82e00823a166747b4c2a07de63b6560b1acb880638547e0cabf825a01e463"}, + {file = "gymnasium-robotics-1.2.4.tar.gz", hash = "sha256:d304192b066f8b800599dfbe3d9d90bba9b761ee884472bdc4d05968a8bc61cb"}, + {file = "gymnasium_robotics-1.2.4-py3-none-any.whl", hash = "sha256:c2cb23e087ca0280ae6802837eb7b3a6d14e5bd24c00803ab09f015fcff3eef5"}, ] +[package.dependencies] +gymnasium = ">=0.26" +imageio = "*" +Jinja2 = ">=3.0.3" +mujoco = ">=2.3.3,<3.0" +numpy = ">=1.21.0" +PettingZoo = ">=1.23.0" + +[package.extras] +mujoco-py = ["cython (<3)", "mujoco-py (>=2.1,<2.2)"] +testing = ["Jinja2 (>=3.0.3)", "PettingZoo (>=1.23.0)", "cython (<3)", "mujoco-py (>=2.1,<2.2)", "pytest (==7.0.1)"] + [[package]] name = "h5py" version = "3.10.0" @@ -1506,25 +1465,6 @@ glfw = "*" numpy = "*" pyopengl = "*" -[[package]] -name = "mujoco-py" -version = "2.1.2.14" -description = "" -optional = false -python-versions = ">=3.6" -files = [ - {file = "mujoco-py-2.1.2.14.tar.gz", hash = "sha256:eb5b14485acf80a3cf8c15f4b080c6a28a9f79e68869aa696d16cbd51ea7706f"}, - {file = "mujoco_py-2.1.2.14-py3-none-any.whl", hash = "sha256:37c0b41bc0153a8a0eb3663103a67c60f65467753f74e4ff6e68b879f3e3a71f"}, -] - -[package.dependencies] -cffi = ">=1.10" -Cython = ">=0.27.2" -fasteners = ">=0.15,<1.0" -glfw = ">=1.4.0" -imageio = ">=2.1.2" -numpy = ">=1.11" - [[package]] name = "networkx" version = "3.2.1" @@ -1940,6 +1880,31 @@ sql-other = ["SQLAlchemy (>=2.0.0)", "adbc-driver-postgresql (>=0.8.0)", "adbc-d test = ["hypothesis (>=6.46.1)", "pytest (>=7.3.2)", "pytest-xdist (>=2.2.0)"] xml = ["lxml (>=4.9.2)"] +[[package]] +name = "pettingzoo" +version = "1.24.3" +description = "Gymnasium for multi-agent reinforcement learning." +optional = false +python-versions = ">=3.8" +files = [ + {file = "pettingzoo-1.24.3-py3-none-any.whl", hash = "sha256:23ed90517d2e8a7098bdaf5e31234b3a7f7b73ca578d70d1ca7b9d0cb0e37982"}, + {file = "pettingzoo-1.24.3.tar.gz", hash = "sha256:91f9094f18e06fb74b98f4099cd22e8ae4396125e51719d50b30c9f1c7ab07e6"}, +] + +[package.dependencies] +gymnasium = ">=0.28.0" +numpy = ">=1.21.0" + +[package.extras] +all = ["box2d-py (==2.3.5)", "chess (==1.9.4)", "multi-agent-ale-py (==0.1.11)", "pillow (>=8.0.1)", "pygame (==2.3.0)", "pymunk (==6.2.0)", "rlcard (==1.0.5)", "scipy (>=1.4.1)", "shimmy[openspiel] (>=1.2.0)"] +atari = ["multi-agent-ale-py (==0.1.11)", "pygame (==2.3.0)"] +butterfly = ["pygame (==2.3.0)", "pymunk (==6.2.0)"] +classic = ["chess (==1.9.4)", "pygame (==2.3.0)", "rlcard (==1.0.5)", "shimmy[openspiel] (>=1.2.0)"] +mpe = ["pygame (==2.3.0)"] +other = ["pillow (>=8.0.1)"] +sisl = ["box2d-py (==2.3.5)", "pygame (==2.3.0)", "pymunk (==6.2.0)", "scipy (>=1.4.1)"] +testing = ["AutoROM", "pre-commit", "pynput", "pytest", "pytest-cov", "pytest-markdown-docs", "pytest-xdist"] + [[package]] name = "pillow" version = "10.2.0" @@ -3510,4 +3475,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "1a45c808e1c48bcbf4319d4cf6876771b7d50f40a5a8968a8b7f3af36192bf34" +content-hash = "99addbfc02bcd35a308f4ecc5b4285c9c5054118f4aadea27650d8bf355d9616" diff --git a/pyproject.toml b/pyproject.toml index 7e9996a..29cec3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ packages = [{include = "lerobot"}] [tool.poetry.dependencies] python = "^3.10" -cython = "^3.0.8" termcolor = "^2.4.0" omegaconf = "^2.3.0" dm-env = "^1.6" @@ -42,9 +41,7 @@ mpmath = "^1.3.0" torch = "^2.2.1" tensordict = {git = "https://github.com/pytorch/tensordict"} torchrl = {git = "https://github.com/pytorch/rl", rev = "13bef426dcfa5887c6e5034a6e9697993fa92c37"} -mujoco = "2.3.7" -mujoco-py = "^2.1.2.14" -gym = "^0.26.2" +mujoco = "^2.3.7" opencv-python = "^4.9.0.80" diffusers = "^0.26.3" torchvision = "^0.17.1" @@ -52,6 +49,8 @@ h5py = "^3.10.0" dm-control = "1.0.14" huggingface-hub = {extras = ["hf-transfer"], version = "^0.21.4"} robomimic = "0.2.0" +gymnasium-robotics = "^1.2.4" +gymnasium = "^0.29.1" [tool.poetry.group.dev.dependencies] @@ -90,7 +89,7 @@ exclude = [ [tool.ruff.lint] select = ["E4", "E7", "E9", "F", "I", "N", "B", "C4", "SIM"] - +ignore-init-module-imports = true [tool.poetry-dynamic-versioning] enable = true diff --git a/tests/data/xarm_lift_medium/replay_buffer/action.memmap b/tests/data/xarm_lift_medium/replay_buffer/action.memmap new file mode 100644 index 0000000..c90afbe --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/action.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:10ec2f944de18f1a2aa3fc2f9a4185c03e3a5afc31148c85c98b58602ac4186e +size 800 diff --git a/tests/data/xarm_lift_medium/replay_buffer/episode.memmap b/tests/data/xarm_lift_medium/replay_buffer/episode.memmap new file mode 100644 index 0000000..7924f02 --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/episode.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1a589cba6bf0dfce138110864b6509508a804d7ea5c519d0a3cd67c4a87fa2d0 +size 200 diff --git a/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap b/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap new file mode 100644 index 0000000..a633d34 --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/frame_id.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6afe7098f30bdc8564526517c085e62613f6cb67194153840567974cfa6f3815 +size 400 diff --git a/tests/data/xarm_lift_medium/replay_buffer/meta.json b/tests/data/xarm_lift_medium/replay_buffer/meta.json new file mode 100644 index 0000000..33dc932 --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/meta.json @@ -0,0 +1 @@ +{"action": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "episode": {"device": "cpu", "shape": [50], "dtype": "torch.int32"}, "frame_id": {"device": "cpu", "shape": [50], "dtype": "torch.int64"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap new file mode 100644 index 0000000..cf5e9cc --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/next/done.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:dab3a9712c413c4bfcd91c645752ab981306b23d25bcd4da4c422911574673f7 +size 50 diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/meta.json b/tests/data/xarm_lift_medium/replay_buffer/next/meta.json new file mode 100644 index 0000000..d69cada --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/next/meta.json @@ -0,0 +1 @@ +{"reward": {"device": "cpu", "shape": [50], "dtype": "torch.float32"}, "done": {"device": "cpu", "shape": [50], "dtype": "torch.bool"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap new file mode 100644 index 0000000..462d011 --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/next/observation/image.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d6f9d1422ce3764e7253f70ed4da278f0c07fafef0d5386479f09d6b6b9d8259 +size 1058400 diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json b/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json new file mode 100644 index 0000000..b13b8ec --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/next/observation/meta.json @@ -0,0 +1 @@ +{"image": {"device": "cpu", "shape": [50, 3, 84, 84], "dtype": "torch.uint8"}, "state": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap new file mode 100644 index 0000000..1dbe602 --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/next/observation/state.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:52e7c1a3c4fb2423b195e66ffee2c9e23b3ea0ad7c8bfc4dec30a35c65cadcbb +size 800 diff --git a/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap b/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap new file mode 100644 index 0000000..9ff5d5a --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/next/reward.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:c4dbe8ea1966e5cc6da6daf5704805b9b5810f4575de7016b8f6cb1da1d7bb8a +size 200 diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap b/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap new file mode 100644 index 0000000..c941694 --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/observation/image.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8fca8ddbda3f7bb2f6e7553297c18f3ab8f8b73d64b5c9f81a3695ad9379d403 +size 1058400 diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json b/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json new file mode 100644 index 0000000..b13b8ec --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/observation/meta.json @@ -0,0 +1 @@ +{"image": {"device": "cpu", "shape": [50, 3, 84, 84], "dtype": "torch.uint8"}, "state": {"device": "cpu", "shape": [50, 4], "dtype": "torch.float32"}, "shape": [50], "device": "cpu", "_type": ""} \ No newline at end of file diff --git a/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap b/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap new file mode 100644 index 0000000..3bae16d --- /dev/null +++ b/tests/data/xarm_lift_medium/replay_buffer/observation/state.memmap @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:7b3e3e12896d553c208ee152f6d447c877c435e15d010c4a6171966d5b8a0c0b +size 800 diff --git a/tests/data/xarm_lift_medium/stats.pth b/tests/data/xarm_lift_medium/stats.pth new file mode 100644 index 0000000..0accffb Binary files /dev/null and b/tests/data/xarm_lift_medium/stats.pth differ diff --git a/tests/test_datasets.py b/tests/test_datasets.py index b7d1e6f..252e004 100644 --- a/tests/test_datasets.py +++ b/tests/test_datasets.py @@ -9,10 +9,8 @@ from .utils import DEVICE, init_config @pytest.mark.parametrize( "env_name,dataset_id", [ - # TODO(rcadene): simxarm is depreciated for now - # ("simxarm", "lift"), + ("simxarm", "lift"), ("pusht", "pusht"), - # TODO(aliberts): add aloha when dataset is available on hub ("aloha", "sim_insertion_human"), ("aloha", "sim_insertion_scripted"), ("aloha", "sim_transfer_cube_human"), diff --git a/tests/test_envs.py b/tests/test_envs.py index 7776ba3..1db83af 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -7,7 +7,7 @@ from lerobot.common.datasets.factory import make_offline_buffer from lerobot.common.envs.factory import make_env from lerobot.common.envs.pusht.env import PushtEnv -from lerobot.common.envs.simxarm import SimxarmEnv +from lerobot.common.envs.simxarm.env import SimxarmEnv from .utils import DEVICE, init_config @@ -39,19 +39,19 @@ def print_spec_rollout(env): print("data from rollout:", simple_rollout(100)) -@pytest.mark.skip(reason="Simxarm is deprecated") @pytest.mark.parametrize( "task,from_pixels,pixels_only", [ ("lift", False, False), ("lift", True, False), ("lift", True, True), - ("reach", False, False), - ("reach", True, False), - ("push", False, False), - ("push", True, False), - ("peg_in_box", False, False), - ("peg_in_box", True, False), + # TODO(aliberts): Add simxarm other tasks + # ("reach", False, False), + # ("reach", True, False), + # ("push", False, False), + # ("push", True, False), + # ("peg_in_box", False, False), + # ("peg_in_box", True, False), ], ) def test_simxarm(task, from_pixels, pixels_only): @@ -84,7 +84,7 @@ def test_pusht(from_pixels, pixels_only): @pytest.mark.parametrize( "env_name", [ - # "simxarm", + "simxarm", "pusht", "aloha", ], diff --git a/tests/test_policies.py b/tests/test_policies.py index 92508da..cd08fc4 100644 --- a/tests/test_policies.py +++ b/tests/test_policies.py @@ -19,12 +19,13 @@ from .utils import DEVICE, init_config [ ("simxarm", "tdmpc", ["policy.mpc=true"]), ("pusht", "tdmpc", ["policy.mpc=false"]), - ("simxarm", "diffusion", []), ("pusht", "diffusion", []), ("aloha", "act", ["env.task=sim_insertion_scripted"]), ("aloha", "act", ["env.task=sim_insertion_human"]), ("aloha", "act", ["env.task=sim_transfer_cube_scripted"]), ("aloha", "act", ["env.task=sim_transfer_cube_human"]), + # TODO(aliberts): simxarm not working with diffusion + # ("simxarm", "diffusion", []), ], ) def test_concrete_policy(env_name, policy_name, extra_overrides): @@ -45,13 +46,6 @@ def test_concrete_policy(env_name, policy_name, extra_overrides): # Check that we can make the policy object. policy = make_policy(cfg) # Check that we run select_actions and get the appropriate output. - if env_name == "simxarm": - # TODO(rcadene): Not implemented - return - if policy_name == "tdmpc": - # TODO(alexander-soare): TDMPC does not use n_obs_steps but the environment requires this. - with open_dict(cfg): - cfg["n_obs_steps"] = 1 offline_buffer = make_offline_buffer(cfg) env = make_env(cfg, transform=offline_buffer.transform)