Version v2.1.1

docs/changelog.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Version 2.1.1 (under development)
+
+TODO
+
 ## Version 2.1.0
 
 !!! danger "API-breaking changes"

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 = [
@@ -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

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 <name>.tar + <name>.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 <path/to/asset_dir>
+set -e
+
+if [ $# -ne 1 ]; then
+    echo "Usage: $0 <path/to/asset_dir>" >&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"))"

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.
 

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

src/flygym/compose/world/base_world.py

@@ -82,11 +82,25 @@ 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(
         self,
         fly: BaseFly,
-        spawn_position: Vec3,
+        spawn_position: Vec3 | tuple[float, float, float],
         spawn_rotation: Rotation3D,
         *args,
         **kwargs,
@@ -125,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,
@@ -246,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: (

src/flygym/simulation.py

@@ -7,10 +7,11 @@
 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
+from flygym.utils.typing import n_jointdofs, n_actuators, n_tendon_actuators
 
 
 class Simulation:
@@ -68,7 +69,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)
 
@@ -152,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:
@@ -165,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:
@@ -206,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:
@@ -344,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.
 
@@ -383,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.
 
@@ -760,3 +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