From 5c13547de02cfcfd3ce83e22abc35a3b53e49cf7 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Wed, 24 Jun 2026 17:42:56 +0200 Subject: [PATCH 01/12] start new branch for v2.1.1 --- docs/changelog.md | 4 ++++ pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/changelog.md b/docs/changelog.md index 0333d494..7ed79576 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -1,5 +1,9 @@ # Changelog +## Version 2.1.1 (under development) + +TODO + ## Version 2.1.0 !!! danger "API-breaking changes" diff --git a/pyproject.toml b/pyproject.toml index 67ebac77..2b95fe0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "flygym" -version = "2.1.0" +version = "2.1.1" description = "NeuroMechFly, a framework for simulating embodied sensorimotor control in adult Drosophila." requires-python = ">=3.12,<3.15" authors = [ diff --git a/uv.lock b/uv.lock index f879f2d4..34b05cab 100644 --- a/uv.lock +++ b/uv.lock @@ -683,7 +683,7 @@ wheels = [ [[package]] name = "flygym" -version = "2.1.0" +version = "2.1.1" source = { editable = "." } dependencies = [ { name = "imageio" }, From aa724e22127301bf3cd30a2d70f110ddad6c1349 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Mon, 29 Jun 2026 15:24:20 +0200 Subject: [PATCH 02/12] bump up mujoco version --- pyproject.toml | 4 ++-- uv.lock | 44 ++++++++++++++++++++++---------------------- 2 files changed, 24 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2b95fe0f..2c24c3f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,7 @@ maintainers = [ readme = "README.md" license = "Apache-2.0" dependencies = [ - "mujoco>=3.9,<3.10", + "mujoco>=3.10,<3.11", "numpy>=2.0,<3.0", "pyyaml>=6.0,<7.0", "mediapy>=1.2,<3.0", @@ -39,7 +39,7 @@ dependencies = [ [project.optional-dependencies] warp = [ "warp-lang>=1.14,<1.15", - "mujoco_warp>=3.9,<3.10", + "mujoco_warp>=3.10,<3.11", ] dev = [ # Testing diff --git a/uv.lock b/uv.lock index 34b05cab..7a2e42dc 100644 --- a/uv.lock +++ b/uv.lock @@ -763,8 +763,8 @@ requires-dist = [ { name = "mkdocs-material", marker = "extra == 'dev'", specifier = ">=9.7,<10.0" }, { name = "mkdocs-section-index", marker = "extra == 'dev'", specifier = "==0.3.12" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'dev'", specifier = ">=1.0.0,<2" }, - { name = "mujoco", specifier = ">=3.9,<3.10" }, - { name = "mujoco-warp", marker = "extra == 'warp'", specifier = ">=3.9,<3.10" }, + { name = "mujoco", specifier = ">=3.10,<3.11" }, + { name = "mujoco-warp", marker = "extra == 'warp'", specifier = ">=3.10,<3.11" }, { name = "nbstripout", marker = "extra == 'dev'", specifier = "==0.9.1" }, { name = "networkx", marker = "extra == 'dev'", specifier = ">=3.0,<4.0" }, { name = "numba", specifier = ">=0.65.1" }, @@ -1811,7 +1811,7 @@ wheels = [ [[package]] name = "mujoco" -version = "3.9.0" +version = "3.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -1820,28 +1820,28 @@ dependencies = [ { name = "numpy" }, { name = "pyopengl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bc/f6/325fa0c1a1d63b1bb6852e265fb8d7fd8a88ddcd3580dc153e1dceb197dd/mujoco-3.9.0.tar.gz", hash = "sha256:0a1ca878c4a80893251420037a0583252d1356a0cc8d661268627c674c7835f7", size = 923500, upload-time = "2026-05-27T14:49:48.52Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a2/3b/c76837b7fdb007f7605ff783689a0bd23a5a49b065928bec2f1fa7ea3d67/mujoco-3.10.0.tar.gz", hash = "sha256:c9e8d5d87d82204ed5bccc87d843c0a53e75aaf381de2938ec46d04f1ac6e24e", size = 1094987, upload-time = "2026-06-22T17:40:59.904Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5a/9b/86bc30e294f835a04ed261a3271d6f192595fd70e5f9f0bbe7f31a357ddd/mujoco-3.9.0-cp312-cp312-macosx_10_16_x86_64.whl", hash = "sha256:f1d8736589ad736f8ec1601a701a8ab84aff5a2b29244fe58194d8f0a0984865", size = 7467226, upload-time = "2026-05-27T14:49:18.948Z" }, - { url = "https://files.pythonhosted.org/packages/e6/14/5c8d8d660e23bcd6963f3f7df525504cd082000efc2bbc76812b9b7a4af5/mujoco-3.9.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:2f38a5d011425585d10a1b4dd9ff941072ffccda2db071b462c86967199a3bb0", size = 7461002, upload-time = "2026-05-27T14:49:21.028Z" }, - { url = "https://files.pythonhosted.org/packages/f9/7e/eb1840f486d322cd5285a2dc2e422cbc22f4e273b66edc0cf9d96cbad47c/mujoco-3.9.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8af0cb8c96f5b02e701f9ab7a9c905d1c655626b3226a4c4da7b28bf2771a4e7", size = 6847433, upload-time = "2026-05-27T14:49:22.974Z" }, - { url = "https://files.pythonhosted.org/packages/d5/01/b2a88b6b73df933d5ab38583240c296684b626a8de3c3bb9a7c2fd356f08/mujoco-3.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f8c6769ac1463752ff392d8c4d20014e54220112d772e6ccecdf6bcdb6d43e07", size = 7392160, upload-time = "2026-05-27T14:49:25.543Z" }, - { url = "https://files.pythonhosted.org/packages/66/22/7ce98722ecfb38aa2e3ca9da7bbe2b662ef95ec13acad1c8a63f22e14f39/mujoco-3.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:812406661161ebff1335e5293a6fb2ae4b193a9ea03e3f6579875aef345a3c8f", size = 5974891, upload-time = "2026-05-27T14:49:27.147Z" }, - { url = "https://files.pythonhosted.org/packages/f8/63/9d2cd488ccb8500aeb76dda6d9f8c3ad1e55b0d72066caa728e577bdf178/mujoco-3.9.0-cp313-cp313-macosx_10_16_x86_64.whl", hash = "sha256:a0b2c213362091e9b7868ed0d329cd092c06a020f2b76da21981653e1a50504c", size = 7467496, upload-time = "2026-05-27T14:49:29.278Z" }, - { url = "https://files.pythonhosted.org/packages/b1/88/a29bca408a4c2db6c5bcf58a8b92c464660b7f846c559abd9110783574cb/mujoco-3.9.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2e55da7485025b8266ee8ab715eb0eae9c5d2b14619a8c8cebb6cce88f9545ff", size = 7461035, upload-time = "2026-05-27T14:49:31.281Z" }, - { url = "https://files.pythonhosted.org/packages/14/22/342c0298d0d1f36393528ac1b34c308b79ffe1893fd21ddcc0dafc2ea3d5/mujoco-3.9.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fc9a82c103b15086e43aa4278bcd52cb11a3e5e3da2cba2239cd80a08a4d51fa", size = 6847101, upload-time = "2026-05-27T14:49:33.229Z" }, - { url = "https://files.pythonhosted.org/packages/97/df/a1495de78c1da3e8e93978dd177b04d18aaa7361452e30a3467c41c3b19e/mujoco-3.9.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ddf821ec7f7d991534da4a3440a57adc59bacf6c0a0cf3b5b05691613ccac91f", size = 7392158, upload-time = "2026-05-27T14:49:35.067Z" }, - { url = "https://files.pythonhosted.org/packages/8f/34/37133c8ff51f9c92cb47be537c2e8a1ea7ab201d5e66241a686efe2438cb/mujoco-3.9.0-cp313-cp313-win_amd64.whl", hash = "sha256:4e9d5602c0c337dbd1856b600df4388811fb753c8e84b835c17f3f37afab42d1", size = 5974865, upload-time = "2026-05-27T14:49:37.165Z" }, - { url = "https://files.pythonhosted.org/packages/98/6f/0ce33544e068183e053d88d18d14c77f0b98afa0347e830ad546e647b2df/mujoco-3.9.0-cp314-cp314-macosx_10_16_x86_64.whl", hash = "sha256:4395896d56618d5020cf694e749e13946ef27023d611aa0066cd2477956d3f63", size = 7523973, upload-time = "2026-05-27T14:49:39.427Z" }, - { url = "https://files.pythonhosted.org/packages/8e/96/c17b8b1d81dd8be2e74e1831e409bea0e6503855be6e6d4367c57d671465/mujoco-3.9.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:2f65ccdc6fe3027d27cc00761b761c8914bb3701d1a5a9e202aedda8151819f4", size = 7503471, upload-time = "2026-05-27T14:49:41.436Z" }, - { url = "https://files.pythonhosted.org/packages/41/b0/ce0e6431338d68cb54a7d5f1ac050521203b261d047315afb02ccf395469/mujoco-3.9.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:46182fce59757da66970b4573363f7442e1c8afeb43c4d61f28ef9718982e0e7", size = 6858506, upload-time = "2026-05-27T14:49:43.067Z" }, - { url = "https://files.pythonhosted.org/packages/5b/10/d3574267555b9923e88443b4524bb6fd7859b82adeaade10c2d32dafe5f9/mujoco-3.9.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:aa860cf246b0bf340e5469f54d48b4adb40c191e1c239341783b8935bebd1853", size = 7392326, upload-time = "2026-05-27T14:49:44.96Z" }, - { url = "https://files.pythonhosted.org/packages/5d/fb/690f5b7a2a8f794dea942120f4a794ec3167d1ab4ffd2791d45d6fd284cc/mujoco-3.9.0-cp314-cp314-win_amd64.whl", hash = "sha256:ff40abf8d5154fd847d610176dd3a7edd6f3e7392f957151fb0f0f4e8c20e771", size = 6313119, upload-time = "2026-05-27T14:49:46.831Z" }, + { url = "https://files.pythonhosted.org/packages/03/a2/4dd9f4cec6ce92f836a8b2de1cc799c4458af1467d7a044ef8014217bdb4/mujoco-3.10.0-cp312-cp312-macosx_10_16_x86_64.whl", hash = "sha256:47d4a22b7667c60e24e7ef6acb027c13abe9abba9acf17cc8db6fb250ba275ea", size = 7772567, upload-time = "2026-06-22T17:40:16.539Z" }, + { url = "https://files.pythonhosted.org/packages/5d/7d/ebe5342c136de27e0c430ba781f829df2cd66c00ed22627c1964fbd5d7fe/mujoco-3.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:a4d35e9d0b13ff9ad3196294a7dac363f1d0cdaa988832d0b687d42d98f4ee29", size = 19380823, upload-time = "2026-06-22T17:40:19.211Z" }, + { url = "https://files.pythonhosted.org/packages/d8/5d/43d1b2b9fe97676e5af03020e132ac497b45a0333a4c61de657d0d52170a/mujoco-3.10.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cb7f0d7c148a588f3633020807fc0ec3f3a9aff1f647406e3e0ffe96b05dfd57", size = 19705628, upload-time = "2026-06-22T17:40:22.549Z" }, + { url = "https://files.pythonhosted.org/packages/c3/11/c69199e4123935f98068ab6ab6b35955b4de0f6a91d3f9883805a5789394/mujoco-3.10.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:966d12f88e77e2b188e7530667b519d6963b9b906cff83bab534e5e4279325a0", size = 20904309, upload-time = "2026-06-22T17:40:25.861Z" }, + { url = "https://files.pythonhosted.org/packages/47/13/07bf2550c7dcd69ee8c7fd1f5c400a4ba2e4ede0a29a463ad3ac4cc9da90/mujoco-3.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:708edb5aceee96f2767b1072641523060043b2c67000e39e6e9797addf073696", size = 17865123, upload-time = "2026-06-22T17:40:28.996Z" }, + { url = "https://files.pythonhosted.org/packages/c0/ea/a9a19f025d3b323aca9d00e34fdc177e7290c5f23adf6341109a7e961920/mujoco-3.10.0-cp313-cp313-macosx_10_16_x86_64.whl", hash = "sha256:396a49c1baeb4cb1efd1fc213e54baf6d2a5ca42bcd99da82278501a32cb5d42", size = 7772943, upload-time = "2026-06-22T17:40:31.951Z" }, + { url = "https://files.pythonhosted.org/packages/68/ea/15dfc7cc10f44f46286a0634499f3a63e03269adb4788903f3751e964be0/mujoco-3.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:db794d8b8a64e3bb6a0f226a9200293a4890f490c93411fcc531a145f373521b", size = 19381290, upload-time = "2026-06-22T17:40:34.251Z" }, + { url = "https://files.pythonhosted.org/packages/76/82/ab50abb150805e711e02b7b7424b86bd12297bc5476fb1455ce739fa0d1b/mujoco-3.10.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:720174e7132ccf5f7e9eb86a0855c1b1e93330755f734c21fa58e7bde65bd140", size = 19705727, upload-time = "2026-06-22T17:40:36.968Z" }, + { url = "https://files.pythonhosted.org/packages/52/49/f6f0c7c54c646adda647aa7495bdc69572419480c37d5cd0bee132ebacd6/mujoco-3.10.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e0a5725cc9c9a4da1d5f2b8314ad6d39007e09430ac00aa8d023995ce1f2a3e9", size = 20904596, upload-time = "2026-06-22T17:40:40.423Z" }, + { url = "https://files.pythonhosted.org/packages/05/ea/53c7bcef3ef4de63563d9536e4fe6728996f0fc13e82040b1614ad0e32d9/mujoco-3.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:928dca15b06161b88153df2dc2f7c141d41f39e08e20d1804df8d53f42d9f62c", size = 17865523, upload-time = "2026-06-22T17:40:43.233Z" }, + { url = "https://files.pythonhosted.org/packages/24/69/a55911e462181ec5793cdcd6e4a8bf35dedbe4c8cb8f71f51ad38f724f55/mujoco-3.10.0-cp314-cp314-macosx_10_16_x86_64.whl", hash = "sha256:d8c9a865b04030933de15842823856ce0ef20a12bf5c972a3782e27e78158e41", size = 7835852, upload-time = "2026-06-22T17:40:46.062Z" }, + { url = "https://files.pythonhosted.org/packages/3c/6d/e96037c529e3639f86ac7a62252f464bea2683ef534833bc76658a44f8f6/mujoco-3.10.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1cdc8f3487e5d73fc330c446fd6e521e81fb28b39b4f735e1cc934c09044ec5b", size = 19400276, upload-time = "2026-06-22T17:40:48.199Z" }, + { url = "https://files.pythonhosted.org/packages/4f/c5/f5937fc94d6922d85c750434642a0c02be361de06e5ae35b3b5b4426907a/mujoco-3.10.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0f4ebc3f494cb2dabcdab499154ba7e02b2d2e98f83612961a4fb6365abb79ba", size = 19719062, upload-time = "2026-06-22T17:40:51.418Z" }, + { url = "https://files.pythonhosted.org/packages/ff/37/1df28d466b09f57f414a5c165075789e123b25b3127c9faf2ded5deca616/mujoco-3.10.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:04b57f3f9a58b35b99d5966f8be12f68158ad0c7393b219378af0686f3856dbe", size = 20905473, upload-time = "2026-06-22T17:40:54.665Z" }, + { url = "https://files.pythonhosted.org/packages/51/d2/14e60c0d3ea248782e8b5ce4844f8970a4d812078b2536aff61385876e94/mujoco-3.10.0-cp314-cp314-win_amd64.whl", hash = "sha256:cb48d042682d227188ba22b51c31b72975bf7ea7d89f0cafed3d40634337f88e", size = 18572894, upload-time = "2026-06-22T17:40:57.73Z" }, ] [[package]] name = "mujoco-warp" -version = "3.9.0.1" +version = "3.10.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "absl-py" }, @@ -1850,9 +1850,9 @@ dependencies = [ { name = "numpy" }, { name = "warp-lang" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/da/df/713e655aae5498e60fc92f6cad868f9cb5f24f7e7e67ebd88fbcdaded72c/mujoco_warp-3.9.0.1.tar.gz", hash = "sha256:f6efacb2ac23041d34cee5208c6c95b4b34a071c8b5df5e223aa7a48a7aac431", size = 1981867, upload-time = "2026-05-27T16:10:19.096Z" } +sdist = { url = "https://files.pythonhosted.org/packages/fc/e5/37d982ba8bdbad8455fb74da827aa4e8ea1b984db66171d4b66c458da985/mujoco_warp-3.10.0.1.tar.gz", hash = "sha256:adca3e740112bca1aaaa732556aad0f36fb3dea99f86abd22ee1232dc5e2b28d", size = 2046299, upload-time = "2026-06-26T15:59:41.217Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/38/9a/15fd259527550066b9f9a296fc24ac8c53019dcdcf817a3f89671da0c381/mujoco_warp-3.9.0.1-py3-none-any.whl", hash = "sha256:0ca40c4b4a8de11a7c2d293b049db73e0b2060a8787a37a81351324e60a55856", size = 2062165, upload-time = "2026-05-27T16:10:17.124Z" }, + { url = "https://files.pythonhosted.org/packages/16/32/9adafac9cf8133cebe5825b306be01ab9209ab0fb5a057cb29e4f5b23bd5/mujoco_warp-3.10.0.1-py3-none-any.whl", hash = "sha256:7a5871a5522796fa36ec5614dc214a180d4fc7de861e413ee1559d2cc4ec79f9", size = 2123960, upload-time = "2026-06-26T15:59:39.77Z" }, ] [[package]] From 8f2d7c253410d92a74d5dc58d5dcfc8a17d974e2 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:50:33 +0200 Subject: [PATCH 03/12] add fly alias in world and simulation --- src/flygym/compose/world/base_world.py | 14 ++++++++++++++ src/flygym/simulation.py | 12 +++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/flygym/compose/world/base_world.py b/src/flygym/compose/world/base_world.py index 25b2bad7..75cefa32 100644 --- a/src/flygym/compose/world/base_world.py +++ b/src/flygym/compose/world/base_world.py @@ -81,6 +81,20 @@ def mjcf_root(self) -> mj.MjSpec: def fly_lookup(self) -> dict[str, BaseFly]: """Lookup for `Fly` objects in the world, keyed by fly name.""" return self._fly_lookup + + @property + def fly(self) -> BaseFly: + """Get the single fly in the world. + + Raises: + ValueError: If there is not exactly one fly in the world. + """ + if len(self.fly_lookup) != 1: + raise ValueError( + "World contains multiple flies. " + "`.fly` is ambiguous; use `.fly_lookup` instead." + ) + return next(iter(self.fly_lookup.values())) @abstractmethod def _attach_fly_mjcf( diff --git a/src/flygym/simulation.py b/src/flygym/simulation.py index d8565dce..733acee4 100644 --- a/src/flygym/simulation.py +++ b/src/flygym/simulation.py @@ -7,7 +7,7 @@ from jaxtyping import Float from flygym.anatomy import BodySegment -from flygym.compose.fly import ActuatorType +from flygym.compose.fly import BaseFly, ActuatorType from flygym.compose.world import BaseWorld from flygym.rendering import Renderer from flygym.utils.profiling import print_perf_report @@ -760,3 +760,13 @@ def close(self): self.eye_renderer = None # Don't destruct self.retina and self.eye_renderer_scene_option: they can be # reused and retina init requires some IO ops. + + @property + def fly(self) -> BaseFly: + """Return the single fly in the world, or raise an error if there are multiple.""" + return self.world.fly + + @property + def fly_lookup(self) -> dict[str, BaseFly]: + """Return the fly lookup dictionary from the world.""" + return self.world.fly_lookup \ No newline at end of file From eeb31ad2d381043f841ef7f1889fce18d4c24daf Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:51:49 +0200 Subject: [PATCH 04/12] improve typehint for spawn position --- src/flygym/compose/world/base_world.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/flygym/compose/world/base_world.py b/src/flygym/compose/world/base_world.py index 75cefa32..386b1486 100644 --- a/src/flygym/compose/world/base_world.py +++ b/src/flygym/compose/world/base_world.py @@ -100,7 +100,7 @@ def fly(self) -> BaseFly: def _attach_fly_mjcf( self, fly: BaseFly, - spawn_position: Vec3, + spawn_position: Vec3 | tuple[float, float, float], spawn_rotation: Rotation3D, *args, **kwargs, @@ -139,7 +139,7 @@ def _add_skybox(self): def add_fly( self, fly: BaseFly, - spawn_position: Vec3, + spawn_position: Vec3 | tuple[float, float, float], spawn_rotation: Rotation3D, *args: Any, **kwargs: Any, @@ -260,7 +260,7 @@ class _GroundContactMixin: def _attach_fly_mjcf( self, fly: BaseFly, - spawn_position: Vec3, + spawn_position: Vec3 | tuple[float, float, float], spawn_rotation: Rotation3D, *, bodysegs_with_ground_contact: ( From e2f7d610cd35f0a1831c11284b29d1b22fff0393 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:52:00 +0200 Subject: [PATCH 05/12] add shorthand of number of dofs, etc --- src/flygym/compose/fly/base_fly.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/flygym/compose/fly/base_fly.py b/src/flygym/compose/fly/base_fly.py index 3d63c183..9f8bf3d2 100644 --- a/src/flygym/compose/fly/base_fly.py +++ b/src/flygym/compose/fly/base_fly.py @@ -269,10 +269,20 @@ def get_bodysegs_order(self) -> list[BodySegment]: """ return list(self.bodyseg_to_mjcfbody.keys()) + @property + def n_bodysegs(self) -> int: + """Number of body segments in this fly.""" + return len(self.bodyseg_to_mjcfbody) + def get_jointdofs_order(self) -> list[JointDOF]: """Same as `get_bodysegs_order()`, but for joint DoFs instead of body segments.""" return list(self.jointdof_to_mjcfjoint.keys()) + @property + def n_jointdofs(self) -> int: + """Number of joint DoFs in this fly.""" + return len(self.jointdof_to_mjcfjoint) + def get_actuated_jointdofs_order( self, actuator_type: "ActuatorType | str" ) -> list[JointDOF]: @@ -282,6 +292,19 @@ def get_actuated_jointdofs_order( actuator_type = ActuatorType(actuator_type) return list(self.jointdof_to_mjcfactuator_by_type[actuator_type].keys()) + @property + def n_actuated_jointdofs(self) -> int: + raise RuntimeError( + "`n_actuated_jointdofs` is ambiguous because there might be different " + "actuator types. Use `get_n_actuated_jointdofs(actuator_type)` instead, " + "similar to `fly.get_actuated_jointdofs_order(actuator_type)`." + ) + + def get_n_actuated_jointdofs(self, actuator_type: "ActuatorType | str") -> int: + """Number of joint DoFs actuated by the specified actuator type.""" + actuator_type = ActuatorType(actuator_type) + return len(self.jointdof_to_mjcfactuator_by_type[actuator_type]) + def get_legs_order(self) -> list[str]: """Get the ordered list of leg position identifiers (same as `anatomy.LEGS`).""" return LEGS From aae031ffee00370f00f8332f71bb2df3fe0e10c0 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:52:31 +0200 Subject: [PATCH 06/12] update world feature stripping function for mujoco-warp 3.10 --- src/flygym/warp/rendering.py | 34 ++++++++++++++++++---------------- 1 file changed, 18 insertions(+), 16 deletions(-) diff --git a/src/flygym/warp/rendering.py b/src/flygym/warp/rendering.py index be88be53..7a900e35 100644 --- a/src/flygym/warp/rendering.py +++ b/src/flygym/warp/rendering.py @@ -423,21 +423,23 @@ def modify_world_for_batch_rendering(world: BaseWorld) -> bool: material.texrepeat = tuple(tr / 1000 for tr in material.texrepeat) is_modified = True - # Add light above each fly explicitly - for body in world.mjcf_root.bodies: - if body.name.split("/")[-1] == "c_thorax": - warnings.warn(f"Adding overhead light for body {body.name}") - body.add_light( - name=body.name.replace("/", "-") + "-overheadlight", - mode=mj.mjtCamLight.mjCAMLIGHT_TRACK, - targetbody=body.name, - pos=(0, 0, 30), - dir=(0, 0, -1), - type=mj.mjtLightType.mjLIGHT_DIRECTIONAL, - ambient=(10, 10, 10), - diffuse=(10, 10, 10), - specular=(0.3, 0.3, 0.3), - ) - is_modified = True + # Add light above each fly explicitly (only until MuJoCo Warp 3.9) + mujoco_warp_version = tuple(int(x) for x in mjw.__version__.split(".")[:2]) + if mujoco_warp_version < (3, 10): + for body in world.mjcf_root.bodies: + if body.name.split("/")[-1] == "c_thorax": + warnings.warn(f"Adding overhead light for body {body.name}") + body.add_light( + name=body.name.replace("/", "-") + "-overheadlight", + mode=mj.mjtCamLight.mjCAMLIGHT_TRACK, + targetbody=body.name, + pos=(0, 0, 30), + dir=(0, 0, -1), + type=mj.mjtLightType.mjLIGHT_DIRECTIONAL, + ambient=(10, 10, 10), + diffuse=(10, 10, 10), + specular=(0.3, 0.3, 0.3), + ) + is_modified = True return is_modified From e1772a2cbd3373ab545dfb1a9efa3a8289052ba5 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 00:53:49 +0200 Subject: [PATCH 07/12] Allow lazy getters to write to existing buffers --- src/flygym/warp/simulation.py | 44 ++++++++++++++++++++++++----------- 1 file changed, 30 insertions(+), 14 deletions(-) diff --git a/src/flygym/warp/simulation.py b/src/flygym/warp/simulation.py index 36c46351..7d8b33ec 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -77,19 +77,22 @@ def reset(self) -> None: @override def get_joint_angles( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_jointdofs"]: """Get joint angles for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_jointdofs)`` in radians, ordered as in ``fly.get_jointdofs_order()``. """ indices = self._wp_intern_qposadrs_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) wp.launch( wp_gather_indexed_cols_2d, dim=(self.n_worlds, indices.size), @@ -99,19 +102,22 @@ def get_joint_angles( @override def get_joint_velocities( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_jointdofs"]: """Get joint velocities for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_jointdofs)`` in radians per second, ordered as in ``fly.get_jointdofs_order()``. """ indices = self._wp_intern_qveladrs_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) wp.launch( wp_gather_indexed_cols_2d, dim=(self.n_worlds, indices.size), @@ -121,19 +127,22 @@ def get_joint_velocities( @override def get_body_positions( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_bodies 3"]: """Get global body positions for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_bodies, 3)`` in mm, ordered as in ``fly.get_bodysegs_order()``. """ indices = self._wp_internal_bodyids_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) wp.launch( wp_gather_indexed_rows_vec3f, dim=(self.n_worlds, indices.size), @@ -143,19 +152,22 @@ def get_body_positions( @override def get_body_rotations( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_bodies 4"]: """Get global body orientations as quaternions for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_bodies, 4)`` (w, x, y, z), ordered as in ``fly.get_bodysegs_order()``. """ indices = self._wp_internal_bodyids_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size, 4), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size, 4), dtype=wp.float32) wp.launch( wp_gather_indexed_rows_quatf, dim=(self.n_worlds, indices.size), @@ -165,19 +177,22 @@ def get_body_rotations( @override def get_site_positions( - self, fly_name: str + self, fly_name: str, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_sites 3"]: """Get global anatomical-joint site positions for all parallel worlds. Args: fly_name: Name of the fly. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_sites, 3)`` in mm, ordered as in ``fly.get_sites_order()``. """ indices = self._wp_internal_siteids_by_fly[fly_name] - dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size, 3), dtype=wp.float32) wp.launch( wp_gather_indexed_rows_vec3f, dim=(self.n_worlds, indices.size), @@ -192,22 +207,23 @@ def time(self) -> float: @override def get_actuator_forces( - self, - fly_name: str, - actuator_type: ActuatorType, + self, fly_name: str, actuator_type: ActuatorType, dst: wp.array | None = None ) -> Float[wp.array, "n_worlds n_actuators"]: """Get actuator forces for all parallel worlds. Args: fly_name: Name of the fly. actuator_type: Type of actuator to query. + dst: Optional warp array to store the result. If not specified, a new array + is allocated. Returns: Warp array of shape ``(n_worlds, n_actuators)``, ordered as in ``fly.get_actuated_jointdofs_order(actuator_type)``. """ indices = self._wp_intern_actuatorids_by_type_by_fly[actuator_type][fly_name] - dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) + if dst is None: + dst = wp.zeros((self.n_worlds, indices.size), dtype=wp.float32) wp.launch( wp_gather_indexed_cols_2d, dim=(self.n_worlds, indices.size), From 5fa6000e4a88aae14eccf237be8dc9f19f27fe12 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 19:56:51 +0200 Subject: [PATCH 08/12] Implement in-place mjw_data reset --- src/flygym/simulation.py | 12 +- src/flygym/warp/simulation.py | 20 +++- src/flygym/warp/utils.py | 199 +++++++++++++++++++++++++++++++++- tests/warp/test_simulation.py | 83 +++++++++++++- tests/warp/test_utils.py | 134 +++++++++++++++++++++++ 5 files changed, 438 insertions(+), 10 deletions(-) diff --git a/src/flygym/simulation.py b/src/flygym/simulation.py index d8565dce..a86ff9d7 100644 --- a/src/flygym/simulation.py +++ b/src/flygym/simulation.py @@ -68,7 +68,17 @@ def __init__(self, world: BaseWorld, *, timestep: float | None = None) -> None: self._total_render_time_ns = 0 def reset(self) -> None: - """Reset simulation and renderer to the neutral keyframe.""" + """Reset simulation and renderer to the neutral keyframe. + + !!! warning + + `reset()` does not update derived kinematic quantities (`xpos`, + `xquat`, `site_xpos`, ...) -- it only restores state fields + (`qpos`, `qvel`, `act`, `ctrl`, `mocap`, `time`). Reading + derived quantities right after `reset()`, before calling `step()`, + does not reflect the reset state. This is consistent with the behavior of + MuJoCo's native `mj_resetData`/`mj_resetDataKeyframe` functions. + """ # Reset physics mj.mj_resetDataKeyframe(self.mj_model, self.mj_data, self._neutral_keyframe_id) diff --git a/src/flygym/warp/simulation.py b/src/flygym/warp/simulation.py index 36c46351..8e11cd0b 100644 --- a/src/flygym/warp/simulation.py +++ b/src/flygym/warp/simulation.py @@ -21,6 +21,7 @@ wp_gather_indexed_cols_2d, wp_gather_indexed_rows_vec3f, wp_gather_indexed_rows_quatf, + reset_data_keyframe, ) @@ -68,12 +69,21 @@ def __init__( @override def reset(self) -> None: - """Reset all parallel worlds to the neutral keyframe.""" + """Reset all parallel worlds to the neutral keyframe. + + !!! warning + + `reset()` does not update derived kinematic quantities (`xpos`, + `xquat`, `site_xpos`, ...); it only restores state fields + (`qpos`, `qvel`, `act`, `ctrl`, `mocap`, `time`). Reading + derived quantities right after `reset()`, before calling `step()`, + returns stale values from before the reset. This is consistent with the + behavior of MuJoCo Warp's native `reset_data()` function. + """ super().reset() - # The superclass call resets CPU-side MuJoCo structs to the neutral keyframe, - # so we need to recreate GPU-side structs to reflect that reset. - self.mjw_model, self.mjw_data = self._mj_structs_to_mjw_structs() - # ... don't call mjw.reset_data() here! That loses the keyframe reset. + reset_data_keyframe( + self.mj_model, self.mjw_model, self.mjw_data, self._neutral_keyframe_id + ) @override def get_joint_angles( diff --git a/src/flygym/warp/utils.py b/src/flygym/warp/utils.py index 5dd6253a..85ba8916 100644 --- a/src/flygym/warp/utils.py +++ b/src/flygym/warp/utils.py @@ -1,6 +1,8 @@ -import warp as wp +from typing import Optional -from mujoco_warp._src.types import RenderContext +import warp as wp +import mujoco as mj +import mujoco_warp as mjw @wp.kernel @@ -153,7 +155,7 @@ def unpack_rgb_kernel_selected_worlds_and_cameras( def get_rgb_selected_worlds_and_cameras( - rc: RenderContext, + rc: mjw.RenderContext, worldids: wp.array(dtype=int), # type: ignore camids: wp.array(dtype=int), # type: ignore rgb_out: wp.array4d(dtype=wp.vec3), # type: ignore @@ -203,3 +205,194 @@ def check_gpu(): "You can specify which GPU to use by setting the 'CUDA_VISIBLE_DEVICES' " "environment variable." ) + + +def reset_data_keyframe( + mj_model: mj.MjModel, + mjw_model: mjw.Model, + mjw_data: mjw.Data, + key: int, + reset: Optional[wp.array] = None, +): + """In-place equivalent of ``mj_resetDataKeyframe`` for a batched MJWarp ``Data``. + + This functionality is not provided natively by MuJoCo Warp, so this is a custom + implementation. Note: this function differs from other Warp utils in that it + requires both the `mujoco.MjModel` object and the `mujoco_warp.types.Model` object. + This is because `key_{qpos,qvel,act,mpos,mquat,ctrl}` are not tracked by the + GPU-side model. + + Args: + mj_model: CPU-side MuJoCo model holding the keyframe to reset to. + mjw_model: GPU-side MuJoCo-Warp model instance. + mjw_data: GPU-side MuJoCo-Warp data instance. + key: Index of the keyframe (in `mj_model`) to reset to. + reset: Optional per-world boolean mask, shape `(mjw_data.nworld,)`. + """ + mjw.reset_data(mjw_model, mjw_data, reset) + + @wp.kernel(module="unique", enable_backward=False) + def reset_time( + # From MjModel: + target_time: float, + # In: + reset_in: wp.array[bool], + # Data out: + time_out: wp.array[float], + ): + worldid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + time_out[worldid] = target_time + + @wp.kernel(module="unique", enable_backward=False) + def reset_qpos( + # From MjModel: + target_qpos: wp.array[float], + # In: + reset_in: wp.array[bool], + # Data out: + qpos_out: wp.array2d[float], + ): + worldid, qid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + qpos_out[worldid, qid] = target_qpos[qid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_qvel( + # From MjModel: + target_qvel: wp.array[float], + # In: + reset_in: wp.array[bool], + # Data out: + qvel_out: wp.array2d[float], + ): + worldid, vid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + qvel_out[worldid, vid] = target_qvel[vid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_activation( + # From MjModel: + target_act: wp.array[float], + # In: + reset_in: wp.array[bool], + # Data out: + act_out: wp.array2d[float], + ): + worldid, aid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + act_out[worldid, aid] = target_act[aid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_mocap( + # From MjModel: + target_mpos: wp.array[wp.vec3], + target_mquat: wp.array[wp.quat], + # From mjwarp Model: + body_mocapid: wp.array[int], + # In: + reset_in: wp.array[bool], + # Data out: + mocap_pos_out: wp.array2d[wp.vec3], + mocap_quat_out: wp.array2d[wp.quat], + ): + worldid, bodyid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + mocapid = body_mocapid[bodyid] + + if mocapid >= 0: + mocap_pos_out[worldid, mocapid] = target_mpos[mocapid] + mocap_quat_out[worldid, mocapid] = target_mquat[mocapid] + + @wp.kernel(module="unique", enable_backward=False) + def reset_control( + # From MjModel: + target_ctrl: wp.array[float], + # In: + reset_in: wp.array[bool], + # Data out: + ctrl_out: wp.array2d[float], + ): + worldid, cid = wp.tid() + + if wp.static(reset is not None): + if not reset_in[worldid]: + return + + ctrl_out[worldid, cid] = target_ctrl[cid] + + reset_input = reset or wp.ones(mjw_data.nworld, dtype=bool) + + target_time = mj_model.key_time[key] + wp.launch( + reset_time, + dim=mjw_data.nworld, + inputs=[target_time, reset_input], + outputs=[mjw_data.time], + ) + + target_qpos = wp.array(mj_model.key_qpos[key], dtype=float) + wp.launch( + reset_qpos, + dim=(mjw_data.nworld, mjw_model.nq), + inputs=[target_qpos, reset_input], + outputs=[mjw_data.qpos], + ) + + target_qvel = wp.array(mj_model.key_qvel[key], dtype=float) + wp.launch( + reset_qvel, + dim=(mjw_data.nworld, mjw_model.nv), + inputs=[target_qvel, reset_input], + outputs=[mjw_data.qvel], + ) + + target_act = wp.array(mj_model.key_act[key], dtype=float) + wp.launch( + reset_activation, + dim=(mjw_data.nworld, mjw_model.na), + inputs=[target_act, reset_input], + outputs=[mjw_data.act], + ) + + target_mpos = wp.array(mj_model.key_mpos[key], dtype=wp.vec3) + target_mquat = wp.array(mj_model.key_mquat[key], dtype=wp.quat) + wp.launch( + reset_mocap, + dim=(mjw_data.nworld, mjw_model.nbody), + inputs=[ + target_mpos, + target_mquat, + mjw_model.body_mocapid, + reset_input, + ], + outputs=[mjw_data.mocap_pos, mjw_data.mocap_quat], + ) + + target_ctrl = wp.array(mj_model.key_ctrl[key], dtype=float) + wp.launch( + reset_control, + dim=(mjw_data.nworld, mjw_model.nu), + inputs=[target_ctrl, reset_input], + outputs=[mjw_data.ctrl], + ) diff --git a/tests/warp/test_simulation.py b/tests/warp/test_simulation.py index 822e86ae..c3136437 100644 --- a/tests/warp/test_simulation.py +++ b/tests/warp/test_simulation.py @@ -3,6 +3,7 @@ import warnings import pytest import numpy as np +import mujoco as mj # These tests require the optional warp (GPU) extra; tag them so they can be # excluded with ``-m "not warp"``, and skip the whole module if warp is absent. @@ -165,6 +166,79 @@ def test_reset_clears_physics_time(self, gpu_bundle): sim.reset() assert sim._total_physics_time_ns == 0 + def test_reset_does_not_reallocate_gpu_structs(self, gpu_bundle): + """reset() should reset mjw_data in place, not rebuild it from scratch.""" + sim, fly, cam = gpu_bundle + mjw_model_before = sim.mjw_model + mjw_data_before = sim.mjw_data + sim.step() + sim.reset() + assert sim.mjw_model is mjw_model_before + assert sim.mjw_data is mjw_data_before + + def test_reset_restores_joint_angles_without_a_subsequent_step(self, gpu_bundle): + """State fields (qpos, ...) are restored by reset() itself, since + reset_data_keyframe writes them directly -- no step() needed.""" + sim, fly, cam = gpu_bundle + sim.reset() + baseline = sim.get_joint_angles(fly.name).numpy().copy() + + # Drive the position actuators away from the neutral pose so stepping + # actually moves the joints (they'd otherwise be held at neutral). + displaced = baseline + 0.3 + sim.set_actuator_inputs(fly.name, ActuatorType.POSITION, displaced) + for _ in range(20): + sim.step() + moved = sim.get_joint_angles(fly.name).numpy() + assert not np.allclose(moved, baseline, atol=1e-3) + + sim.reset() + restored = sim.get_joint_angles(fly.name).numpy() + np.testing.assert_allclose(restored, baseline, atol=1e-5) + + def test_body_positions_are_stale_immediately_after_reset(self, gpu_bundle): + """Mirrors mj_resetDataKeyframe: reset() only restores state fields (qpos, + qvel, act, ctrl, mocap, time). Derived kinematic quantities like xpos are + not recomputed, so immediately after reset() -- before any step() -- they + still hold whatever was there before the reset, not the neutral pose. + """ + sim, fly, cam = gpu_bundle + sim.reset() + sim.step() + after_step = sim.get_body_positions(fly.name).numpy().copy() + + sim.reset() + still_stale = sim.get_body_positions(fly.name).numpy() + np.testing.assert_allclose(still_stale, after_step) + + def test_body_positions_correct_after_reset_and_step(self, gpu_bundle): + """Once step() runs (which begins with a forward pass), derived kinematic + quantities catch up to the reset state.""" + sim, fly, cam = gpu_bundle + for _ in range(20): + sim.step() + sim.reset() + sim.step() + gpu_positions = sim.get_body_positions(fly.name).numpy() + + mj.mj_resetDataKeyframe(sim.mj_model, sim.mj_data, sim._neutral_keyframe_id) + mj.mj_step(sim.mj_model, sim.mj_data) + internal_ids = sim._internal_bodyids_by_fly[fly.name] + cpu_positions = sim.mj_data.xpos[internal_ids, :] + + for world_positions in gpu_positions: + np.testing.assert_allclose(world_positions, cpu_positions, atol=1e-4) + + def test_reset_restores_qpos_to_neutral_keyframe(self, gpu_bundle): + sim, fly, cam = gpu_bundle + for _ in range(5): + sim.step() + sim.reset() + expected = np.tile( + sim.mj_model.key_qpos[sim._neutral_keyframe_id], (sim.n_worlds, 1) + ) + np.testing.assert_allclose(sim.mjw_data.qpos.numpy(), expected, atol=1e-5) + # ============================================================================== # State queries @@ -250,8 +324,15 @@ def test_world0_matches_cpu_site_xpos(self, gpu_bundle_with_joint_sites): sim.reset() sim.step() gpu_site_xpos = sim.get_site_positions(fly.name).numpy()[0] + + # sim.reset() does not forward-compute derived quantities on the CPU side + # either (mirroring mj_resetDataKeyframe), so explicitly reset + step the CPU + # reference here rather than relying on sim.mj_data being pre-populated. + mj.mj_resetDataKeyframe(sim.mj_model, sim.mj_data, sim._neutral_keyframe_id) + mj.mj_step(sim.mj_model, sim.mj_data) cpu_site_xpos = sim.mj_data.site_xpos[sim._internal_siteids_by_fly[fly.name], :] - np.testing.assert_allclose(gpu_site_xpos, cpu_site_xpos, atol=1e-6) + + np.testing.assert_allclose(gpu_site_xpos, cpu_site_xpos, atol=1e-4) # ============================================================================== diff --git a/tests/warp/test_utils.py b/tests/warp/test_utils.py index 45688969..687e9629 100644 --- a/tests/warp/test_utils.py +++ b/tests/warp/test_utils.py @@ -7,6 +7,8 @@ # excluded with ``-m "not warp"``, and skip the whole module if warp is absent. pytestmark = pytest.mark.warp wp = pytest.importorskip("warp") +mjw = pytest.importorskip("mujoco_warp") +mj = pytest.importorskip("mujoco") # ============================================================================== @@ -389,3 +391,135 @@ def test_camids_size_mismatch_raises(self): with pytest.raises(ValueError, match="camids"): get_rgb_selected_worlds_and_cameras(rc, worldids, camids, rgb_out) + + +# ============================================================================== +# reset_data_keyframe +# ============================================================================== + +# A free joint (nq=7, nv=6) plus a hinge joint (nq=1, nv=1) gives a model +# where nq != nv, which is the case that exposed the qpos/qvel kernel bug. +# One actuator (with activation state) exercises act/ctrl, and one mocap +# body exercises mocap_pos/mocap_quat. +_RESET_KEYFRAME_XML = """ + + + + + + + + + + + + + + + + + + + + + +""" + + +@pytest.fixture +def mj_and_mjw(): + mj_model = mj.MjModel.from_xml_string(_RESET_KEYFRAME_XML) + mj_data = mj.MjData(mj_model) + # sanity-check the model actually has nq != nv, which is what this test + # suite cares about exercising. + assert mj_model.nq != mj_model.nv + mjw_model = mjw.put_model(mj_model) + mjw_data = mjw.put_data(mj_model, mj_data, nworld=3) + return mj_model, mjw_model, mjw_data + + +class TestResetDataKeyframe: + def test_resets_all_worlds_to_keyframe_values(self, mj_and_mjw): + from flygym.warp.utils import reset_data_keyframe + + mj_model, mjw_model, mjw_data = mj_and_mjw + + # Perturb data away from both the keyframe and the model defaults so + # a no-op kernel launch would be caught by the assertions below. + mjw_data.qpos.fill_(-1.0) + mjw_data.qvel.fill_(-1.0) + mjw_data.time.fill_(-1.0) + + reset_data_keyframe(mj_model, mjw_model, mjw_data, key=0) + + n = mjw_data.nworld + np.testing.assert_allclose( + mjw_data.qpos.numpy(), np.tile(mj_model.key_qpos[0], (n, 1)) + ) + np.testing.assert_allclose( + mjw_data.qvel.numpy(), np.tile(mj_model.key_qvel[0], (n, 1)) + ) + np.testing.assert_allclose( + mjw_data.act.numpy(), np.tile(mj_model.key_act[0], (n, 1)) + ) + np.testing.assert_allclose( + mjw_data.ctrl.numpy(), np.tile(mj_model.key_ctrl[0], (n, 1)) + ) + np.testing.assert_allclose( + mjw_data.mocap_pos.numpy()[:, 0, :], + np.tile(mj_model.key_mpos[0], (n, 1)), + ) + np.testing.assert_allclose( + mjw_data.mocap_quat.numpy()[:, 0, :], + np.tile(mj_model.key_mquat[0], (n, 1)), + ) + np.testing.assert_allclose(mjw_data.time.numpy(), [0.5] * n) + + def test_qvel_write_does_not_overrun_into_neighboring_world(self, mj_and_mjw): + """Regression test: qpos has nq=8 columns but qvel only has nv=7. + + A kernel that loops over nq columns and writes both qpos_out and + qvel_out at the same column index writes one element past the end + of each world's qvel row. Because qvel is stored as a flat + (nworld, nv) buffer, that out-of-bounds column aliases column 0 of + the next world's row. qpos and qvel are reset by separate kernels, + each launched with its own dimension (nq vs. nv). + """ + from flygym.warp.utils import reset_data_keyframe + + mj_model, mjw_model, mjw_data = mj_and_mjw + + mjw_data.qvel.fill_(-1.0) + reset_input = wp.array([True, False, True], dtype=bool) + reset_data_keyframe(mj_model, mjw_model, mjw_data, key=0, reset=reset_input) + + qvel = mjw_data.qvel.numpy() + np.testing.assert_allclose(qvel[0], mj_model.key_qvel[0]) + np.testing.assert_allclose(qvel[2], mj_model.key_qvel[0]) + # World 1 was excluded from the reset and must be untouched. + np.testing.assert_allclose(qvel[1], np.full(mj_model.nv, -1.0)) + + def test_partial_reset_leaves_unselected_worlds_untouched(self, mj_and_mjw): + from flygym.warp.utils import reset_data_keyframe + + mj_model, mjw_model, mjw_data = mj_and_mjw + + mjw_data.qpos.fill_(0.0) + mjw_data.time.fill_(99.0) + reset_input = wp.array([True, False, True], dtype=bool) + + reset_data_keyframe(mj_model, mjw_model, mjw_data, key=0, reset=reset_input) + + qpos = mjw_data.qpos.numpy() + time = mjw_data.time.numpy() + np.testing.assert_allclose(qpos[0], mj_model.key_qpos[0]) + np.testing.assert_allclose(qpos[2], mj_model.key_qpos[0]) + np.testing.assert_allclose(qpos[1], np.zeros(mj_model.nq)) + assert time[1] == pytest.approx(99.0) + assert time[0] == pytest.approx(0.5) + assert time[2] == pytest.approx(0.5) From 44f1fd5aeed26a84603abcee327a59cb95203335 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Thu, 2 Jul 2026 19:59:17 +0200 Subject: [PATCH 09/12] add test for nonzero position after reset --- tests/core/test_simulation.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/core/test_simulation.py b/tests/core/test_simulation.py index 2376bc60..4451e509 100644 --- a/tests/core/test_simulation.py +++ b/tests/core/test_simulation.py @@ -153,6 +153,14 @@ def test_correct_number_of_bodies(self, simulation, fly_with_adhesion): pos = simulation.get_body_positions(fly_with_adhesion.name) assert pos.shape[0] == len(ALL_SEGMENT_NAMES) + def test_positions_correct_after_reset_and_step( + self, simulation, fly_with_adhesion + ): + simulation.reset() + simulation.step() + pos = simulation.get_body_positions(fly_with_adhesion.name) + assert not np.allclose(pos, 0.0) + # ============================================================================== # get_body_rotations From 88bb5a1a015e034fa1f06d53127de5a63acd9a6b Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Fri, 3 Jul 2026 13:56:03 +0200 Subject: [PATCH 10/12] Use tarred assets on the remote s3 bucket. This simplifies lazy asset loading. Add script to generate tarballs and checksum files --- .../dev/make_tar_for_lazy_loaded_assets.sh | 34 +++ src/flygym/utils/assets_lazy_loading.py | 244 ++++++------------ tests/core/test_assets.py | 113 +++----- 3 files changed, 152 insertions(+), 239 deletions(-) create mode 100755 scripts/dev/make_tar_for_lazy_loaded_assets.sh diff --git a/scripts/dev/make_tar_for_lazy_loaded_assets.sh b/scripts/dev/make_tar_for_lazy_loaded_assets.sh new file mode 100755 index 00000000..eada1eed --- /dev/null +++ b/scripts/dev/make_tar_for_lazy_loaded_assets.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Package a lazily-loaded asset directory (see +# src/flygym/utils/assets_lazy_loading.py) into the .tar + .checksum +# pair expected on the S3 bucket. The two output files are written next to the +# input folder, named after it, ready to be uploaded as-is. +# +# Usage: scripts/dev/make_tar_for_lazy_loaded_assets.sh +set -e + +if [ $# -ne 1 ]; then + echo "Usage: $0 " >&2 + exit 1 +fi + +src_dir="$1" +if [ ! -d "$src_dir" ]; then + echo "Error: '$src_dir' is not a directory." >&2 + exit 1 +fi + +# Strip a trailing slash so basename gives the asset's name, not "". +src_dir="${src_dir%/}" +name="$(basename "$src_dir")" +out_dir="$(dirname "$src_dir")" +tar_path="$out_dir/$name.tar" +checksum_path="$out_dir/$name.checksum" + +# Archive the directory's *contents*, not the directory itself, so archive +# members are bare file names (lazy_load_asset_dir extracts straight into the +# cache dir and expects files there directly, e.g. `mesh_dir / "a.stl"`). +tar -cf "$tar_path" -C "$src_dir" . +sha256sum "$tar_path" | cut -d' ' -f1 > "$checksum_path" + +echo "Wrote $tar_path ($(du -h "$tar_path" | cut -f1)) and $checksum_path ($(cat "$checksum_path"))" diff --git a/src/flygym/utils/assets_lazy_loading.py b/src/flygym/utils/assets_lazy_loading.py index 3bea4dc4..e161d9f7 100644 --- a/src/flygym/utils/assets_lazy_loading.py +++ b/src/flygym/utils/assets_lazy_loading.py @@ -1,39 +1,29 @@ """Lazy download of large asset files from a public S3 bucket. -Most asset files (configs, poses, and the simplified default NeuroMechFly meshes) -are small enough to ship inside the ``flygym`` package. The high-resolution -``fullsize`` meshes -- especially the FlyBody ``.obj`` meshes, which are an order -of magnitude larger than everything else combined -- would bloat the package and -the git repository, so they are hosted on an institution-managed S3 bucket and -pulled in *the first time they are needed*, similar to how PyTorch downloads -pretrained weights. - -Downloaded files are cached on disk (see :func:`get_cache_root`) so the download -happens only once per machine. The bucket is public and served over a standard -S3-compatible HTTP endpoint, so plain ``urllib`` is enough -- no extra -dependencies (boto3 etc.) are required. - -The bucket stores each remotely hosted asset directory as a flat, *versioned* -sub-prefix of :data:`S3_ROOT_PREFIX`, so future revisions can be uploaded under a -new name without disturbing existing releases. Bump the version constants below to -point a release at a new version. Example: - - bucket: flygym_assets/neuromechfly_fullsize_meshes_20260623a/ - cache: ~/.cache/flygym_assets/neuromechfly_fullsize_meshes_20260623a/ +The high-resolution `fullsize` meshes are too large to ship inside the +`flygym` package, so they are hosted on a public S3 bucket and downloaded the +first time they are needed, then cached on disk (see :func:`get_cache_root`). +The bucket is served over plain HTTP(S), so `urllib` is enough -- no boto3. + +Each asset directory lives on the bucket as a `.tar` archive plus a +`.checksum` sidecar holding the tar's sha256 hex digest, both generated +by `scripts/dev/make_tar_for_lazy_loaded_assets.sh`. Names are versioned +(e.g. `neuromechfly_fullsize_meshes_20260623a`) so new revisions can be +uploaded without disturbing existing releases; bump the version constants in +the fly model modules to point a release at a new asset set. """ import hashlib import os -import shutil +import tarfile import tempfile from pathlib import Path from urllib.parse import quote from urllib.request import urlopen -from xml.etree import ElementTree from loguru import logger -__all__ = ["get_cache_root", "lazy_load_asset_dir", "prefetch_meshes"] +__all__ = ["get_cache_root", "lazy_load_asset_dir", "download_all_assets"] #: Base HTTP(S) endpoint of the S3-compatible object store. S3_ENDPOINT = "https://datasets.epfl.ch" @@ -42,29 +32,26 @@ #: Top-level key prefix within the bucket under which all assets live. S3_ROOT_PREFIX = "flygym_assets" - -# S3 ListObjectsV2 responses are namespaced; this is the namespace MinIO/S3 use. -_S3_XML_NS = {"s3": "http://s3.amazonaws.com/doc/2006-03-01/"} +#: Read/write files in chunks of this size while streaming a download. +_CHUNK_SIZE = 1024 * 1024 +#: How many times to (re)try a download before giving up. A dropped connection +#: yields a truncated tar that fails the checksum; the endpoint is flaky enough +#: that a single such failure shouldn't abort the whole run. +_MAX_ATTEMPTS = 3 +#: Per-request timeout (seconds). Bounds how long a stalled connection can hang +#: before it errors out and the attempt is retried, rather than blocking forever. +_TIMEOUT = 30 def get_cache_root() -> Path: - """Return the directory under which downloaded assets are cached. - - Resolution order: - - 1. ``$FLYGYM_ASSET_CACHE_DIR`` if set (useful for CI caching or shared, - read-only installs); - 2. ``$XDG_CACHE_HOME/flygym_assets`` if ``XDG_CACHE_HOME`` is set; - 3. ``~/.cache/flygym_assets`` otherwise. - - The directory is named ``flygym_assets`` to match the bucket's top-level - prefix (:data:`S3_ROOT_PREFIX`). + """Return the directory under which downloaded assets are cached: + `$FLYGYM_ASSET_CACHE_DIR` if set (useful for CI caching), else + `$XDG_CACHE_HOME/flygym_assets`, else `~/.cache/flygym_assets`. """ env = os.environ.get("FLYGYM_ASSET_CACHE_DIR") if env: return Path(env).expanduser() - # Per the XDG Base Directory spec, a relative XDG_CACHE_HOME is invalid and - # must be ignored (as is an unset/empty value). + # Per the XDG spec, a relative XDG_CACHE_HOME is invalid and must be ignored. xdg = os.environ.get("XDG_CACHE_HOME") if xdg and os.path.isabs(xdg): return Path(xdg) / S3_ROOT_PREFIX @@ -76,106 +63,39 @@ def _object_url(key: str) -> str: return f"{S3_ENDPOINT}/{S3_BUCKET}/{quote(key)}" -def _list_s3_prefix(prefix: str) -> list[dict]: - """List every object under ``prefix`` via the public ListObjectsV2 API. - - Returns a list of ``{"key", "size", "etag"}`` dicts. Handles pagination via - continuation tokens. The bucket is public, so the request is unsigned. - """ - if not prefix.endswith("/"): - prefix += "/" - objects: list[dict] = [] - continuation_token: str | None = None - while True: - url = f"{S3_ENDPOINT}/{S3_BUCKET}?list-type=2&prefix={quote(prefix, safe='')}" - if continuation_token is not None: - url += f"&continuation-token={quote(continuation_token, safe='')}" - with urlopen(url) as response: - tree = ElementTree.fromstring(response.read()) - for contents in tree.findall("s3:Contents", _S3_XML_NS): - key = contents.findtext("s3:Key", namespaces=_S3_XML_NS) - if key is None or key.endswith("/"): - continue # skip "directory" placeholder keys - size = int(contents.findtext("s3:Size", default="0", namespaces=_S3_XML_NS)) - etag = contents.findtext("s3:ETag", default="", namespaces=_S3_XML_NS) - objects.append({"key": key, "size": size, "etag": etag.strip('"')}) - is_truncated = ( - tree.findtext("s3:IsTruncated", default="false", namespaces=_S3_XML_NS) - == "true" - ) - if not is_truncated: - break - continuation_token = tree.findtext( - "s3:NextContinuationToken", namespaces=_S3_XML_NS - ) - if not continuation_token: - break - return objects - - -def _is_up_to_date(path: Path, size: int, etag: str) -> bool: - """Return True if ``path`` already holds the object described by (size, etag). - - For non-multipart uploads the S3 ETag is the MD5 hex digest of the content, - which we verify. Multipart ETags contain a ``-`` and are not plain MD5, so we - fall back to a size check for those. +def _download_tar(name: str, dest: Path) -> None: + """Download `.tar` to `dest` and verify it against `.checksum`, + retrying on transient network errors and truncated (checksum-mismatched) + downloads. """ - if not path.is_file(): - return False - if path.stat().st_size != size: - return False - if etag and "-" not in etag: - digest = hashlib.md5(path.read_bytes()).hexdigest() - return digest == etag - return True - - -def _download_object(key: str, dest: Path, size: int, etag: str) -> None: - """Download a single object to ``dest`` atomically.""" - dest.parent.mkdir(parents=True, exist_ok=True) - fd, tmp_name = tempfile.mkstemp(dir=dest.parent, suffix=".part") - tmp_path = Path(tmp_name) - try: - with os.fdopen(fd, "wb") as out, urlopen(_object_url(key)) as response: - shutil.copyfileobj(response, out) - if not _is_up_to_date(tmp_path, size, etag): - raise OSError( - f"Downloaded asset failed integrity check: {key} " - f"(expected {size} bytes, etag {etag!r})" - ) - tmp_path.replace(dest) - finally: - tmp_path.unlink(missing_ok=True) - + checksum_url = _object_url(f"{S3_ROOT_PREFIX}/{name}.checksum") + with urlopen(checksum_url, timeout=_TIMEOUT) as response: + expected = response.read().decode().split()[0] -def _download_prefix(s3_prefix: str, dest_dir: Path) -> Path: - """Download every object under ``s3_prefix`` into ``dest_dir`` (skipping files - that are already present and up to date). Returns ``dest_dir``. - """ - objects = _list_s3_prefix(s3_prefix) - if not objects: - raise FileNotFoundError( - f"No assets found on S3 under prefix '{s3_prefix}'. The bucket may be " - "unreachable or the asset may have been moved." - ) - prefix = s3_prefix if s3_prefix.endswith("/") else s3_prefix + "/" - pending = [] - for obj in objects: - rel_key = obj["key"][len(prefix) :] - dest = dest_dir / rel_key - if not _is_up_to_date(dest, obj["size"], obj["etag"]): - pending.append((obj, dest)) - - if pending: - total_mb = sum(obj["size"] for obj, _ in pending) / 1e6 - logger.info( - f"Downloading {len(pending)} FlyGym asset file(s) " - f"({total_mb:.1f} MB) from S3 to {dest_dir} (one-time download)..." - ) - for obj, dest in pending: - _download_object(obj["key"], dest, obj["size"], obj["etag"]) - logger.info("Finished downloading FlyGym assets.") - return dest_dir + tar_url = _object_url(f"{S3_ROOT_PREFIX}/{name}.tar") + for attempt in range(1, _MAX_ATTEMPTS + 1): + digest = hashlib.sha256() + try: + with ( + urlopen(tar_url, timeout=_TIMEOUT) as response, + open(dest, "wb") as out, + ): + while chunk := response.read(_CHUNK_SIZE): + digest.update(chunk) + out.write(chunk) + except OSError as e: + reason = f"download failed ({e})" + else: + if digest.hexdigest() == expected: + return + reason = ( + f"integrity check failed " + f"(expected sha256 {expected}, got {digest.hexdigest()})" + ) + dest.unlink(missing_ok=True) + if attempt == _MAX_ATTEMPTS: + raise OSError(f"Could not download {name}.tar: {reason}") + logger.warning(f"Retrying {name}.tar ({attempt}/{_MAX_ATTEMPTS}): {reason}") def lazy_load_asset_dir(rel_path: os.PathLike | str) -> Path: @@ -183,43 +103,45 @@ def lazy_load_asset_dir(rel_path: os.PathLike | str) -> Path: from S3 on first use. Args: - rel_path: Path of the directory within the bucket, relative to - :data:`S3_ROOT_PREFIX` (e.g. ``"neuromechfly_fullsize_meshes_20260623a"``, - as defined by each fly model's ``*_FULLSIZE_MESH_DIR`` constant). - - The directory is cached under :func:`get_cache_root` keyed by ``rel_path``. If - the cached copy already exists it is returned as-is (no network access); - otherwise the whole directory is downloaded into a temporary location and moved - into place atomically, so an interrupted or concurrent download never leaves a - partial cache. - - Raises: - FileNotFoundError: If ``rel_path`` does not exist in the bucket. + rel_path: Name of the asset set on the bucket, i.e. the shared stem of + `.tar` and `.checksum` under + :data:`S3_ROOT_PREFIX` (as defined by each fly model's + `*_MESH_DIR` constant). + + If the cached copy already exists it is returned as-is (no network access). + Otherwise the tar is downloaded, verified, and extracted inside a temporary + directory that is moved into place atomically, so an interrupted or + concurrent download never leaves a partial cache. """ - rel_path = Path(rel_path) - cache_dir = get_cache_root() / rel_path + name = Path(rel_path).as_posix() + cache_dir = get_cache_root() / name if cache_dir.is_dir(): return cache_dir + logger.info(f"Downloading FlyGym asset '{name}' from S3 (one-time download)...") cache_dir.parent.mkdir(parents=True, exist_ok=True) - staging = Path(tempfile.mkdtemp(dir=cache_dir.parent, suffix=".partial")) - try: - _download_prefix(f"{S3_ROOT_PREFIX}/{rel_path.as_posix()}", staging) + with tempfile.TemporaryDirectory( + dir=cache_dir.parent, suffix=".partial" + ) as staging: + staging = Path(staging) + _download_tar(name, staging / "asset.tar") + extracted = staging / "extracted" + with tarfile.open(staging / "asset.tar") as tar: + tar.extractall(extracted, filter="data") try: - staging.replace(cache_dir) + extracted.replace(cache_dir) except OSError: # Another process finished downloading the same asset while we were - # working: os.replace cannot move onto the now-populated directory. - # Their copy is equivalent to ours, so use it instead of failing. + # working, so cache_dir is now populated and cannot be replaced. + # Their copy is equivalent to ours: use it instead of failing. if not cache_dir.is_dir(): raise - finally: - shutil.rmtree(staging, ignore_errors=True) + logger.info(f"Finished downloading FlyGym asset '{name}'.") return cache_dir -def prefetch_meshes() -> list[Path]: - """Eagerly download all remotely hosted meshes into the cache. +def download_all_assets() -> list[Path]: + """Eagerly download all remotely hosted assets into the cache. Useful for warming a CI cache or preparing an offline environment. Returns the list of local directories that now hold the assets. diff --git a/tests/core/test_assets.py b/tests/core/test_assets.py index d5512497..a3046c83 100644 --- a/tests/core/test_assets.py +++ b/tests/core/test_assets.py @@ -7,6 +7,7 @@ import hashlib import io +import tarfile from pathlib import Path import pytest @@ -56,99 +57,64 @@ def test_get_cache_root_ignores_relative_xdg(self, monkeypatch): ) -class TestIsUpToDate: - def test_missing_file(self, tmp_path): - assert not assets_lazy_loading._is_up_to_date(tmp_path / "nope", 1, "x") - - def test_size_mismatch(self, tmp_path): - p = tmp_path / "f" - p.write_bytes(b"abc") - assert not assets_lazy_loading._is_up_to_date(p, 99, "") - - def test_md5_match_and_mismatch(self, tmp_path): - p = tmp_path / "f" - data = b"hello world" - p.write_bytes(data) - good = hashlib.md5(data).hexdigest() - assert assets_lazy_loading._is_up_to_date(p, len(data), good) - assert not assets_lazy_loading._is_up_to_date(p, len(data), "0" * 32) - - def test_multipart_etag_falls_back_to_size(self, tmp_path): - p = tmp_path / "f" - data = b"hello world" - p.write_bytes(data) - # Multipart ETags contain a dash and are not a plain MD5; size match wins. - assert assets_lazy_loading._is_up_to_date(p, len(data), "deadbeef-2") - - # --------------------------------------------------------------------------- # Download path, exercised offline via a fake "remote" # --------------------------------------------------------------------------- -# Flat, versioned S3 sub-prefix -- mirrors the real bucket layout. +# Flat, versioned S3 key stem -- mirrors the real bucket layout. _DEMO_VERSION = "demo_fullsize_meshes_vtest" -_DEMO_PREFIX = f"flygym_assets/{_DEMO_VERSION}/" + + +def _make_tar_bytes(files: dict) -> bytes: + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w") as tar: + for name, data in files.items(): + info = tarfile.TarInfo(name=name) + info.size = len(data) + tar.addfile(info, io.BytesIO(data)) + return buf.getvalue() @pytest.fixture def fake_remote(monkeypatch): """Serve a small in-memory object store through the assets module's network - seam, so the full download/verify/cache flow runs without touching S3. + seam, so the full download/verify/extract/cache flow runs without touching S3. """ + tar_bytes = _make_tar_bytes({"a.stl": b"aaaa", "b.stl": b"bbbbbb"}) + checksum = hashlib.sha256(tar_bytes).hexdigest() store = { - f"{_DEMO_PREFIX}a.stl": b"aaaa", - f"{_DEMO_PREFIX}b.stl": b"bbbbbb", + f"flygym_assets/{_DEMO_VERSION}.tar": tar_bytes, + f"flygym_assets/{_DEMO_VERSION}.checksum": checksum.encode("ascii"), } - objects = [ - {"key": k, "size": len(v), "etag": hashlib.md5(v).hexdigest()} - for k, v in store.items() - ] - def fake_list(prefix): - if not prefix.endswith("/"): - prefix += "/" - return [o for o in objects if o["key"].startswith(prefix)] - - def fake_urlopen(url): + def fake_urlopen(url, **kwargs): # url is the object URL; recover the key after the bucket name. marker = f"/{assets_lazy_loading.S3_BUCKET}/" key = url.split(marker, 1)[1] return io.BytesIO(store[key]) - monkeypatch.setattr(assets_lazy_loading, "_list_s3_prefix", fake_list) monkeypatch.setattr(assets_lazy_loading, "urlopen", fake_urlopen) return store -def test_download_prefix_writes_and_is_idempotent(fake_remote, tmp_path): - dest = tmp_path / "out" - assets_lazy_loading._download_prefix(_DEMO_PREFIX, dest) - assert (dest / "a.stl").read_bytes() == b"aaaa" - assert (dest / "b.stl").read_bytes() == b"bbbbbb" - - # Re-running must not re-download (no .part temp files left behind, content - # unchanged) since everything is already up to date. - assets_lazy_loading._download_prefix(_DEMO_PREFIX, dest) - assert sorted(p.name for p in dest.iterdir()) == ["a.stl", "b.stl"] - - -def test_download_object_integrity_check(fake_remote, tmp_path, monkeypatch): - dest = tmp_path / "corrupt.stl" - # Lie about the expected size so the post-download check fails. +def test_download_tar_integrity_check(fake_remote, tmp_path, monkeypatch): + # Corrupt the checksum so the post-download check fails. + fake_remote[f"flygym_assets/{_DEMO_VERSION}.checksum"] = b"0" * 64 + dest = tmp_path / "out.tar" with pytest.raises(OSError, match="integrity check"): - assets_lazy_loading._download_object( - f"{_DEMO_PREFIX}a.stl", dest, size=999, etag="" - ) + assets_lazy_loading._download_tar(_DEMO_VERSION, dest) assert not dest.exists() # nothing left behind on failure def test_lazy_load_asset_dir_downloads_and_caches(fake_remote, tmp_path, monkeypatch): monkeypatch.setenv("FLYGYM_ASSET_CACHE_DIR", str(tmp_path / "cache")) - # First call downloads from S3 and caches under the versioned dir name. + # First call downloads from S3, verifies, extracts, and caches under the + # versioned dir name. out = assets_lazy_loading.lazy_load_asset_dir(_DEMO_VERSION) assert out == tmp_path / "cache" / _DEMO_VERSION assert (out / "a.stl").read_bytes() == b"aaaa" + assert (out / "b.stl").read_bytes() == b"bbbbbb" assert not list(out.parent.glob("*.partial")), "staging dir not cleaned up" @@ -162,7 +128,7 @@ def test_lazy_load_asset_dir_returns_cache_without_network(tmp_path, monkeypatch def boom(*args, **kwargs): raise AssertionError("must not hit the network when already cached") - monkeypatch.setattr(assets_lazy_loading, "_download_prefix", boom) + monkeypatch.setattr(assets_lazy_loading, "_download_tar", boom) assert assets_lazy_loading.lazy_load_asset_dir(_DEMO_VERSION) == cached @@ -175,7 +141,7 @@ def test_lazy_load_asset_dir_interrupted_download_leaves_no_cache( def boom(*args, **kwargs): raise RuntimeError("network died mid-download") - monkeypatch.setattr(assets_lazy_loading, "_download_prefix", boom) + monkeypatch.setattr(assets_lazy_loading, "_download_tar", boom) with pytest.raises(RuntimeError): assets_lazy_loading.lazy_load_asset_dir(_DEMO_VERSION) assert not (tmp_path / "cache" / _DEMO_VERSION).exists() @@ -222,7 +188,7 @@ def test_prefetch_meshes_covers_all_remote_sets(monkeypatch): "lazy_load_asset_dir", lambda rel: requested.append(str(rel)) or Path("/cache") / rel, ) - assets_lazy_loading.prefetch_meshes() + assets_lazy_loading.download_all_assets() assert set(requested) == set(_remote_mesh_dirs().values()) @@ -230,21 +196,12 @@ def test_prefetch_meshes_covers_all_remote_sets(monkeypatch): @pytest.mark.parametrize( "mesh_dir", _remote_mesh_dirs().values(), ids=_remote_mesh_dirs().keys() ) -def test_real_s3_roundtrip(mesh_dir): - """Opt-in: list and download a single small object from the live bucket, for - each model's remote mesh set.""" - prefix = f"{assets_lazy_loading.S3_ROOT_PREFIX}/{mesh_dir}/" +def test_real_s3_roundtrip(mesh_dir, tmp_path): + """Opt-in: download and checksum-verify the real tar for each model's remote + mesh set from the live bucket.""" try: - objects = assets_lazy_loading._list_s3_prefix(prefix) + dest = tmp_path / f"{mesh_dir}.tar" + assets_lazy_loading._download_tar(mesh_dir, dest) except Exception as e: # network unavailable in this environment pytest.skip(f"S3 endpoint unreachable: {e}") - assert objects, f"expected objects under prefix {prefix!r}" - smallest = min(objects, key=lambda o: o["size"]) - import tempfile - - with tempfile.TemporaryDirectory() as d: - dest = Path(d) / "obj.stl" - assets_lazy_loading._download_object( - smallest["key"], dest, smallest["size"], smallest["etag"] - ) - assert dest.stat().st_size == smallest["size"] + assert dest.stat().st_size > 0 From 4c8640ba310bb1dd87ce13150fe40b7bcd694db4 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Fri, 3 Jul 2026 14:54:08 +0200 Subject: [PATCH 11/12] update warp type hints to be #484 compliant --- scripts/replay_behavior_gpu.py | 8 ++-- src/flygym/warp/utils.py | 46 +++++++++---------- .../benchmark/time_gpu_simulation.py | 8 ++-- 3 files changed, 31 insertions(+), 31 deletions(-) diff --git a/scripts/replay_behavior_gpu.py b/scripts/replay_behavior_gpu.py index edb960bd..54b1119d 100644 --- a/scripts/replay_behavior_gpu.py +++ b/scripts/replay_behavior_gpu.py @@ -79,10 +79,10 @@ @wp.kernel def record_joint_angles_kernel( - qpos: wp.array2d(dtype=wp.float32), # type: ignore # (n_worlds, nq) - qpos_adrs: wp.array(dtype=wp.int32), # type: ignore # (n_jointdofs,) - step_counter: wp.array(dtype=wp.int32), # type: ignore - recorded: wp.array3d(dtype=wp.float32), # type: ignore # (n_steps, n_worlds, n_dofs) + qpos: wp.array2d[float], # (n_worlds, nq) + qpos_adrs: wp.array[int], # (n_jointdofs,) + step_counter: wp.array[int], + recorded: wp.array3d[float], # (n_steps, n_worlds, n_dofs) ): """Gather this step's joint angles into a pre-allocated, GPU-resident buffer. diff --git a/src/flygym/warp/utils.py b/src/flygym/warp/utils.py index 85ba8916..628c51d9 100644 --- a/src/flygym/warp/utils.py +++ b/src/flygym/warp/utils.py @@ -7,9 +7,9 @@ @wp.kernel def wp_gather_indexed_rows_3d( - src: wp.array3d(dtype=wp.float32), # type: ignore - dst: wp.array3d(dtype=wp.float32), # type: ignore - rows: wp.array(dtype=wp.int32), # type: ignore + src: wp.array3d[float], + dst: wp.array3d[float], + rows: wp.array[int], ): """Gather specific rows (dim 1) from a 3D Warp array into a narrower destination. @@ -30,9 +30,9 @@ def wp_gather_indexed_rows_3d( @wp.kernel def wp_gather_indexed_rows_vec3f( - src: wp.array2d(dtype=wp.vec3f), # type: ignore - dst: wp.array3d(dtype=wp.float32), # type: ignore - rows: wp.array(dtype=wp.int32), # type: ignore + src: wp.array2d[wp.vec3], + dst: wp.array3d[float], + rows: wp.array[int], ): """Gather specific rows from a 2D ``vec3f`` array into a ``(n_worlds, n_rows_narrow, 3)`` ``float32`` destination. @@ -57,9 +57,9 @@ def wp_gather_indexed_rows_vec3f( @wp.kernel def wp_gather_indexed_rows_quatf( - src: wp.array2d(dtype=wp.quatf), # type: ignore - dst: wp.array3d(dtype=wp.float32), # type: ignore - rows: wp.array(dtype=wp.int32), # type: ignore + src: wp.array2d[wp.quat], + dst: wp.array3d[float], + rows: wp.array[int], ): """Gather specific rows from a 2D ``quatf`` array into a ``(n_worlds, n_rows_narrow, 4)`` ``float32`` destination. @@ -85,9 +85,9 @@ def wp_gather_indexed_rows_quatf( @wp.kernel def wp_scatter_indexed_cols_2d( - src: wp.array2d(dtype=wp.float32), # type: ignore - dst: wp.array2d(dtype=wp.float32), # type: ignore - cols: wp.array(dtype=wp.int32), # type: ignore + src: wp.array2d[float], + dst: wp.array2d[float], + cols: wp.array[int], ): """Scatter a 2D Warp array into specific columns of a wider destination array. @@ -108,9 +108,9 @@ def wp_scatter_indexed_cols_2d( @wp.kernel def wp_gather_indexed_cols_2d( - src: wp.array2d(dtype=wp.float32), # type: ignore - dst: wp.array2d(dtype=wp.float32), # type: ignore - cols: wp.array(dtype=wp.int32), # type: ignore + src: wp.array2d[float], + dst: wp.array2d[float], + cols: wp.array[int], ): """Gather specific columns from a 2D Warp array into a narrower destination array. @@ -132,12 +132,12 @@ def wp_gather_indexed_cols_2d( @wp.kernel def unpack_rgb_kernel_selected_worlds_and_cameras( # In: - packed: wp.array2d(dtype=wp.uint32), # type: ignore - rgb_adr: wp.array(dtype=int), # type: ignore - worldids_to_render: wp.array(dtype=int), # type: ignore - camids_to_render: wp.array(dtype=int), # type: ignore + packed: wp.array2d[wp.uint32], + rgb_adr: wp.array[int], + worldids_to_render: wp.array[int], + camids_to_render: wp.array[int], # Out: - rgb_out: wp.array4d(dtype=wp.vec3), # type: ignore + rgb_out: wp.array4d[wp.vec3], ): """Unpack ABGR uint32 packed pixel data into separate R, G, and B channels.""" idx_within_worldids, idx_within_camids, pixelid = wp.tid() @@ -156,9 +156,9 @@ def unpack_rgb_kernel_selected_worlds_and_cameras( def get_rgb_selected_worlds_and_cameras( rc: mjw.RenderContext, - worldids: wp.array(dtype=int), # type: ignore - camids: wp.array(dtype=int), # type: ignore - rgb_out: wp.array4d(dtype=wp.vec3), # type: ignore + worldids: wp.array[int], + camids: wp.array[int], + rgb_out: wp.array4d[wp.vec3], ): """Get the RGB data output from the render context buffers for the selected worlds and cameras. diff --git a/src/flygym_demo/benchmark/time_gpu_simulation.py b/src/flygym_demo/benchmark/time_gpu_simulation.py index 5693bc29..57d309b0 100644 --- a/src/flygym_demo/benchmark/time_gpu_simulation.py +++ b/src/flygym_demo/benchmark/time_gpu_simulation.py @@ -88,9 +88,9 @@ def make_target_angles_all_worlds(self, n_worlds: int, sim_steps: int): @wp.kernel def update_target_angles_kernel( - dof_angles_all_worlds_gpu: wp.array3d(dtype=wp.float32), # type: ignore - step_counter_gpu: wp.array(dtype=wp.int32), # type: ignore - curr_target_angles_gpu: wp.array2d(dtype=wp.float32), # type: ignore + dof_angles_all_worlds_gpu: wp.array3d[float], + step_counter_gpu: wp.array[int], + curr_target_angles_gpu: wp.array2d[float], ): world_id, actuator_id = wp.tid() step = step_counter_gpu[0] @@ -100,7 +100,7 @@ def update_target_angles_kernel( @wp.kernel def increment_counter_kernel( - step_counter_gpu: wp.array(dtype=wp.int32), # type: ignore + step_counter_gpu: wp.array[int], ): step_counter_gpu[0] = step_counter_gpu[0] + 1 From 1ecf8bf7e438dc2b765846659839494269ca1299 Mon Sep 17 00:00:00 2001 From: Sibo Wang Date: Fri, 3 Jul 2026 15:24:21 +0200 Subject: [PATCH 12/12] add array dimensions as typevars --- src/flygym/compose/world/base_world.py | 2 +- src/flygym/simulation.py | 17 ++++++++-------- src/flygym/utils/typing.py | 19 ++++++++++++++++++ src/flygym/warp/utils.py | 20 +++++-------------- .../benchmark/time_gpu_simulation.py | 4 +--- 5 files changed, 35 insertions(+), 27 deletions(-) create mode 100644 src/flygym/utils/typing.py diff --git a/src/flygym/compose/world/base_world.py b/src/flygym/compose/world/base_world.py index 386b1486..e3eb8b44 100644 --- a/src/flygym/compose/world/base_world.py +++ b/src/flygym/compose/world/base_world.py @@ -81,7 +81,7 @@ def mjcf_root(self) -> mj.MjSpec: def fly_lookup(self) -> dict[str, BaseFly]: """Lookup for `Fly` objects in the world, keyed by fly name.""" return self._fly_lookup - + @property def fly(self) -> BaseFly: """Get the single fly in the world. diff --git a/src/flygym/simulation.py b/src/flygym/simulation.py index 970d211a..f1296fe9 100644 --- a/src/flygym/simulation.py +++ b/src/flygym/simulation.py @@ -11,6 +11,7 @@ from flygym.compose.world import BaseWorld from flygym.rendering import Renderer from flygym.utils.profiling import print_perf_report +from flygym.utils.typing import n_jointdofs, n_actuators, n_tendon_actuators class Simulation: @@ -162,7 +163,7 @@ def render_as_needed_with_profile(self) -> bool: self._frames_rendered += 1 return render_done - def get_joint_angles(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]: # noqa: F821 + def get_joint_angles(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]: """Get current joint angles ordered by the fly's skeleton. Args: @@ -175,7 +176,7 @@ def get_joint_angles(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]: internal_ids = self._intern_qposadrs_by_fly[fly_name] return self.mj_data.qpos[internal_ids] - def get_joint_velocities(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]: # noqa: F821 + def get_joint_velocities(self, fly_name: str) -> Float[np.ndarray, "n_jointdofs"]: """Get current joint angular velocities ordered by the fly's skeleton. Args: @@ -216,7 +217,7 @@ def get_body_rotations(self, fly_name: str) -> Float[np.ndarray, "n_bodies 4"]: def get_actuator_forces( self, fly_name: str, actuator_type: ActuatorType - ) -> Float[np.ndarray, "n_actuators"]: # noqa: F821 + ) -> Float[np.ndarray, "n_actuators"]: """Get actuator forces for the given actuator type. Args: @@ -354,7 +355,7 @@ def set_actuator_inputs( self, fly_name: str, actuator_type: ActuatorType, - inputs: Float[np.ndarray, "n_actuators"], # noqa: F821 + inputs: Float[np.ndarray, "n_actuators"], ) -> None: """Set control inputs for the given actuator type. @@ -393,7 +394,7 @@ def set_leg_adhesion_states( def set_tendon_actuator_inputs( self, fly_name: str, - inputs: Float[np.ndarray, "n_tendon_actuators"], # noqa: F821 + inputs: Float[np.ndarray, "n_tendon_actuators"], ) -> None: """Set control inputs for tendon actuators. @@ -770,13 +771,13 @@ def close(self): self.eye_renderer = None # Don't destruct self.retina and self.eye_renderer_scene_option: they can be # reused and retina init requires some IO ops. - + @property def fly(self) -> BaseFly: """Return the single fly in the world, or raise an error if there are multiple.""" return self.world.fly - + @property def fly_lookup(self) -> dict[str, BaseFly]: """Return the fly lookup dictionary from the world.""" - return self.world.fly_lookup \ No newline at end of file + return self.world.fly_lookup diff --git a/src/flygym/utils/typing.py b/src/flygym/utils/typing.py new file mode 100644 index 00000000..e931b968 --- /dev/null +++ b/src/flygym/utils/typing.py @@ -0,0 +1,19 @@ +"""Canonical axis-name bindings for jaxtyping shape hints. + +Referencing these (instead of bare string identifiers) lets pyflakes/ruff +resolve the forward references inside shape strings like +``Float[np.ndarray, "n_bodies"]`` instead of flagging them as undefined +names (F821) — no lint ignores needed. +""" + +from typing import TypeVar + +n_worlds = TypeVar("n_worlds") +n_jointdofs = TypeVar("n_jointdofs") +n_actuators = TypeVar("n_actuators") +n_tendon_actuators = TypeVar("n_tendon_actuators") +n_bodies = TypeVar("n_bodies") +n_sites = TypeVar("n_sites") +n_bodysegments = TypeVar("n_bodysegments") +n_cameras = TypeVar("n_cameras") +n_ommatidia = TypeVar("n_ommatidia") diff --git a/src/flygym/warp/utils.py b/src/flygym/warp/utils.py index 628c51d9..80c409d0 100644 --- a/src/flygym/warp/utils.py +++ b/src/flygym/warp/utils.py @@ -7,9 +7,7 @@ @wp.kernel def wp_gather_indexed_rows_3d( - src: wp.array3d[float], - dst: wp.array3d[float], - rows: wp.array[int], + src: wp.array3d[float], dst: wp.array3d[float], rows: wp.array[int] ): """Gather specific rows (dim 1) from a 3D Warp array into a narrower destination. @@ -30,9 +28,7 @@ def wp_gather_indexed_rows_3d( @wp.kernel def wp_gather_indexed_rows_vec3f( - src: wp.array2d[wp.vec3], - dst: wp.array3d[float], - rows: wp.array[int], + src: wp.array2d[wp.vec3], dst: wp.array3d[float], rows: wp.array[int] ): """Gather specific rows from a 2D ``vec3f`` array into a ``(n_worlds, n_rows_narrow, 3)`` ``float32`` destination. @@ -57,9 +53,7 @@ def wp_gather_indexed_rows_vec3f( @wp.kernel def wp_gather_indexed_rows_quatf( - src: wp.array2d[wp.quat], - dst: wp.array3d[float], - rows: wp.array[int], + src: wp.array2d[wp.quat], dst: wp.array3d[float], rows: wp.array[int] ): """Gather specific rows from a 2D ``quatf`` array into a ``(n_worlds, n_rows_narrow, 4)`` ``float32`` destination. @@ -85,9 +79,7 @@ def wp_gather_indexed_rows_quatf( @wp.kernel def wp_scatter_indexed_cols_2d( - src: wp.array2d[float], - dst: wp.array2d[float], - cols: wp.array[int], + src: wp.array2d[float], dst: wp.array2d[float], cols: wp.array[int] ): """Scatter a 2D Warp array into specific columns of a wider destination array. @@ -108,9 +100,7 @@ def wp_scatter_indexed_cols_2d( @wp.kernel def wp_gather_indexed_cols_2d( - src: wp.array2d[float], - dst: wp.array2d[float], - cols: wp.array[int], + src: wp.array2d[float], dst: wp.array2d[float], cols: wp.array[int] ): """Gather specific columns from a 2D Warp array into a narrower destination array. diff --git a/src/flygym_demo/benchmark/time_gpu_simulation.py b/src/flygym_demo/benchmark/time_gpu_simulation.py index 57d309b0..a601082e 100644 --- a/src/flygym_demo/benchmark/time_gpu_simulation.py +++ b/src/flygym_demo/benchmark/time_gpu_simulation.py @@ -99,9 +99,7 @@ def update_target_angles_kernel( @wp.kernel -def increment_counter_kernel( - step_counter_gpu: wp.array[int], -): +def increment_counter_kernel(step_counter_gpu: wp.array[int]): step_counter_gpu[0] = step_counter_gpu[0] + 1