jaxsim 0.4.3.dev12__py3-none-any.whl → 0.4.3.dev18__py3-none-any.whl

jaxsim/integrators/common.py

@@ -1,6 +1,6 @@
 import abc
 import dataclasses
-from typing import Any, ClassVar, Generic, Protocol, Type, TypeVar
+from typing import Any, ClassVar, Generic, Protocol, TypeVar
 
 import jax
 import jax.numpy as jnp
@@ -64,7 +64,7 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
 
     @classmethod
     def build(
-        cls: Type[Self],
+        cls: type[Self],
         *,
         dynamics: SystemDynamics[State, StateDerivative],
         **kwargs,
@@ -109,11 +109,14 @@ class Integrator(JaxsimDataclass, abc.ABC, Generic[State, StateDerivative]):
             integrator.params = params
 
         with integrator.mutable_context(mutability=Mutability.MUTABLE):
-            xf = integrator(x0, t0, dt, **kwargs)
+            xf, aux_dict = integrator(x0, t0, dt, **kwargs)
 
-        return xf, integrator.params | {
-            Integrator.AfterInitKey: jnp.array(False).astype(bool)
-        }
+        return (
+            xf,
+            integrator.params
+            | {Integrator.AfterInitKey: jnp.array(False).astype(bool)}
+            | aux_dict,
+        )
 
     @abc.abstractmethod
     def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
@@ -224,7 +227,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
     @override
     @classmethod
     def build(
-        cls: Type[Self],
+        cls: type[Self],
         *,
         dynamics: SystemDynamics[State, StateDerivative],
         fsal_enabled_if_supported: jtp.BoolLike = True,
@@ -277,15 +280,19 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         return integrator
 
-    def __call__(self, x0: State, t0: Time, dt: TimeStep, **kwargs) -> NextState:
+    def __call__(
+        self, x0: State, t0: Time, dt: TimeStep, **kwargs
+    ) -> tuple[NextState, dict[str, Any]]:
 
         # Here z is a batched state with as many batch elements as b.T rows.
         # Note that z has multiple batches only if b.T has more than one row,
         # e.g. in Butcher tableau of embedded schemes.
-        z = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
+        z, aux_dict = self._compute_next_state(x0=x0, t0=t0, dt=dt, **kwargs)
 
         # The next state is the batch element located at the configured index of solution.
-        return jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
+        next_state = jax.tree_util.tree_map(lambda l: l[self.row_index_of_solution], z)
+
+        return next_state, aux_dict
 
     @classmethod
     def integrate_rk_stage(
@@ -343,7 +350,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
     def _compute_next_state(
         self, x0: State, t0: Time, dt: TimeStep, **kwargs
-    ) -> NextState:
+    ) -> tuple[NextState, dict[str, Any]]:
         """
         Compute the next state of the system, returning all the output states.
 
@@ -373,19 +380,21 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         )
 
         # Apply FSAL property by passing ẋ0 = f(x0, t0) from the previous iteration.
-        get_ẋ0 = lambda: self.params.get("dxdt0", f(x0, t0)[0])
+        get_ẋ0_and_aux_dict = lambda: self.params.get("dxdt0", f(x0, t0))
 
         # We use a `jax.lax.scan` to compile the `f` function only once.
         # Otherwise, if we compute e.g. for RK4 sequentially, the jit-compiled code
         # would include 4 repetitions of the `f` logic, making everything extremely slow.
-        def scan_body(carry: jax.Array, i: int | jax.Array) -> tuple[jax.Array, None]:
+        def scan_body(
+            carry: jax.Array, i: int | jax.Array
+        ) -> tuple[jax.Array, dict[str, Any]]:
             """"""
 
             # Unpack the carry, i.e. the stacked kᵢ vectors.
             K = carry
 
             # Define the computation of the Runge-Kutta stage.
-            def compute_ki() -> jax.Array:
+            def compute_ki() -> tuple[jax.Array, dict[str, Any]]:
 
                 # Compute ∑ⱼ aᵢⱼ kⱼ.
                 op_sum_ak = lambda k: jnp.einsum("s,s...->...", A[i], k)
@@ -398,13 +407,13 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
                 # Compute the next time for the kᵢ evaluation.
                 ti = t0 + c[i] * Δt
 
-                # This is kᵢ = f(xᵢ, tᵢ).
-                return f(xi, ti)[0]
+                # This is kᵢ, aux_dict = f(xᵢ, tᵢ).
+                return f(xi, ti)
 
             # This selector enables FSAL property in the first iteration (i=0).
-            ki = jax.lax.cond(
+            ki, aux_dict = jax.lax.cond(
                 pred=jnp.logical_and(i == 0, self.has_fsal),
-                true_fun=get_ẋ0,
+                true_fun=get_ẋ0_and_aux_dict,
                 false_fun=compute_ki,
             )
 
@@ -413,10 +422,10 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             K = jax.tree_util.tree_map(op, K, ki)
 
             carry = K
-            return carry, None
+            return carry, aux_dict
 
         # Compute the state derivatives kᵢ.
-        K, _ = jax.lax.scan(
+        K, aux_dict = jax.lax.scan(
             f=scan_body,
             init=carry0,
             xs=jnp.arange(c.size),
@@ -439,7 +448,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
             lambda xf: self.post_process_state(x0=x0, t0=t0, xf=xf, dt=dt)
         )(z)
 
-        return z_transformed
+        return z_transformed, aux_dict
 
     @staticmethod
     def butcher_tableau_is_valid(

jaxsim/integrators/variable_step.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Any, ClassVar, Generic, Type
+from typing import Any, ClassVar, Generic
 
 try:
     from typing import Self
@@ -495,7 +495,7 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
 
     @classmethod
     def build(
-        cls: Type[Self],
+        cls: type[Self],
         *,
         dynamics: SystemDynamics[State, StateDerivative],
         fsal_enabled_if_supported: jtp.BoolLike = True,

jaxsim/logging.py

@@ -1,6 +1,5 @@
 import enum
 import logging
-from typing import Union
 
 import coloredlogs
 
@@ -20,7 +19,7 @@ def _logger() -> logging.Logger:
     return logging.getLogger(name=LOGGER_NAME)
 
 
-def set_logging_level(level: Union[int, LoggingLevel] = LoggingLevel.WARNING):
+def set_logging_level(level: int | LoggingLevel = LoggingLevel.WARNING):
     if isinstance(level, int):
         level = LoggingLevel(level)
 

jaxsim/math/inertia.py

@@ -1,5 +1,3 @@
-from typing import Tuple
-
 import jax.numpy as jnp
 
 import jaxsim.typing as jtp
@@ -39,7 +37,7 @@ class Inertia:
         return M
 
     @staticmethod
-    def to_params(M: jtp.Matrix) -> Tuple[jtp.Float, jtp.Vector, jtp.Matrix]:
+    def to_params(M: jtp.Matrix) -> tuple[jtp.Float, jtp.Vector, jtp.Matrix]:
         """
         Convert a 6x6 inertia matrix to mass, center of mass, and inertia matrix.
 

jaxsim/math/joint_model.py

@@ -107,7 +107,7 @@ class JointModel:
             λ_H_pre=λ_H_pre,
             suc_H_i=suc_H_i,
             # Static attributes
-            joint_dofs=tuple([base_dofs] + [int(1) for _ in ordered_joints]),
+            joint_dofs=tuple([base_dofs] + [1 for _ in ordered_joints]),
             joint_names=tuple(["world_to_base"] + [j.name for j in ordered_joints]),
             joint_types=tuple([JointType.Fixed] + [j.jtype for j in ordered_joints]),
             joint_axis=tuple(JointGenericAxis(axis=j.axis) for j in ordered_joints),

jaxsim/math/rotation.py

@@ -1,5 +1,3 @@
-from typing import Tuple
-
 import jax
 import jax.numpy as jnp
 import jaxlie
@@ -64,7 +62,7 @@ class Rotation:
         vector = vector.squeeze()
         theta = jnp.linalg.norm(vector)
 
-        def theta_is_not_zero(theta_and_v: Tuple[jtp.Float, jtp.Vector]) -> jtp.Matrix:
+        def theta_is_not_zero(theta_and_v: tuple[jtp.Float, jtp.Vector]) -> jtp.Matrix:
             theta, v = theta_and_v
 
             s = jnp.sin(theta)

jaxsim/mujoco/loaders.py

@@ -4,7 +4,8 @@ import dataclasses
 import pathlib
 import tempfile
 import warnings
-from typing import Any, Sequence
+from collections.abc import Sequence
+from typing import Any
 
 import mujoco as mj
 import numpy as np

jaxsim/mujoco/model.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import functools
 import pathlib
-from typing import Any, Callable
+from collections.abc import Callable
+from typing import Any
 
 import mujoco as mj
 import numpy as np

jaxsim/mujoco/visualizer.py

@@ -1,6 +1,6 @@
 import contextlib
 import pathlib
-from typing import ContextManager, Sequence
+from collections.abc import Sequence
 
 import mediapy as media
 import mujoco as mj
@@ -172,7 +172,7 @@ class MujocoVisualizer:
         distance: float | int | npt.NDArray | None = None,
         azimut: float | int | npt.NDArray | None = None,
         elevation: float | int | npt.NDArray | None = None,
-    ) -> ContextManager[mujoco.viewer.Handle]:
+    ) -> contextlib.AbstractContextManager[mujoco.viewer.Handle]:
         """
         Context manager to open the Mujoco passive viewer.
 

jaxsim/parsers/descriptions/model.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import dataclasses
 import itertools
-from typing import Sequence
+from collections.abc import Sequence
 
 from jaxsim import logging
 

jaxsim/parsers/kinematic_graph.py

@@ -3,7 +3,8 @@ from __future__ import annotations
 import copy
 import dataclasses
 import functools
-from typing import Any, Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
+from typing import Any
 
 import numpy as np
 import numpy.typing as npt
@@ -444,7 +445,7 @@ class KinematicGraph(Sequence[LinkDescription]):
                 msg.format(
                     link_to_remove.name,
                     self.joints_connection_dict[
-                        (parent_of_link_to_remove.name, link_to_remove.name)
+                        parent_of_link_to_remove.name, link_to_remove.name
                     ].name,
                     parent_of_link_to_remove.name,
                 )
@@ -852,7 +853,7 @@ class KinematicGraphTransforms:
 
             # Get the joint between the link and its parent.
             parent_joint = self.graph.joints_connection_dict[
-                (link.parent.name, link.name)
+                link.parent.name, link.name
             ]
 
             # Get the transform of the parent joint.

jaxsim/parsers/rod/parser.py

@@ -1,6 +1,6 @@
 import dataclasses
 import pathlib
-from typing import Dict, List, NamedTuple, Optional, Union
+from typing import NamedTuple
 
 import jax.numpy as jnp
 import numpy as np
@@ -23,17 +23,17 @@ class SDFData(NamedTuple):
     fixed_base: bool
     base_link_name: str
 
-    link_descriptions: List[descriptions.LinkDescription]
-    joint_descriptions: List[descriptions.JointDescription]
-    frame_descriptions: List[descriptions.LinkDescription]
-    collision_shapes: List[descriptions.CollisionShape]
+    link_descriptions: list[descriptions.LinkDescription]
+    joint_descriptions: list[descriptions.JointDescription]
+    frame_descriptions: list[descriptions.LinkDescription]
+    collision_shapes: list[descriptions.CollisionShape]
 
     sdf_model: rod.Model | None = None
     model_pose: kinematic_graph.RootPose = kinematic_graph.RootPose()
 
 
 def extract_model_data(
-    model_description: Union[pathlib.Path, str, rod.Model],
+    model_description: pathlib.Path | str | rod.Model,
     model_name: str | None = None,
     is_urdf: bool | None = None,
 ) -> SDFData:
@@ -114,7 +114,7 @@ def extract_model_data(
     ]
 
     # Create a dictionary to find easily links.
-    links_dict: Dict[str, descriptions.LinkDescription] = {l.name: l for l in links}
+    links_dict: dict[str, descriptions.LinkDescription] = {l.name: l for l in links}
 
     # ============
     # Parse frames
@@ -304,7 +304,7 @@ def extract_model_data(
     # ================
 
     # Initialize the collision shapes
-    collisions: List[descriptions.CollisionShape] = []
+    collisions: list[descriptions.CollisionShape] = []
 
     # Parse the collisions
     for link in sdf_model.links():
@@ -339,8 +339,8 @@ def extract_model_data(
 
 
 def build_model_description(
-    model_description: Union[pathlib.Path, str, rod.Model],
-    is_urdf: Optional[bool] = False,
+    model_description: pathlib.Path | str | rod.Model,
+    is_urdf: bool | None = False,
 ) -> descriptions.ModelDescription:
     """
     Builds a model description from an SDF/URDF resource.

jaxsim/rbda/contacts/common.py

@@ -5,6 +5,7 @@ from typing import Any
 
 import jaxsim.terrain
 import jaxsim.typing as jtp
+from jaxsim.utils import JaxsimDataclass
 
 
 class ContactsState(abc.ABC):
@@ -42,7 +43,7 @@ class ContactsState(abc.ABC):
         pass
 
 
-class ContactsParams(abc.ABC):
+class ContactsParams(JaxsimDataclass):
     """
     Abstract class representing the parameters of a contact model.
     """
@@ -67,7 +68,7 @@ class ContactsParams(abc.ABC):
         pass
 
 
-class ContactModel(abc.ABC):
+class ContactModel(JaxsimDataclass):
     """
     Abstract class representing a contact model.