jaxsim 0.4.3.dev327__py3-none-any.whl → 0.4.3.dev350__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev327'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev327')
+__version__ = version = '0.4.3.dev350'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev350')

jaxsim/api/joint.py

@@ -53,9 +53,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str
     """
 
     exceptions.raise_value_error_if(
-        condition=jnp.array(
-            [joint_index < 0, joint_index >= model.number_of_joints()]
-        ).any(),
+        condition=joint_index < 0,
         msg="Invalid joint index '{idx}'",
         idx=joint_index,
     )
@@ -123,10 +121,7 @@ def position_limit(
     """
 
     if model.number_of_joints() == 0:
-        s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min
-        s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max
-
-        return jnp.atleast_1d(s_min).astype(float), jnp.atleast_1d(s_max).astype(float)
+        return jnp.empty(0).astype(float), jnp.empty(0).astype(float)
 
     exceptions.raise_value_error_if(
         condition=jnp.array(
@@ -136,8 +131,12 @@ def position_limit(
         idx=joint_index,
     )
 
-    s_min = model.kin_dyn_parameters.joint_parameters.position_limits_min[joint_index]
-    s_max = model.kin_dyn_parameters.joint_parameters.position_limits_max[joint_index]
+    s_min = jnp.atleast_1d(
+        model.kin_dyn_parameters.joint_parameters.position_limits_min
+    )[joint_index]
+    s_max = jnp.atleast_1d(
+        model.kin_dyn_parameters.joint_parameters.position_limits_max
+    )[joint_index]
 
     return s_min.astype(float), s_max.astype(float)
 

jaxsim/api/kin_dyn_parameters.py

@@ -438,7 +438,9 @@ class KynDynParameters(JaxsimDataclass):
     # Helpers to update parameters
     # ============================
 
-    def set_link_mass(self, link_index: int, mass: jtp.FloatLike) -> KynDynParameters:
+    def set_link_mass(
+        self, link_index: jtp.IntLike, mass: jtp.FloatLike
+    ) -> KynDynParameters:
         """
         Set the mass of a link.
 
@@ -457,7 +459,7 @@ class KynDynParameters(JaxsimDataclass):
         return self.replace(link_parameters=link_parameters)
 
     def set_link_inertia(
-        self, link_index: int, inertia: jtp.MatrixLike
+        self, link_index: jtp.IntLike, inertia: jtp.MatrixLike
     ) -> KynDynParameters:
         r"""
         Set the inertia tensor of a link.
@@ -593,10 +595,10 @@ class LinkParameters(JaxsimDataclass):
         """
 
         # Extract the link parameters from the 6D spatial inertia.
-        m, L_p_CoM, I = Inertia.to_params(M=M)
+        m, L_p_CoM, I_CoM = Inertia.to_params(M=M)
 
         # Extract only the necessary elements of the inertia tensor.
-        inertia_elements = I[jnp.triu_indices(3)]
+        inertia_elements = I_CoM[jnp.triu_indices(3)]
 
         return LinkParameters(
             index=jnp.array(index).squeeze().astype(int),

jaxsim/api/link.py

@@ -4,6 +4,7 @@ from collections.abc import Sequence
 import jax
 import jax.numpy as jnp
 import jax.scipy.linalg
+import numpy as np
 
 import jaxsim.api as js
 import jaxsim.rbda
@@ -54,9 +55,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, link_index: jtp.IntLike) -> str:
     """
 
     exceptions.raise_value_error_if(
-        condition=jnp.array(
-            [link_index < 0, link_index >= model.number_of_links()]
-        ).any(),
+        condition=link_index < 0,
         msg="Invalid link index '{idx}'",
         idx=link_index,
     )
@@ -98,7 +97,7 @@ def idxs_to_names(
         The names of the links.
     """
 
-    return tuple(idx_to_name(model=model, link_index=idx) for idx in link_indices)
+    return tuple(np.array(model.kin_dyn_parameters.link_names)[list(link_indices)])
 
 
 # =========

jaxsim/api/model.py

@@ -304,7 +304,7 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.model_name
 
-    def number_of_links(self) -> jtp.Int:
+    def number_of_links(self) -> int:
         """
         Return the number of links in the model.
 
@@ -317,7 +317,7 @@ class JaxSimModel(JaxsimDataclass):
 
         return self.kin_dyn_parameters.number_of_links()
 
-    def number_of_joints(self) -> jtp.Int:
+    def number_of_joints(self) -> int:
         """
         Return the number of joints in the model.
 
@@ -419,7 +419,7 @@ class JaxSimModel(JaxsimDataclass):
 def reduce(
     model: JaxSimModel,
     considered_joints: tuple[str, ...],
-    locked_joint_positions: dict[str, jtp.Float] | None = None,
+    locked_joint_positions: dict[str, jtp.FloatLike] | None = None,
 ) -> JaxSimModel:
     """
     Reduce the model by lumping together the links connected by removed joints.
@@ -1038,12 +1038,7 @@ def forward_dynamics_aba(
     C_v̇_WB = to_active(
         W_v̇_WB=W_v̇_WB,
         W_H_C=W_H_C,
-        W_v_WB=jnp.hstack(
-            [
-                data.state.physics_model.base_linear_velocity,
-                data.state.physics_model.base_angular_velocity,
-            ]
-        ),
+        W_v_WB=W_v_WB,
         W_v_WC=W_v_WC,
     )
 
@@ -2274,16 +2269,12 @@ def step(
             # Raise runtime error for not supported case in which Rigid contacts and
             # Baumgarte stabilization are enabled and used with ForwardEuler integrator.
             jaxsim.exceptions.raise_runtime_error_if(
-                condition=jnp.logical_and(
-                    isinstance(
-                        integrator,
-                        jaxsim.integrators.fixed_step.ForwardEuler
-                        | jaxsim.integrators.fixed_step.ForwardEulerSO3,
-                    ),
-                    jnp.array(
-                        [data_tf.contacts_params.K, data_tf.contacts_params.D]
-                    ).any(),
-                ),
+                condition=isinstance(
+                    integrator,
+                    jaxsim.integrators.fixed_step.ForwardEuler
+                    | jaxsim.integrators.fixed_step.ForwardEulerSO3,
+                )
+                & ((data_tf.contacts_params.K > 0) | (data_tf.contacts_params.D > 0)),
                 msg="Baumgarte stabilization is not supported with ForwardEuler integrators",
             )
 

jaxsim/api/references.py

@@ -503,7 +503,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
         ]
 
         exceptions.raise_value_error_if(
-            condition=jnp.logical_not(data.valid(model=model)),
+            condition=~data.valid(model=model),
             msg="The provided data is not valid for the model",
         )
         W_H_Fi = jax.vmap(

jaxsim/integrators/common.py

@@ -319,7 +319,7 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
         # Initialize the carry of the for loop with the stacked kᵢ vectors.
-        carry0 = jax.tree_map(
+        carry0 = jax.tree.map(
             lambda l: jnp.zeros((c.size, *l.shape), dtype=l.dtype), x0
         )
 
@@ -507,7 +507,7 @@ class ExplicitRungeKuttaSO3Mixin:
 
         # We assume that the initial quaternion is already unary.
         exceptions.raise_runtime_error_if(
-            condition=jnp.logical_not(jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0)),
+            condition=~jnp.allclose(W_Q_B_t0.dot(W_Q_B_t0), 1.0),
             msg="The SO(3) integrator received a quaternion at t0 that is not unary.",
         )
 

jaxsim/integrators/variable_step.py

@@ -152,7 +152,7 @@ def compute_pytree_scale(
     """
 
     # Consider a zero second pytree, if not given.
-    x2 = jax.tree.map(lambda l: jnp.zeros_like(l), x1) if x2 is None else x2
+    x2 = jax.tree.map(jnp.zeros_like, x1) if x2 is None else x2
 
     # Compute the scaling factors of the initial state and its derivative.
     compute_scale = lambda l1, l2: atol + jnp.maximum(jnp.abs(l1), jnp.abs(l2)) * rtol
@@ -199,9 +199,7 @@ def local_error_estimation(
 
     # Consider a zero estimated final state, if not given.
     xf_estimate = (
-        jax.tree.map(lambda l: jnp.zeros_like(l), xf)
-        if xf_estimate is None
-        else xf_estimate
+        jax.tree.map(jnp.zeros_like, xf) if xf_estimate is None else xf_estimate
     )
 
     # Estimate the error.
@@ -483,14 +481,10 @@ class EmbeddedRungeKutta(ExplicitRungeKutta[PyTreeType], Generic[PyTreeType]):
                 metadata_next,
                 discarded_steps,
             ) = jax.lax.cond(
-                pred=jnp.array(
-                    [
-                        discarded_steps >= self.max_step_rejections,
-                        local_error <= 1.0,
-                        Δt_next < self.dt_min,
-                        integrator_init,
-                    ]
-                ).any(),
+                pred=discarded_steps
+                >= self.max_step_rejections | local_error
+                <= 1.0 | Δt_next
+                < self.dt_min | integrator_init,
                 true_fun=accept_step,
                 false_fun=reject_step,
             )

jaxsim/mujoco/loaders.py

@@ -1,6 +1,3 @@
-from __future__ import annotations
-
-import dataclasses
 import pathlib
 import tempfile
 import warnings
@@ -9,10 +6,14 @@ from typing import Any
 
 import mujoco as mj
 import numpy as np
-import numpy.typing as npt
 import rod.urdf.exporter
 from lxml import etree as ET
-from scipy.spatial.transform import Rotation
+
+from .utils import MujocoCamera
+
+MujocoCameraType = (
+    MujocoCamera | Sequence[MujocoCamera] | dict[str, str] | Sequence[dict[str, str]]
+)
 
 
 def load_rod_model(
@@ -167,12 +168,7 @@ class RodModelToMjcf:
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
         heightmap_samples_xy: tuple[int, int] = (101, 101),
-        cameras: (
-            MujocoCamera
-            | Sequence[MujocoCamera]
-            | dict[str, str]
-            | Sequence[dict[str, str]]
-        ) = (),
+        cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a ROD model to a Mujoco MJCF string.
@@ -533,12 +529,7 @@ class UrdfToMjcf:
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
-        cameras: (
-            MujocoCamera
-            | Sequence[MujocoCamera]
-            | dict[str, str]
-            | Sequence[dict[str, str]]
-        ) = (),
+        cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a URDF file to a Mujoco MJCF string.
@@ -580,12 +571,7 @@ class SdfToMjcf:
         model_name: str | None = None,
         plane_normal: tuple[float, float, float] = (0, 0, 1),
         heightmap: bool | None = None,
-        cameras: (
-            MujocoCamera
-            | Sequence[MujocoCamera]
-            | dict[str, str]
-            | Sequence[dict[str, str]]
-        ) = (),
+        cameras: MujocoCameraType = (),
     ) -> tuple[str, dict[str, Any]]:
         """
         Converts a SDF file to a Mujoco MJCF string.
@@ -617,118 +603,3 @@ class SdfToMjcf:
             heightmap=heightmap,
             cameras=cameras,
         )
-
-
-@dataclasses.dataclass
-class MujocoCamera:
-    """
-    Helper class storing parameters of a Mujoco camera.
-
-    Refer to the official documentation for more details:
-    https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-camera
-    """
-
-    mode: str = "fixed"
-
-    target: str | None = None
-    fovy: str = "45"
-    pos: str = "0 0 0"
-
-    quat: str | None = None
-    axisangle: str | None = None
-    xyaxes: str | None = None
-    zaxis: str | None = None
-    euler: str | None = None
-
-    name: str | None = None
-
-    @classmethod
-    def build(cls, **kwargs) -> MujocoCamera:
-
-        if not all(isinstance(value, str) for value in kwargs.values()):
-            raise ValueError(f"Values must be strings: {kwargs}")
-
-        return cls(**kwargs)
-
-    @staticmethod
-    def build_from_target_view(
-        camera_name: str,
-        lookat: Sequence[float | int] | npt.NDArray = (0, 0, 0),
-        distance: float | int | npt.NDArray = 3,
-        azimut: float | int | npt.NDArray = 90,
-        elevation: float | int | npt.NDArray = -45,
-        fovy: float | int | npt.NDArray = 45,
-        degrees: bool = True,
-        **kwargs,
-    ) -> MujocoCamera:
-        """
-        Create a custom camera that looks at a target point.
-
-        Note:
-            The choice of the parameters is easier if we imagine to consider a target
-            frame `T` whose origin is located over the lookat point and having the same
-            orientation of the world frame `W`. We also introduce a camera frame `C`
-            whose origin is located over the lower-left corner of the image, and having
-            the x-axis pointing right and the y-axis pointing up in image coordinates.
-            The camera renders what it sees in the -z direction of frame `C`.
-
-        Args:
-            camera_name: The name of the camera.
-            lookat: The target point to look at (origin of `T`).
-            distance:
-                The distance from the target point (displacement between the origins
-                of `T` and `C`).
-            azimut:
-                The rotation around z of the camera. With an angle of 0, the camera
-                would loot at the target point towards the positive x-axis of `T`.
-            elevation:
-                The rotation around the x-axis of the camera frame `C`. Note that if
-                you want to lift the view angle, the elevation is negative.
-            fovy: The field of view of the camera.
-            degrees: Whether the angles are in degrees or radians.
-            **kwargs: Additional camera parameters.
-
-        Returns:
-            The custom camera.
-        """
-
-        # Start from a frame whose origin is located over the lookat point.
-        # We initialize a -90 degrees rotation around the z-axis because due to
-        # the default camera coordinate system (x pointing right, y pointing up).
-        W_H_C = np.eye(4)
-        W_H_C[0:3, 3] = np.array(lookat)
-        W_H_C[0:3, 0:3] = Rotation.from_euler(
-            seq="ZX", angles=[-90, 90], degrees=True
-        ).as_matrix()
-
-        # Process the azimut.
-        R_az = Rotation.from_euler(seq="Y", angles=azimut, degrees=degrees).as_matrix()
-        W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_az
-
-        # Process elevation.
-        R_el = Rotation.from_euler(
-            seq="X", angles=elevation, degrees=degrees
-        ).as_matrix()
-        W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_el
-
-        # Process distance.
-        tf_distance = np.eye(4)
-        tf_distance[2, 3] = distance
-        W_H_C = W_H_C @ tf_distance
-
-        # Extract the position and the quaternion.
-        p = W_H_C[0:3, 3]
-        Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat(scalar_first=True)
-
-        return MujocoCamera.build(
-            name=camera_name,
-            mode="fixed",
-            fovy=f"{fovy if degrees else np.rad2deg(fovy)}",
-            pos=" ".join(p.astype(str).tolist()),
-            quat=" ".join(Q.astype(str).tolist()),
-            **kwargs,
-        )
-
-    def asdict(self) -> dict[str, str]:
-
-        return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}

jaxsim/mujoco/utils.py

@@ -1,7 +1,14 @@
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Sequence
+
 import mujoco as mj
 import numpy as np
+import numpy.typing as npt
+from scipy.spatial.transform import Rotation
 
-from . import MujocoModelHelper
+from .model import MujocoModelHelper
 
 
 def mujoco_data_from_jaxsim(
@@ -99,3 +106,118 @@ def mujoco_data_from_jaxsim(
     mj.mj_forward(mujoco_model, model_helper.data)
 
     return model_helper.data
+
+
+@dataclasses.dataclass
+class MujocoCamera:
+    """
+    Helper class storing parameters of a Mujoco camera.
+
+    Refer to the official documentation for more details:
+    https://mujoco.readthedocs.io/en/stable/XMLreference.html#body-camera
+    """
+
+    mode: str = "fixed"
+
+    target: str | None = None
+    fovy: str = "45"
+    pos: str = "0 0 0"
+
+    quat: str | None = None
+    axisangle: str | None = None
+    xyaxes: str | None = None
+    zaxis: str | None = None
+    euler: str | None = None
+
+    name: str | None = None
+
+    @classmethod
+    def build(cls, **kwargs) -> MujocoCamera:
+
+        if not all(isinstance(value, str) for value in kwargs.values()):
+            raise ValueError(f"Values must be strings: {kwargs}")
+
+        return cls(**kwargs)
+
+    @staticmethod
+    def build_from_target_view(
+        camera_name: str,
+        lookat: Sequence[float | int] | npt.NDArray = (0, 0, 0),
+        distance: float | int | npt.NDArray = 3,
+        azimut: float | int | npt.NDArray = 90,
+        elevation: float | int | npt.NDArray = -45,
+        fovy: float | int | npt.NDArray = 45,
+        degrees: bool = True,
+        **kwargs,
+    ) -> MujocoCamera:
+        """
+        Create a custom camera that looks at a target point.
+
+        Note:
+            The choice of the parameters is easier if we imagine to consider a target
+            frame `T` whose origin is located over the lookat point and having the same
+            orientation of the world frame `W`. We also introduce a camera frame `C`
+            whose origin is located over the lower-left corner of the image, and having
+            the x-axis pointing right and the y-axis pointing up in image coordinates.
+            The camera renders what it sees in the -z direction of frame `C`.
+
+        Args:
+            camera_name: The name of the camera.
+            lookat: The target point to look at (origin of `T`).
+            distance:
+                The distance from the target point (displacement between the origins
+                of `T` and `C`).
+            azimut:
+                The rotation around z of the camera. With an angle of 0, the camera
+                would loot at the target point towards the positive x-axis of `T`.
+            elevation:
+                The rotation around the x-axis of the camera frame `C`. Note that if
+                you want to lift the view angle, the elevation is negative.
+            fovy: The field of view of the camera.
+            degrees: Whether the angles are in degrees or radians.
+            **kwargs: Additional camera parameters.
+
+        Returns:
+            The custom camera.
+        """
+
+        # Start from a frame whose origin is located over the lookat point.
+        # We initialize a -90 degrees rotation around the z-axis because due to
+        # the default camera coordinate system (x pointing right, y pointing up).
+        W_H_C = np.eye(4)
+        W_H_C[0:3, 3] = np.array(lookat)
+        W_H_C[0:3, 0:3] = Rotation.from_euler(
+            seq="ZX", angles=[-90, 90], degrees=True
+        ).as_matrix()
+
+        # Process the azimut.
+        R_az = Rotation.from_euler(seq="Y", angles=azimut, degrees=degrees).as_matrix()
+        W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_az
+
+        # Process elevation.
+        R_el = Rotation.from_euler(
+            seq="X", angles=elevation, degrees=degrees
+        ).as_matrix()
+        W_H_C[0:3, 0:3] = W_H_C[0:3, 0:3] @ R_el
+
+        # Process distance.
+        tf_distance = np.eye(4)
+        tf_distance[2, 3] = distance
+        W_H_C = W_H_C @ tf_distance
+
+        # Extract the position and the quaternion.
+        p = W_H_C[0:3, 3]
+        Q = Rotation.from_matrix(W_H_C[0:3, 0:3]).as_quat(scalar_first=True)
+
+        return MujocoCamera.build(
+            name=camera_name,
+            mode="fixed",
+            fovy=f"{fovy if degrees else np.rad2deg(fovy)}",
+            pos=" ".join(p.astype(str).tolist()),
+            quat=" ".join(Q.astype(str).tolist()),
+            **kwargs,
+        )
+
+    def asdict(self) -> dict[str, str]:
+
+        return {k: v for k, v in dataclasses.asdict(self).items() if v is not None}

jaxsim/parsers/descriptions/joint.py

@@ -100,32 +100,7 @@ class JointDescription(JaxsimDataclass):
         if not isinstance(other, JointDescription):
             return False
 
-        if not (
-            self.name == other.name
-            and self.jtype == other.jtype
-            and self.child == other.child
-            and self.parent == other.parent
-            and self.index == other.index
-            and all(
-                np.allclose(getattr(self, attr), getattr(other, attr))
-                for attr in [
-                    "axis",
-                    "pose",
-                    "friction_static",
-                    "friction_viscous",
-                    "position_limit_damper",
-                    "position_limit_spring",
-                    "position_limit",
-                    "initial_position",
-                    "motor_inertia",
-                    "motor_viscous_friction",
-                    "motor_gear_ratio",
-                ]
-            ),
-        ):
-            return False
-
-        return True
+        return hash(self) == hash(other)
 
     def __hash__(self) -> int:
 

jaxsim/parsers/kinematic_graph.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import copy
 import dataclasses
 import functools
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Iterable, Iterator, Sequence
 from typing import Any
 
 import numpy as np
@@ -82,7 +82,7 @@ class KinematicGraph(Sequence[LinkDescription]):
         default_factory=list, hash=False, compare=False
     )
 
-    root_pose: RootPose = dataclasses.field(default_factory=lambda: RootPose())
+    root_pose: RootPose = dataclasses.field(default_factory=RootPose)
 
     # Private attribute storing optional additional info.
     _extra_info: dict[str, Any] = dataclasses.field(
@@ -700,7 +700,7 @@ class KinematicGraph(Sequence[LinkDescription]):
     # Sequence protocol
     # =================
 
-    def __iter__(self) -> Iterable[LinkDescription]:
+    def __iter__(self) -> Iterator[LinkDescription]:
         yield from KinematicGraph.breadth_first_search(root=self.root)
 
     def __reversed__(self) -> Iterable[LinkDescription]:

jaxsim/parsers/rod/parser.py

@@ -85,10 +85,7 @@ def extract_model_data(
 
     # Log type of base link.
     logging.debug(
-        msg="Model '{}' is {}".format(
-            sdf_model.name,
-            "fixed-base" if sdf_model.is_fixed_base() else "floating-base",
-        )
+        msg=f"Model '{sdf_model.name}' is {'fixed-base' if sdf_model.is_fixed_base() else 'floating-base'}"
     )
 
     # Log detected base link.
@@ -175,7 +172,7 @@ def extract_model_data(
             for j in sdf_model.joints()
             if j.type == "fixed"
             and j.parent == "world"
-            and j.child in links_dict.keys()
+            and j.child in links_dict
             and j.pose.relative_to in {"__model__", "world", None}
         ]
 
@@ -287,7 +284,7 @@ def extract_model_data(
         for j in sdf_model.joints()
         if j.type in {"revolute", "continuous", "prismatic", "fixed"}
         and j.parent != "world"
-        and j.child in links_dict.keys()
+        and j.child in links_dict
     ]
 
     # Create a dictionary to find the parent joint of the links.

jaxsim/parsers/rod/utils.py

@@ -179,7 +179,7 @@ def create_sphere_collision(
 
     r = collision.geometry.sphere.radius
     sphere_points = r * fibonacci_sphere(
-        samples=int(os.getenv(key="JAXSIM_COLLISION_SPHERE_POINTS", default="250"))
+        samples=int(os.getenv(key="JAXSIM_COLLISION_SPHERE_POINTS", default="50"))
     )
 
     H = collision.pose.transform() if collision.pose is not None else np.eye(4)

jaxsim/rbda/jacobian.py

@@ -205,7 +205,7 @@ def jacobian_full_doubly_left(
     # Convert adjoints to SE(3) transforms.
     # Returning them here prevents calling FK in case the output representation
     # of the Jacobian needs to be changed.
-    B_H_L = jax.vmap(lambda B_X_L: Adjoint.to_transform(B_X_L))(B_X_i)
+    B_H_L = jax.vmap(Adjoint.to_transform)(B_X_i)
 
     # Adjust shape of doubly-left free-floating full Jacobian.
     B_J_full_WL_B = J.squeeze().astype(float)
@@ -322,7 +322,7 @@ def jacobian_derivative_full_doubly_left(
     # Convert adjoints to SE(3) transforms.
     # Returning them here prevents calling FK in case the output representation
     # of the Jacobian needs to be changed.
-    B_H_L = jax.vmap(lambda B_X_L: Adjoint.to_transform(B_X_L))(B_X_i)
+    B_H_L = jax.vmap(Adjoint.to_transform)(B_X_i)
 
     # Adjust shape of doubly-left free-floating full Jacobian derivative.
     B_J̇_full_WL_B = J̇.squeeze().astype(float)

jaxsim/rbda/utils.py

@@ -135,7 +135,7 @@ def process_inputs(
     # Check that the quaternion is unary since our RBDAs make this assumption in order
     # to prevent introducing additional normalizations that would affect AD.
     exceptions.raise_value_error_if(
-        condition=jnp.logical_not(jnp.allclose(W_Q_B.dot(W_Q_B), 1.0)),
+        condition=~jnp.allclose(W_Q_B.dot(W_Q_B), 1.0),
         msg="A RBDA received a quaternion that is not normalized.",
     )
 

jaxsim/terrain/terrain.py

@@ -8,6 +8,7 @@ import jax_dataclasses
 import numpy as np
 
 import jaxsim.typing as jtp
+from jaxsim import exceptions
 
 
 class Terrain(abc.ABC):
@@ -108,7 +109,9 @@ class PlaneTerrain(FlatTerrain):
             _normal=tuple(normal.tolist()),
         )
 
-    def normal(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Vector:
+    def normal(
+        self, x: jtp.FloatLike | None = None, y: jtp.FloatLike | None = None
+    ) -> jtp.Vector:
         """
         Compute the normal vector of the terrain at a specific (x, y) location.
 
@@ -141,6 +144,11 @@ class PlaneTerrain(FlatTerrain):
         # Get the plane equation coefficients from the terrain normal.
         A, B, C = self._normal
 
+        exceptions.raise_value_error_if(
+            condition=jnp.allclose(C, 0.0),
+            msg="The z component of the normal cannot be zero.",
+        )
+
         # Compute the final coefficient D considering the terrain height.
         D = -C * self._height
 

jaxsim/utils/tracing.py

@@ -8,15 +8,9 @@ import jax.interpreters.partial_eval
 def tracing(var: Any) -> bool | jax.Array:
     """Returns True if the variable is being traced by JAX, False otherwise."""
 
-    return jax.numpy.array(
-        [
-            isinstance(var, t)
-            for t in (
-                jax._src.core.Tracer,
-                jax.interpreters.partial_eval.DynamicJaxprTracer,
-            )
-        ]
-    ).any()
+    return isinstance(
+        var, jax._src.core.Tracer | jax.interpreters.partial_eval.DynamicJaxprTracer
+    )
 
 
 def not_tracing(var: Any) -> bool | jax.Array:

jaxsim/utils/wrappers.py

@@ -49,7 +49,7 @@ class CustomHashedObject(Generic[T]):
 
     obj: T
 
-    hash_function: Callable[[T], int] = dataclasses.field(default=lambda obj: hash(obj))
+    hash_function: Callable[[T], int] = hash
 
     def get(self: CustomHashedObject[T]) -> T:
         return self.obj

{jaxsim-0.4.3.dev327.dist-info → jaxsim-0.4.3.dev350.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev327
+Version: 0.4.3.dev350
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.4.3.dev327.dist-info → jaxsim-0.4.3.dev350.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=crnjdWEkz6VH0LOon-LPCudsgKRsebWeeDMlTRY3AB0,428
+jaxsim/_version.py,sha256=56tTuqXBlX9UQVJyJ_A9hRmvozgLzRGJ-9ZCppehae8,428
 jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -9,17 +9,17 @@ jaxsim/api/common.py,sha256=SNgxq42r6eF_-aPszvOjUYkGwXOzz4hKmhDwEUkscFQ,6650
 jaxsim/api/contact.py,sha256=D6RucrH9gnoUFLdmAEYwLGrimU0wLmuoDeOONu4ni74,25658
 jaxsim/api/data.py,sha256=ThRpoBlbdwf1N3xs8SWrY5d8RbfdYRwFcmkdIPgtee4,29004
 jaxsim/api/frame.py,sha256=yPSgNygHkvWlln4wShNt7vZm_fFobVEm7phsklNNyH8,12922
-jaxsim/api/joint.py,sha256=Vl9VJs66_es88zjBBXqPzKkA3oAktV-DNiTkXwxOSsI,7562
-jaxsim/api/kin_dyn_parameters.py,sha256=eFcRJbfati3YsUhiJc5ZgngE-eqzYhqtud9L3ntQ-Uw,29632
-jaxsim/api/link.py,sha256=au47jV7bNdH2itvZ4qngKZyLi0UTIKzqVjjKosnlMsU,12858
-jaxsim/api/model.py,sha256=H87gCwt3_J8NGF_L3CBlqSNxiUmJWPHzP2v-XdrUx3A,79862
+jaxsim/api/joint.py,sha256=8rCIxRMeAidsaBbw7kkGp6z3-UmBPtqmYmV_arHDQJ8,7365
+jaxsim/api/kin_dyn_parameters.py,sha256=Y9wnMshz83Zm4UEPOAOTINdtfkBZ86w853c8Yi2qaVs,29670
+jaxsim/api/link.py,sha256=nHjffhNdi_xGkteMsqdb_hC9mdV9rNw7k3pl89Uhw_8,12798
+jaxsim/api/model.py,sha256=A88AaBZpWvQ-L9blFyl1GHvTWI05rvVFKbSaHzD77_k,79563
 jaxsim/api/ode.py,sha256=_t18avoCJngQk6eMFTGpaeahbpchQP20qJnUOCPkz8s,15360
 jaxsim/api/ode_data.py,sha256=1SD-x-lYk_YSEnVpxTLd69uOKC0mFUj44ZqpSmEDOxw,20190
-jaxsim/api/references.py,sha256=crPzgUCDHkTKUuXQcj9ygFAI0pewIx-fTPEsU8fBiU4,20555
+jaxsim/api/references.py,sha256=eIOk3MAOc9LJSKfI8M4WA8gGD-meo50vRfhXdea4sNI,20539
 jaxsim/integrators/__init__.py,sha256=hxvOD-VK_mmd6v31wtC-nb28AYve1gLuZCNLV9wS-Kg,103
-jaxsim/integrators/common.py,sha256=8Y3PGVOydPBJrqYtH33rtng48T12ySWfXTk-ED_2wao,18297
+jaxsim/integrators/common.py,sha256=ohISUnUWTaNHt2kweg1JyzwYGZgIH_wc-01qJWJsO80,18281
 jaxsim/integrators/fixed_step.py,sha256=KpjRd6hHtapxDoo6D1kyDrVDSHnke2TepI5grFH7_bM,2693
-jaxsim/integrators/variable_step.py,sha256=hGYKG3Sq3QITgzIePmCVCrrirwagqsKnB3aYifAcKR4,22848
+jaxsim/integrators/variable_step.py,sha256=Tqz5ySSgyKak_k6cTXpmtqdPNaFlO7N6zj7jBIlChyM,22681
 jaxsim/math/__init__.py,sha256=8oPITEoGwgRcOeG8KxtqxPQ8b5uku1HNRMokpCoi9Tc,352
 jaxsim/math/adjoint.py,sha256=V7r5VrTCKPLEL5gavNSx9U7xSsrb11a5e4gWqJ2MuRo,4375
 jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
@@ -31,28 +31,28 @@ jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=KXzQgOnCfAtbXCwxhplpJ3F0JT3oEyeLVby1_uRAryQ,2892
 jaxsim/mujoco/__init__.py,sha256=fZyRWre49pIhOrYdf6yJk_hOax8qWGe8OCmoq-dMVq8,201
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
-jaxsim/mujoco/loaders.py,sha256=CkFGydgOku5P_Pz7wdWlM2SCJRs71ePF-vsY9i90-I0,25350
+jaxsim/mujoco/loaders.py,sha256=_CZekIqZNe8oFeH7zSv4gGZAZENRISwMd8dt640zjRI,20860
 jaxsim/mujoco/model.py,sha256=5_7rWk_WBkNKDHqeewIFj0t2ZGqJpE6RDXHSbRvw4e4,16493
-jaxsim/mujoco/utils.py,sha256=bGbLMSzcdqbinIwHHJHt8ZN1uup_6DLdB2dWqKiXwO4,3955
+jaxsim/mujoco/utils.py,sha256=vZ8afASNOSxnxVW9p_1U1J_n-9nVhnBDqlV5k8c1GkM,8256
 jaxsim/mujoco/visualizer.py,sha256=nD6SNWmn-nxjjjIY9oPAHvL2j8q93DJDjZeepzke_DQ,6988
 jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxsim/parsers/kinematic_graph.py,sha256=wT2bgaCS8VQJTHy2H9sENkVPDOiMkRikxEF1t_WaahQ,34748
+jaxsim/parsers/kinematic_graph.py,sha256=MJkJ7AW1TdLZmxibuiVrTfn6jHjh3OVhEF20DqwsCnM,34748
 jaxsim/parsers/descriptions/__init__.py,sha256=PbIlunVfb59pB5jSX97YVpMAANRZPRkJ0X-hS14rzv4,221
 jaxsim/parsers/descriptions/collision.py,sha256=BQeIG-TKi4SVny23w6riDrQ5itC6VRwEMBX6HgAXHxA,3973
-jaxsim/parsers/descriptions/joint.py,sha256=VSb6C0FBBKMqwrHBKfc-Bbn4rl_J0RzUxMQlhIEvOPM,5185
+jaxsim/parsers/descriptions/joint.py,sha256=2KWLP4ILPMV8q1X0J7aS3GGFeZn4zXan0dqGOWc7XuQ,4365
 jaxsim/parsers/descriptions/link.py,sha256=Eh0W5qL7_Uw0GV-BkNKXhm9Q2dRTfIWCX5D-87zQkxA,3711
 jaxsim/parsers/descriptions/model.py,sha256=I2Vsbv8Josl4Le7b5rIvhqA2k9Bbv5JxMqwytayxds0,9833
 jaxsim/parsers/rod/__init__.py,sha256=G2vqlLajBLUc4gyzXwsEI2Wsi4TMOIF9bLDFeT6KrGU,92
-jaxsim/parsers/rod/parser.py,sha256=UAL6vkAab2y6PgIv3tAL8UodNbq-mmNBuBotxHJOSVU,14035
-jaxsim/parsers/rod/utils.py,sha256=5DsF3OeePZGidOJ5GiFSZx-51uIdnFvMW9EK6SgOW6Q,5698
+jaxsim/parsers/rod/parser.py,sha256=EXcbtr_vMjAaUzQjfQlD1zLLYLAZXrNeFHaiZVlLwFI,13976
+jaxsim/parsers/rod/utils.py,sha256=czQ2Y1_I9zGO0y2XDotHSqDorVH6zEcPhkuelApjs3k,5697
 jaxsim/rbda/__init__.py,sha256=kmy4G9aMkrqPNGdLSaSV3k15dpF52vBEUQXDFDuKIxU,337
 jaxsim/rbda/aba.py,sha256=w7ciyxB0IsmueatT0C7PcBQEl9dyiH9oqJgIi3xeTUE,8983
 jaxsim/rbda/collidable_points.py,sha256=0PFLzxWKtRg8-JtfNhGlSjBMv1J98tiLymOdvlvAak4,5325
 jaxsim/rbda/crba.py,sha256=bXkXESnVbv-lxhU1Y_i0rViEcQA4z2t2_jHwdVj5CBo,5049
 jaxsim/rbda/forward_kinematics.py,sha256=2GmEoWsrioVl_SAbKRKfhOLz57pY4aR81PKRdulqStA,3458
-jaxsim/rbda/jacobian.py,sha256=p0EV_8cLzLVV-93VKznT7VPuRj8W7h7rQWkPlWJXfCA,11023
+jaxsim/rbda/jacobian.py,sha256=L6Vn4Kf9I6wj-MYcFY6o67mgIfLFaaW4i2wNQJ2PDL0,10981
 jaxsim/rbda/rnea.py,sha256=CLfqs9XFVaD-hvkLABshDAfdw5bm_AMV3UVAQ_IvURQ,7542
-jaxsim/rbda/utils.py,sha256=eeT21Y4DiiyhrdF0lUE_VvRuwru5-rR7yOlOlWzCCWE,5381
+jaxsim/rbda/utils.py,sha256=GLt7XIl1ROkx0_fnBCKUHYdB9_IBF3Yi4OnkHSX3gxA,5365
 jaxsim/rbda/contacts/__init__.py,sha256=L5MM-2pv76YPGzxExdz2EErgGBATuAjYnNHlq5QOySs,503
 jaxsim/rbda/contacts/common.py,sha256=ai49HeLQOsnckG0H2tUKW2KQ0Au_v9jRuNdnqie-YBk,11234
 jaxsim/rbda/contacts/relaxed_rigid.py,sha256=tbyskONuUhC6BZnZSpNUnlCjkI7LR6mCtmU_HimOAVE,20893
@@ -60,13 +60,13 @@ jaxsim/rbda/contacts/rigid.py,sha256=MSzkU6SFbW6CryNlyyxQ7K0-U-8k6VROGKv_DQrwqiw
 jaxsim/rbda/contacts/soft.py,sha256=t6bqBfGAtV1AWoevY82LAcXy2XW8w_uu7bNywcyxF0s,17001
 jaxsim/rbda/contacts/visco_elastic.py,sha256=vQkfMuqQ3Qu8nbDTPY4jWBZjV3U7qtoRK1Aya3O3oFA,41424
 jaxsim/terrain/__init__.py,sha256=f7lVX-iNpH_wkkjef9Qpjh19TTAUOQw76EiLYJDVizc,78
-jaxsim/terrain/terrain.py,sha256=K91HEzPqTSyNrc_j1KfAAEF_5oDeuk_-jnnZGrcMEcY,5015
+jaxsim/terrain/terrain.py,sha256=_G1QS3zWycj089R8fTP5s2VjcZpEdJxREjXZJ-oXIvc,5248
 jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
-jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
-jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev327.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev327.dist-info/METADATA,sha256=hgujWPHg9YNK1R6nB4QvTd985MXx1jElYRRbwc_v2W0,17513
-jaxsim-0.4.3.dev327.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-jaxsim-0.4.3.dev327.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev327.dist-info/RECORD,,
+jaxsim/utils/tracing.py,sha256=eEY28MZW0Lm_jJNt1NkFqZz0ek01tvhR46OXZYCo7tc,532
+jaxsim/utils/wrappers.py,sha256=ZY7olSORzZRvSzkdeNLj8yjwUIAt9L0Douwl7wItjpk,4008
+jaxsim-0.4.3.dev350.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev350.dist-info/METADATA,sha256=qyh1wWUq5dTCw9iznLWlS4DlKa6kfMnAMqZgYbldbCA,17513
+jaxsim-0.4.3.dev350.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+jaxsim-0.4.3.dev350.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev350.dist-info/RECORD,,