fix: make numpy.typing import explicit

mjinx/components/tasks/_com_task.py

@@ -137,15 +137,15 @@ def target_com(self) -> jnp.ndarray:
         return self._target_com
 
     @target_com.setter
-    def target_com(self, value: Sequence):
+    def target_com(self, value: jnp.ndarray | Sequence):
         """
         Set the target center of mass for the task.
 
         :param value: The new target center of mass as a sequence of values.
         """
         self.update_target_com(value)
 
-    def update_target_com(self, target_com: Sequence):
+    def update_target_com(self, target_com: jnp.ndarray | Sequence):
         """
         Update the target center of mass for the task.
 

mjinx/components/tasks/_joint_task.py

@@ -177,15 +177,15 @@ def target_q(self) -> jnp.ndarray:
         return self._target_q
 
     @target_q.setter
-    def target_q(self, value: Sequence):
+    def target_q(self, value: jnp.ndarray | Sequence):
         """
         Set the target joint positions for the task.
 
         :param value: The new target joint positions as a sequence of values.
         """
         self.update_target_q(value)
 
-    def update_target_q(self, target_q: Sequence):
+    def update_target_q(self, target_q: jnp.ndarray | Sequence):
         """
         Update the target joint positions for the task.
 

mjinx/typing.py

@@ -11,6 +11,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import numpy.typing
 from mujoco.mjx._src.dataclasses import PyTreeNode
 
 ndarray: TypeAlias = np.ndarray | jnp.ndarray

pyproject.toml

@@ -57,4 +57,4 @@ line-length = 120
 "__init__.py" = ["F401"]
 
 [tool.pytest]
-filterwarnings = "ignore:.*U.*mode is deprecated:DeprecationWarning"
+filterwarnings = ["ignore:.*U.*mode is deprecated:DeprecationWarning"]

tests/unit_tests/configuration/test_configuration.py

@@ -305,8 +305,9 @@ def test_get_distance_with_hfield(self) -> None:
        </mujoco>
        """
         )
-        hfield_data = mjx.make_data(hfield_model)
-        hfield_data = mjx.kinematics(hfield_model, hfield_data)
+        hfield_model_mjx = mjx.put_model(hfield_model)
+        hfield_data = mjx.make_data(hfield_model_mjx)
+        hfield_data = mjx.kinematics(hfield_model_mjx, hfield_data)
 
         collision_pairs: list[CollisionPair] = [(0, 1)]  # hfield and sphere