jaxsim 0.6.2.dev281__py3-none-any.whl → 0.6.2.dev296__py3-none-any.whl

jaxsim/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.2.dev281'
-__version_tuple__ = version_tuple = (0, 6, 2, 'dev281')
+__version__ = version = '0.6.2.dev296'
+__version_tuple__ = version_tuple = (0, 6, 2, 'dev296')

jaxsim/api/kin_dyn_parameters.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+from typing import ClassVar
 
 import jax.lax
 import jax.numpy as jnp
@@ -9,8 +10,10 @@ import numpy as np
 import numpy.typing as npt
 from jax_dataclasses import Static
 
+import jaxsim
 import jaxsim.typing as jtp
-from jaxsim.math import Adjoint, Inertia, JointModel, supported_joint_motion
+from jaxsim.math import Inertia, JointModel, supported_joint_motion
+from jaxsim.math.adjoint import Adjoint
 from jaxsim.parsers.descriptions import JointDescription, JointType, ModelDescription
 from jaxsim.utils import HashedNumpyArray, JaxsimDataclass
 
@@ -30,6 +33,7 @@ class KinDynParameters(JaxsimDataclass):
         contact_parameters: The parameters of the collidable points.
         joint_model: The joint model of the model.
         joint_parameters: The parameters of the joints.
+        hw_link_metadata: The hardware parameters of the model links.
     """
 
     # Static
@@ -51,6 +55,9 @@ class KinDynParameters(JaxsimDataclass):
     joint_model: JointModel
     joint_parameters: JointParameters | None
 
+    # Model hardware parameters
+    hw_link_metadata: HwLinkMetadata | None = dataclasses.field(default=None)
+
     @property
     def motion_subspaces(self) -> jtp.Matrix:
         r"""
@@ -197,7 +204,6 @@ class KinDynParameters(JaxsimDataclass):
             carry0 = κb, link_index
 
             def scan_body(carry: tuple, i: jtp.Int) -> tuple[tuple, None]:
-
                 κb, active_link_index = carry
 
                 κb, active_link_index = jax.lax.cond(
@@ -224,7 +230,6 @@ class KinDynParameters(JaxsimDataclass):
         )
 
         def motion_subspace(joint_type: int, axis: npt.ArrayLike) -> npt.ArrayLike:
-
             S = {
                 JointType.Fixed: np.zeros(shape=(6, 1)),
                 JointType.Revolute: np.vstack(np.hstack([np.zeros(3), axis.axis])),
@@ -265,14 +270,12 @@ class KinDynParameters(JaxsimDataclass):
         )
 
     def __eq__(self, other: KinDynParameters) -> bool:
-
         if not isinstance(other, KinDynParameters):
             return False
 
         return hash(self) == hash(other)
 
     def __hash__(self) -> int:
-
         return hash(
             (
                 hash(self.number_of_links()),
@@ -671,7 +674,11 @@ class LinkParameters(JaxsimDataclass):
 
         return (
             jnp.hstack(
-                [params.mass, params.center_of_mass.squeeze(), params.inertia_elements]
+                [
+                    params.mass,
+                    params.center_of_mass.squeeze(),
+                    params.inertia_elements,
+                ]
             )
             .squeeze()
             .astype(float)
@@ -882,3 +889,276 @@ class FrameParameters(JaxsimDataclass):
         assert fp.transform.shape[0] == len(fp.body), fp.transform.shape[0]
 
         return fp
+
+
+@dataclasses.dataclass(frozen=True)
+class LinkParametrizableShape:
+    """
+    Enum-like class listing the supported shapes for HW parametrization.
+    """
+
+    Unsupported: ClassVar[int] = -1
+    Box: ClassVar[int] = 0
+    Cylinder: ClassVar[int] = 1
+    Sphere: ClassVar[int] = 2
+
+
+@jax_dataclasses.pytree_dataclass
+class HwLinkMetadata(JaxsimDataclass):
+    """
+    Class storing the hardware parameters of a link.
+
+    Attributes:
+        shape: The shape of the link.
+            0 = box, 1 = sphere, 2 = cylinder, -1 = unsupported.
+        dims: The dimensions of the link.
+            box: [lx,ly,lz], sphere: [r,0,0], cylinder: [r,l,0]
+        density: The density of the link.
+        L_H_G: The homogeneous transformation matrix from the link frame to the CoM frame G.
+        L_H_vis: The homogeneous transformation matrix from the link frame to the visual frame.
+        L_H_pre_mask: The mask indicating the link's child joint indices.
+        L_H_pre: The homogeneous transforms for child joints.
+    """
+
+    shape: jtp.Vector
+    dims: jtp.Vector
+    density: jtp.Float
+    L_H_G: jtp.Matrix
+    L_H_vis: jtp.Matrix
+    L_H_pre_mask: jtp.Vector
+    L_H_pre: jtp.Matrix
+
+    @staticmethod
+    def compute_mass_and_inertia(
+        hw_link_metadata: HwLinkMetadata,
+    ) -> tuple[jtp.Float, jtp.Matrix]:
+        """
+        Compute the mass and inertia of a hardware link based on its metadata.
+
+        This function calculates the mass and inertia tensor of a hardware link
+        using its shape, dimensions, and density. The computation is performed
+        by using shape-specific methods.
+
+        Args:
+            hw_link_metadata: Metadata describing the hardware link,
+                including its shape, dimensions, and density.
+
+        Returns:
+            tuple: A tuple containing:
+                - mass: The computed mass of the hardware link.
+                - inertia: The computed inertia tensor of the hardware link.
+        """
+
+        mass, inertia = jax.lax.switch(
+            hw_link_metadata.shape,
+            [
+                HwLinkMetadata._box,
+                HwLinkMetadata._cylinder,
+                HwLinkMetadata._sphere,
+            ],
+            hw_link_metadata.dims,
+            hw_link_metadata.density,
+        )
+        return mass, inertia
+
+    @staticmethod
+    def _box(dims, density) -> tuple[jtp.Float, jtp.Matrix]:
+        lx, ly, lz = dims
+
+        mass = density * lx * ly * lz
+
+        inertia = jnp.array(
+            [
+                [mass * (ly**2 + lz**2) / 12, 0, 0],
+                [0, mass * (lx**2 + lz**2) / 12, 0],
+                [0, 0, mass * (lx**2 + ly**2) / 12],
+            ]
+        )
+        return mass, inertia
+
+    @staticmethod
+    def _cylinder(dims, density) -> tuple[jtp.Float, jtp.Matrix]:
+        r, l, _ = dims
+
+        mass = density * (jnp.pi * r**2 * l)
+
+        inertia = jnp.array(
+            [
+                [mass * (3 * r**2 + l**2) / 12, 0, 0],
+                [0, mass * (3 * r**2 + l**2) / 12, 0],
+                [0, 0, mass * (r**2) / 2],
+            ]
+        )
+
+        return mass, inertia
+
+    @staticmethod
+    def _sphere(dims, density) -> tuple[jtp.Float, jtp.Matrix]:
+        r = dims[0]
+
+        mass = density * (4 / 3 * jnp.pi * r**3)
+
+        inertia = jnp.eye(3) * (2 / 5 * mass * r**2)
+
+        return mass, inertia
+
+    @staticmethod
+    def _convert_scaling_to_3d_vector(
+        shape: jtp.Int, scaling_factors: jtp.Vector
+    ) -> jtp.Vector:
+        """
+        Convert scaling factors for specific shape dimensions into a 3D scaling vector.
+
+        Args:
+            shape: The shape of the link (e.g., box, sphere, cylinder).
+            scaling_factors: The scaling factors for the shape dimensions.
+
+        Returns:
+            A 3D scaling vector to apply to position vectors.
+
+        Note:
+            The scaling factors are applied as follows to generate the 3D scale vector:
+            - Box: [lx, ly, lz]
+            - Cylinder: [r, r, l]
+            - Sphere: [r, r, r]
+        """
+        return jax.lax.switch(
+            shape,
+            branches=[
+                # Box
+                lambda: scaling_factors,
+                # Cylinder
+                lambda: jnp.array(
+                    [
+                        scaling_factors[0],
+                        scaling_factors[0],
+                        scaling_factors[1],
+                    ]
+                ),
+                # Sphere
+                lambda: jnp.array(
+                    [
+                        scaling_factors[0],
+                        scaling_factors[0],
+                        scaling_factors[0],
+                    ]
+                ),
+            ],
+        )
+
+    @staticmethod
+    def compute_inertia_link(I_com, mass, L_H_G) -> jtp.Matrix:
+        """
+        Compute the inertia tensor of the link based on its shape and mass.
+        """
+
+        L_R_G = L_H_G[:3, :3]
+        return L_R_G @ I_com @ L_R_G.T
+
+    @staticmethod
+    def apply_scaling(
+        hw_metadata: HwLinkMetadata, scaling_factors: ScalingFactors
+    ) -> HwLinkMetadata:
+        """
+        Apply scaling to the hardware parameters and return a new HwLinkMetadata object.
+
+        Args:
+            hw_metadata: the original HwLinkMetadata object.
+            scaling_factors: the scaling factors to apply.
+
+        Returns:
+            A new HwLinkMetadata object with updated parameters.
+        """
+
+        # ==================================
+        # Handle unsupported links
+        # ==================================
+        def unsupported_case(hw_metadata, scaling_factors):
+            # Return the metadata unchanged for unsupported links
+            return hw_metadata
+
+        def supported_case(hw_metadata, scaling_factors):
+            # ==================================
+            # Update the kinematics of the link
+            # ==================================
+
+            # Get the nominal transforms
+            L_H_G = hw_metadata.L_H_G
+            L_H_vis = hw_metadata.L_H_vis
+            L_H_pre_array = hw_metadata.L_H_pre
+            L_H_pre_mask = hw_metadata.L_H_pre_mask
+
+            # Compute the 3D scaling vector
+            scale_vector = HwLinkMetadata._convert_scaling_to_3d_vector(
+                hw_metadata.shape, scaling_factors.dims
+            )
+
+            # Express the transforms in the G frame
+            G_H_L = jaxsim.math.Transform.inverse(L_H_G)
+            G_H_vis = G_H_L @ L_H_vis
+            G_H_pre_array = jax.vmap(lambda L_H_pre: G_H_L @ L_H_pre)(L_H_pre_array)
+
+            # Apply the scaling to the position vectors
+            G_H̅_L = G_H_L.at[:3, 3].set(scale_vector * G_H_L[:3, 3])
+            G_H̅_vis = G_H_vis.at[:3, 3].set(scale_vector * G_H_vis[:3, 3])
+            # Apply scaling to the position vectors in G_H_pre_array based on the mask
+            G_H̅_pre_array = jax.vmap(
+                lambda G_H_pre, mask: jnp.where(
+                    # Expand mask for broadcasting
+                    mask[..., None, None],
+                    # Apply scaling
+                    G_H_pre.at[:3, 3].set(scale_vector * G_H_pre[:3, 3]),
+                    # Keep unchanged if mask is False
+                    G_H_pre,
+                )
+            )(G_H_pre_array, L_H_pre_mask)
+
+            # Get back to the link frame
+            L_H̅_G = jaxsim.math.Transform.inverse(G_H̅_L)
+            L_H̅_vis = L_H̅_G @ G_H̅_vis
+            L_H̅_pre_array = jax.vmap(lambda G_H̅_pre: L_H̅_G @ G_H̅_pre)(G_H̅_pre_array)
+
+            # ============================
+            # Update the shape parameters
+            # ============================
+
+            updated_dims = hw_metadata.dims * scaling_factors.dims
+
+            # ==============================
+            # Scale the density of the link
+            # ==============================
+
+            updated_density = hw_metadata.density * scaling_factors.density
+
+            # ============================
+            # Return updated HwLinkMetadata
+            # ============================
+
+            return hw_metadata.replace(
+                dims=updated_dims,
+                density=updated_density,
+                L_H_G=L_H̅_G,
+                L_H_vis=L_H̅_vis,
+                L_H_pre=L_H̅_pre_array,
+            )
+
+        # Use jax.lax.cond to handle unsupported links
+        return jax.lax.cond(
+            hw_metadata.shape == LinkParametrizableShape.Unsupported,
+            lambda: unsupported_case(hw_metadata, scaling_factors),
+            lambda: supported_case(hw_metadata, scaling_factors),
+        )
+
+
+@jax_dataclasses.pytree_dataclass
+class ScalingFactors(JaxsimDataclass):
+    """
+    Class storing scaling factors for hardware parameters.
+
+    Attributes:
+        dims: Scaling factors for shape dimensions.
+        density: Scaling factor for density.
+    """
+
+    dims: jtp.Vector
+    density: jtp.Float

jaxsim/api/model.py

@@ -11,14 +11,26 @@ import jax
 import jax.numpy as jnp
 import jax_dataclasses
 import rod
+import rod.urdf
 from jax_dataclasses import Static
+from rod.urdf.exporter import UrdfExporter
 
 import jaxsim.api as js
 import jaxsim.exceptions
 import jaxsim.terrain
 import jaxsim.typing as jtp
+from jaxsim import logging
+from jaxsim.api.kin_dyn_parameters import (
+    HwLinkMetadata,
+    KinDynParameters,
+    LinkParameters,
+    LinkParametrizableShape,
+    ScalingFactors,
+)
 from jaxsim.math import Adjoint, Cross
 from jaxsim.parsers.descriptions import ModelDescription
+from jaxsim.parsers.descriptions.joint import JointDescription
+from jaxsim.parsers.descriptions.link import LinkDescription
 from jaxsim.utils import JaxsimDataclass, Mutability, wrappers
 
 from .common import VelRepr
@@ -86,7 +98,6 @@ class JaxSimModel(JaxsimDataclass):
         return self._description.get()
 
     def __eq__(self, other: JaxSimModel) -> bool:
-
         if not isinstance(other, JaxSimModel):
             return False
 
@@ -102,7 +113,6 @@ class JaxSimModel(JaxsimDataclass):
         return True
 
     def __hash__(self) -> int:
-
         return hash(
             (
                 hash(self.model_name),
@@ -194,6 +204,12 @@ class JaxSimModel(JaxsimDataclass):
         with model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
             model.built_from = model_description
 
+        # Compute the hw parametrization metadata of the model
+        # TODO: move the building of the metadata to KinDynParameters.build()
+        #       and use the model_description instead of model.built_from.
+        with model.mutable_context(mutability=Mutability.MUTABLE_NO_VALIDATION):
+            model.kin_dyn_parameters.hw_link_metadata = model.compute_hw_link_metadata()
+
         return model
 
     @classmethod
@@ -297,6 +313,274 @@ class JaxSimModel(JaxsimDataclass):
 
         return model
 
+    def compute_hw_link_metadata(self) -> HwLinkMetadata:
+        """
+        Compute the parametric metadata of the links in the model.
+
+        Returns:
+            An instance of HwLinkMetadata containing the metadata of all links.
+        """
+        model_description = self.description
+
+        # Get ordered links and joints from the model description
+        ordered_links: list[LinkDescription] = sorted(
+            list(model_description.links_dict.values()),
+            key=lambda l: l.index,
+        )
+        ordered_joints: list[JointDescription] = sorted(
+            list(model_description.joints_dict.values()),
+            key=lambda j: j.index,
+        )
+
+        # Ensure the model was built from a valid source
+        rod_model = None
+        match self.built_from:
+            case str() | pathlib.Path():
+                rod_model = rod.Sdf.load(sdf=self.built_from).models()[0]
+                assert rod_model.name == self.name()
+            case rod.Model():
+                rod_model = self.built_from
+            case _:
+                logging.debug(
+                    f"Invalid type for model.built_from ({type(self.built_from)})."
+                    "Skipping for hardware parametrization."
+                )
+                return HwLinkMetadata(
+                    shape=jnp.array([]),
+                    dims=jnp.array([]),
+                    density=jnp.array([]),
+                    L_H_G=jnp.array([]),
+                    L_H_vis=jnp.array([]),
+                    L_H_pre_mask=jnp.array([]),
+                    L_H_pre=jnp.array([]),
+                )
+
+        # Use URDF frame convention for consistent pose representation
+        rod_model.switch_frame_convention(
+            frame_convention=rod.FrameConvention.Urdf, explicit_frames=True
+        )
+
+        rod_links_dict = {}
+
+        # Filter links that support parameterization
+        for rod_link in rod_model.links():
+            if len(rod_link.visuals()) != 1:
+                logging.debug(
+                    f"Skipping link '{rod_link.name}' for hardware parametrization due to multiple visuals."
+                )
+                continue
+
+            if not isinstance(
+                rod_link.visual.geometry.geometry(), (rod.Box, rod.Sphere, rod.Cylinder)
+            ):
+                logging.debug(
+                    f"Skipping link '{rod_link.name}' for hardware parametrization due to unsupported geometry."
+                )
+                continue
+
+            rod_links_dict[rod_link.name] = rod_link
+
+        # Initialize lists to collect metadata for all links
+        shapes = []
+        dims = []
+        densities = []
+        L_H_Gs = []
+        L_H_vises = []
+        L_H_pre_masks = []
+        L_H_pres = []
+
+        # Process each link
+        for link_description in ordered_links:
+            link_name = link_description.name
+
+            if link_name not in self.link_names():
+                logging.debug(
+                    f"Skipping link '{link_name}' for hardware parametrization as it is not part of the JaxSim model."
+                )
+                continue
+
+            if link_name not in rod_links_dict:
+                logging.debug(
+                    f"Skipping link '{link_name}' for hardware parametrization as it is not part of the ROD model."
+                )
+                continue
+
+            rod_link = rod_links_dict[link_name]
+            link_index = int(js.link.name_to_idx(model=self, link_name=link_name))
+
+            # Get child joints for the link
+            child_joints_indices = [
+                js.joint.name_to_idx(model=self, joint_name=j.name)
+                for j in ordered_joints
+                if j.parent.name == link_name
+            ]
+
+            # Skip unsupported links
+            if not jnp.allclose(
+                self.kin_dyn_parameters.joint_model.suc_H_i[link_index], jnp.eye(4)
+            ):
+                logging.debug(
+                    f"Skipping link '{link_name}' for hardware parametrization due to unsupported suc_H_link."
+                )
+                continue
+
+            # Compute density and dimensions
+            mass = float(self.kin_dyn_parameters.link_parameters.mass[link_index])
+            geometry = rod_link.visual.geometry.geometry()
+            if isinstance(geometry, rod.Box):
+                lx, ly, lz = geometry.size
+                density = mass / (lx * ly * lz)
+                dims.append([lx, ly, lz])
+                shapes.append(LinkParametrizableShape.Box)
+            elif isinstance(geometry, rod.Sphere):
+                r = geometry.radius
+                density = mass / (4 / 3 * jnp.pi * r**3)
+                dims.append([r, 0, 0])
+                shapes.append(LinkParametrizableShape.Sphere)
+            elif isinstance(geometry, rod.Cylinder):
+                r, l = geometry.radius, geometry.length
+                density = mass / (jnp.pi * r**2 * l)
+                dims.append([r, l, 0])
+                shapes.append(LinkParametrizableShape.Cylinder)
+            else:
+                logging.debug(
+                    f"Skipping link '{link_name}' for hardware parametrization due to unsupported geometry."
+                )
+                continue
+
+            densities.append(density)
+            L_H_Gs.append(rod_link.inertial.pose.transform())
+            L_H_vises.append(rod_link.visual.pose.transform())
+            L_H_pre_masks.append(
+                [
+                    int(joint_index in child_joints_indices)
+                    for joint_index in range(self.number_of_joints())
+                ]
+            )
+            L_H_pres.append(
+                [
+                    (
+                        self.kin_dyn_parameters.joint_model.λ_H_pre[joint_index + 1]
+                        if joint_index in child_joints_indices
+                        else jnp.eye(4)
+                    )
+                    for joint_index in range(self.number_of_joints())
+                ]
+            )
+
+        # Stack collected data into JAX arrays
+        return HwLinkMetadata(
+            shape=jnp.array(shapes, dtype=int),
+            dims=jnp.array(dims, dtype=float),
+            density=jnp.array(densities, dtype=float),
+            L_H_G=jnp.array(L_H_Gs, dtype=float),
+            L_H_vis=jnp.array(L_H_vises, dtype=float),
+            L_H_pre_mask=jnp.array(L_H_pre_masks, dtype=bool),
+            L_H_pre=jnp.array(L_H_pres, dtype=float),
+        )
+
+    def export_updated_model(self) -> str:
+        """
+        Export the JaxSim model to URDF with the current hardware parameters.
+
+        Returns:
+            The URDF string of the updated model.
+
+        Note:
+            This method is not meant to be used in JIT-compiled functions.
+        """
+
+        import numpy as np
+
+        if isinstance(jnp.zeros(0), jax.core.Tracer):
+            raise RuntimeError("This method cannot be used in JIT-compiled functions")
+
+        # Ensure `built_from` is a ROD model and create `rod_model_output`
+        if isinstance(self.built_from, rod.Model):
+            rod_model_output = copy.deepcopy(self.built_from)
+        elif isinstance(self.built_from, (str, pathlib.Path)):
+            rod_model_output = rod.Sdf.load(sdf=self.built_from).models()[0]
+        else:
+            raise ValueError(
+                "The JaxSim model must be built from a valid ROD model source"
+            )
+
+        # Switch to URDF frame convention for easier mapping
+        rod_model_output.switch_frame_convention(
+            frame_convention=rod.FrameConvention.Urdf,
+            explicit_frames=True,
+            attach_frames_to_links=True,
+        )
+
+        # Get links and joints from the ROD model
+        links_dict = {link.name: link for link in rod_model_output.links()}
+        joints_dict = {joint.name: joint for joint in rod_model_output.joints()}
+
+        # Iterate over the hardware metadata to update the ROD model
+        hw_metadata = self.kin_dyn_parameters.hw_link_metadata
+        for link_index, link_name in enumerate(self.link_names()):
+            if link_name not in links_dict:
+                continue
+
+            # Update mass and inertia
+            mass = float(self.kin_dyn_parameters.link_parameters.mass[link_index])
+            center_of_mass = np.array(
+                self.kin_dyn_parameters.link_parameters.center_of_mass[link_index]
+            )
+            inertia_tensor = LinkParameters.unflatten_inertia_tensor(
+                self.kin_dyn_parameters.link_parameters.inertia_elements[link_index]
+            )
+
+            links_dict[link_name].inertial.mass = mass
+            L_H_COM = np.eye(4)
+            L_H_COM[0:3, 3] = center_of_mass
+            links_dict[link_name].inertial.pose = rod.Pose.from_transform(
+                transform=L_H_COM,
+                relative_to=links_dict[link_name].inertial.pose.relative_to,
+            )
+            links_dict[link_name].inertial.inertia = rod.Inertia.from_inertia_tensor(
+                inertia_tensor=inertia_tensor, validate=True
+            )
+
+            # Update visual shape
+            shape = hw_metadata.shape[link_index]
+            dims = hw_metadata.dims[link_index]
+            if shape == LinkParametrizableShape.Box:
+                links_dict[link_name].visual.geometry.box.size = dims.tolist()
+            elif shape == LinkParametrizableShape.Sphere:
+                links_dict[link_name].visual.geometry.sphere.radius = float(dims[0])
+            elif shape == LinkParametrizableShape.Cylinder:
+                links_dict[link_name].visual.geometry.cylinder.radius = float(dims[0])
+                links_dict[link_name].visual.geometry.cylinder.length = float(dims[1])
+            else:
+                logging.debug(f"Skipping unsupported shape for link '{link_name}'")
+                continue
+
+            # Update visual pose
+            links_dict[link_name].visual.pose = rod.Pose.from_transform(
+                transform=np.array(hw_metadata.L_H_vis[link_index]),
+                relative_to=links_dict[link_name].visual.pose.relative_to,
+            )
+
+            # Update joint poses
+            for joint_index in range(self.number_of_joints()):
+                if hw_metadata.L_H_pre_mask[link_index, joint_index]:
+                    joint_name = js.joint.idx_to_name(
+                        model=self, joint_index=joint_index
+                    )
+                    if joint_name in joints_dict:
+                        joints_dict[joint_name].pose = rod.Pose.from_transform(
+                            transform=np.array(
+                                hw_metadata.L_H_pre[link_index, joint_index]
+                            ),
+                            relative_to=joints_dict[joint_name].pose.relative_to,
+                        )
+
+        # Export the URDF string.
+        urdf_string = UrdfExporter(pretty=True).to_urdf_string(sdf=rod_model_output)
+
+        return urdf_string
+
     # ==========
     # Properties
     # ==========
@@ -585,9 +869,7 @@ def generalized_free_floating_jacobian(
     # ======================================================================
 
     match data.velocity_representation:
-
         case VelRepr.Inertial:
-
             W_H_B = data._base_transform
             B_X_W = Adjoint.from_transform(transform=W_H_B, inverse=True)
 
@@ -597,11 +879,9 @@ def generalized_free_floating_jacobian(
             )
 
         case VelRepr.Body:
-
             B_J_full_WX_I = B_J_full_WX_B
 
         case VelRepr.Mixed:
-
             W_R_B = jaxsim.math.Quaternion.to_dcm(data.base_orientation)
             BW_H_B = jnp.eye(4).at[0:3, 0:3].set(W_R_B)
             B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
@@ -633,9 +913,7 @@ def generalized_free_floating_jacobian(
     # =======================================================================
 
     match output_vel_repr:
-
         case VelRepr.Inertial:
-
             W_H_B = data._base_transform
             W_X_B = jaxsim.math.Adjoint.from_transform(W_H_B)
 
@@ -644,7 +922,6 @@ def generalized_free_floating_jacobian(
             )(B_J_WL_I)
 
         case VelRepr.Body:
-
             O_J_WL_I = L_J_WL_I = jax.vmap(  # noqa: F841
                 lambda B_H_L, B_J_WL_I: jaxsim.math.Adjoint.from_transform(
                     B_H_L, inverse=True
@@ -653,7 +930,6 @@ def generalized_free_floating_jacobian(
             )(B_H_L, B_J_WL_I)
 
         case VelRepr.Mixed:
-
             W_H_B = data._base_transform
 
             LW_H_L = jax.vmap(
@@ -738,9 +1014,7 @@ def generalized_free_floating_jacobian_derivative(
     On = jnp.zeros(shape=(model.dofs(), model.dofs()))
 
     match data.velocity_representation:
-
         case VelRepr.Inertial:
-
             B_X_W = jaxsim.math.Adjoint.from_transform(transform=W_H_B, inverse=True)
 
             W_v_WB = data.base_velocity
@@ -752,7 +1026,6 @@ def generalized_free_floating_jacobian_derivative(
             Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_W, On)
 
         case VelRepr.Body:
-
             B_X_B = jaxsim.math.Adjoint.from_rotation_and_translation(
                 translation=jnp.zeros(3), rotation=jnp.eye(3)
             )
@@ -765,7 +1038,6 @@ def generalized_free_floating_jacobian_derivative(
             Ṫ = jax.scipy.linalg.block_diag(B_Ẋ_B, On)
 
         case VelRepr.Mixed:
-
             BW_H_B = W_H_B.at[0:3, 3].set(jnp.zeros(3))
             B_X_BW = jaxsim.math.Adjoint.from_transform(transform=BW_H_B, inverse=True)
 
@@ -788,9 +1060,7 @@ def generalized_free_floating_jacobian_derivative(
     # ======================================================
 
     match output_vel_repr:
-
         case VelRepr.Inertial:
-
             O_X_B = W_X_B = jaxsim.math.Adjoint.from_transform(transform=W_H_B)
 
             with data.switch_velocity_representation(VelRepr.Body):
@@ -799,7 +1069,6 @@ def generalized_free_floating_jacobian_derivative(
             O_Ẋ_B = W_Ẋ_B = W_X_B @ jaxsim.math.Cross.vx(B_v_WB)  # noqa: F841
 
         case VelRepr.Body:
-
             O_X_B = L_X_B = jaxsim.math.Adjoint.from_transform(
                 transform=B_H_L, inverse=True
             )
@@ -817,7 +1086,6 @@ def generalized_free_floating_jacobian_derivative(
             )
 
         case VelRepr.Mixed:
-
             W_H_L = W_H_B @ B_H_L
             LW_H_L = W_H_L.at[:, 0:3, 3].set(jnp.zeros_like(W_H_L[:, 0:3, 3]))
             LW_H_B = LW_H_L @ jaxsim.math.Transform.inverse(B_H_L)
@@ -1190,14 +1458,12 @@ def free_floating_mass_matrix(
             return M_body
 
         case VelRepr.Inertial:
-
             B_X_W = Adjoint.from_transform(transform=data._base_transform, inverse=True)
             invT = jax.scipy.linalg.block_diag(B_X_W, jnp.eye(model.dofs()))
 
             return invT.T @ M_body @ invT
 
         case VelRepr.Mixed:
-
             BW_H_B = data._base_transform.at[0:3, 3].set(jnp.zeros(3))
             B_X_BW = Adjoint.from_transform(transform=BW_H_B, inverse=True)
             invT = jax.scipy.linalg.block_diag(B_X_BW, jnp.eye(model.dofs()))
@@ -1233,7 +1499,6 @@ def free_floating_coriolis_matrix(
     # The Coriolis matrix computed in this representation is converted later
     # to the active representation stored in data.
     with data.switch_velocity_representation(VelRepr.Body):
-
         B_ν = data.generalized_velocity
 
         # Doubly-left free-floating Jacobian.
@@ -1251,7 +1516,6 @@ def free_floating_coriolis_matrix(
 
     # Compute the contribution of each link to the Coriolis matrix.
     def compute_link_contribution(M, v, J, J̇) -> jtp.Array:
-
         return J.T @ ((Cross.vx_star(v) @ M + M @ Cross.vx(v)) @ J + M @ J̇)
 
     C_B_links = jax.vmap(compute_link_contribution)(
@@ -1274,12 +1538,10 @@ def free_floating_coriolis_matrix(
     # Adjust the representation of the Coriolis matrix.
     # Refer to https://github.com/traversaro/traversaro-phd-thesis, Section 3.6.
     match data.velocity_representation:
-
         case VelRepr.Body:
             return C_B
 
         case VelRepr.Inertial:
-
             n = model.dofs()
             W_H_B = data._base_transform
             B_X_W = jaxsim.math.Adjoint.from_transform(W_H_B, inverse=True)
@@ -1299,7 +1561,6 @@ def free_floating_coriolis_matrix(
             return C
 
         case VelRepr.Mixed:
-
             n = model.dofs()
             BW_H_B = data._base_transform.at[0:3, 3].set(jnp.zeros(3))
             B_X_BW = jaxsim.math.Adjoint.from_transform(transform=BW_H_B, inverse=True)
@@ -1720,9 +1981,7 @@ def average_velocity_jacobian(
     G_J = js.com.average_centroidal_velocity_jacobian(model=model, data=data)
 
     match output_vel_repr:
-
         case VelRepr.Inertial:
-
             GW_J = G_J
             W_p_CoM = js.com.com_position(model=model, data=data)
 
@@ -1732,7 +1991,6 @@ def average_velocity_jacobian(
             return W_X_GW @ GW_J
 
         case VelRepr.Body:
-
             GB_J = G_J
             W_p_B = data.base_position
             W_p_CoM = js.com.com_position(model=model, data=data)
@@ -1744,7 +2002,6 @@ def average_velocity_jacobian(
             return B_X_GB @ GB_J
 
         case VelRepr.Mixed:
-
             GW_J = G_J
             W_p_B = data.base_position
             W_p_CoM = js.com.com_position(model=model, data=data)
@@ -2039,6 +2296,104 @@ def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.F
     return jnp.sum((m * W_p̃_CoM)[2] * model.gravity)
 
 
+# ===================
+# Hw parametrization
+# ===================
+
+
+@jax.jit
+@js.common.named_scope
+def update_hw_parameters(
+    model: JaxSimModel, scaling_factors: ScalingFactors
+) -> JaxSimModel:
+    """
+    Update the hardware parameters of the model by scaling the parameters of the links.
+
+    This function applies scaling factors to the hardware metadata of the links,
+    updating their shape, dimensions, density, and other related parameters. It
+    recalculates the mass and inertia tensors of the links based on the updated
+    metadata and adjusts the joint model transforms accordingly.
+
+    Args:
+        model: The JaxSimModel object to update.
+        scaling_factors: A ScalingFactors object containing scaling factors for
+                         dimensions and density of the links.
+
+    Returns:
+        The updated JaxSimModel object with modified hardware parameters.
+    """
+    kin_dyn_params: KinDynParameters = model.kin_dyn_parameters
+    link_parameters: LinkParameters = kin_dyn_params.link_parameters
+    hw_link_metadata: HwLinkMetadata = kin_dyn_params.hw_link_metadata
+
+    # Apply scaling to hw_link_metadata using vmap
+    updated_hw_link_metadata = jax.vmap(HwLinkMetadata.apply_scaling)(
+        hw_link_metadata, scaling_factors
+    )
+
+    # Compute mass and inertia once and unpack the results
+    m_updated, I_com_updated = jax.vmap(HwLinkMetadata.compute_mass_and_inertia)(
+        updated_hw_link_metadata
+    )
+
+    # Rotate the inertia tensor at CoM with the link orientation, and store
+    # it in KynDynParameters.
+    I_L_updated = jax.vmap(
+        lambda metadata, I_com: metadata.L_H_G[:3, :3]
+        @ I_com
+        @ metadata.L_H_G[:3, :3].T
+    )(updated_hw_link_metadata, I_com_updated)
+
+    # Update link parameters
+    updated_link_parameters = link_parameters.replace(
+        mass=m_updated,
+        inertia_elements=jax.vmap(LinkParameters.flatten_inertia_tensor)(I_L_updated),
+        center_of_mass=jax.vmap(lambda metadata: metadata.L_H_G[:3, 3])(
+            updated_hw_link_metadata
+        ),
+    )
+
+    # Update joint model transforms (λ_H_pre)
+    def update_λ_H_pre(joint_index):
+        # Extract the transforms and masks for the current joint index across all links
+        L_H_pre_for_joint = updated_hw_link_metadata.L_H_pre[:, joint_index]
+        L_H_pre_mask_for_joint = updated_hw_link_metadata.L_H_pre_mask[:, joint_index]
+
+        # Use the mask to select the first valid transform or fall back to the original
+        valid_transforms = jnp.where(
+            L_H_pre_mask_for_joint[:, None, None],  # Expand mask for broadcasting
+            L_H_pre_for_joint,  # Use the transform if the mask is True
+            jnp.zeros_like(L_H_pre_for_joint),  # Otherwise, use a zero matrix
+        )
+
+        # Sum the valid transforms (only one will be non-zero due to the mask)
+        selected_transform = jnp.sum(valid_transforms, axis=0)
+
+        # If no valid transform exists, fall back to the original λ_H_pre
+        return jax.lax.cond(
+            jnp.any(L_H_pre_mask_for_joint),
+            lambda: selected_transform,
+            lambda: kin_dyn_params.joint_model.λ_H_pre[joint_index],
+        )
+
+    # Apply the update function to all joint indices
+    updated_λ_H_pre = jax.vmap(update_λ_H_pre)(
+        jnp.arange(kin_dyn_params.number_of_joints() + 1)
+    )
+    # Replace the joint model with the updated transforms
+    updated_joint_model = kin_dyn_params.joint_model.replace(λ_H_pre=updated_λ_H_pre)
+
+    # Replace the kin_dyn_parameters with updated values
+    updated_kin_dyn_params = kin_dyn_params.replace(
+        link_parameters=updated_link_parameters,
+        hw_link_metadata=updated_hw_link_metadata,
+        joint_model=updated_joint_model,
+    )
+
+    # Return the updated model
+    return model.replace(kin_dyn_parameters=updated_kin_dyn_params)
+
+
 # ==========
 # Simulation
 # ==========

jaxsim/math/joint_model.py

@@ -10,10 +10,11 @@ import jaxsim.typing as jtp
 from jaxsim.math import Rotation
 from jaxsim.parsers.descriptions import JointGenericAxis, JointType, ModelDescription
 from jaxsim.parsers.kinematic_graph import KinematicGraphTransforms
+from jaxsim.utils.jaxsim_dataclass import JaxsimDataclass
 
 
 @jax_dataclasses.pytree_dataclass
-class JointModel:
+class JointModel(JaxsimDataclass):
     """
     Class describing the joint kinematics of a robot model.
 

jaxsim/mujoco/utils.py

@@ -71,11 +71,7 @@ def mujoco_data_from_jaxsim(
 
         model_helper.set_joint_positions(
             joint_names=list(jaxsim_model.joint_names()),
-            positions=np.array(
-                jaxsim_data.joint_positions(
-                    model=jaxsim_model, joint_names=jaxsim_model.joint_names()
-                )
-            ),
+            positions=np.array(jaxsim_data.joint_positions),
         )
 
     # Updating these joints is not necessary after the first time.

{jaxsim-0.6.2.dev281.dist-info → jaxsim-0.6.2.dev296.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev281
+Version: 0.6.2.dev296
 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.6.2.dev281.dist-info → jaxsim-0.6.2.dev296.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=EKeysKN-7UswwJLCl7n6qIBKQIVUtYsCMYu_tCoFn7g,3628
-jaxsim/_version.py,sha256=1YD2xmwjc7SXvsmFaPIxqihf7l6csMU3_MsQ0uM_T9Q,528
+jaxsim/_version.py,sha256=Bzgp_Dut_GzvkRH6pABuIcblUVxBbl7RquWxbbiRawk,528
 jaxsim/exceptions.py,sha256=MQ3LRMfVMX2-g3qYj7mUVNV9OLlIA48TANJegbcQyXI,2641
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=7msl8t5Jt09RNYfKdPJtpjLfWurldcycDappb045Eso,761
@@ -12,16 +12,16 @@ jaxsim/api/data.py,sha256=9pxug2gFIZPwqi9kNYXhEziA5IQBB9KNNwIfPfc_kAU,23249
 jaxsim/api/frame.py,sha256=4wg6GsyBQgYhSvc-ry_31JsaL66sZt3TtgwjB7NrHmk,14583
 jaxsim/api/integrators.py,sha256=sHdTWw2Z-Q7jggA8zRkA1KYYd4HNIozXPwNvFwt0YBs,9011
 jaxsim/api/joint.py,sha256=AnqlNWmBOay-gsoo0y4AbfFQ2OCJm-8T1E0IMhZeLoY,7457
-jaxsim/api/kin_dyn_parameters.py,sha256=v1ZDb-Eqn552Zb344sJYy9e5sQT--_SYGIdtNC2JVSg,29828
+jaxsim/api/kin_dyn_parameters.py,sha256=r1rOjKfe7Qg7xBWyCNtlAi_Wg1WTzToQrUdbZ7PO10Q,39073
 jaxsim/api/link.py,sha256=bSZOYJDY9HJMgY25VzevTTsdFZTJc6yRRpslc5FhGHE,12782
-jaxsim/api/model.py,sha256=Do5q7vSyml854wFyN21qJtNfJfFo7TlEfxx_FoTxOW0,69948
+jaxsim/api/model.py,sha256=ppdriJBHWJ-qXey9Vjqnd7IjqdEE0R6W9ZG9Y7KuS2s,85460
 jaxsim/api/ode.py,sha256=fp20_LK9lXw2DfNkQgrfJmtd_iBXDNzZkOn0u5Pm8Qw,6190
 jaxsim/api/references.py,sha256=-vd50y3v-jkXAsILS432etIKV6e2EUE2oOeLHuUrfuQ,20399
 jaxsim/math/__init__.py,sha256=dNozvtm8WsB7nxw4uK29yQQKPcDUEczr2zcHoZfMItc,383
 jaxsim/math/adjoint.py,sha256=Pb0WAiAoN1ge8j_dPcdK307jmC5LzD1-DeUj9Z_NXkI,4667
 jaxsim/math/cross.py,sha256=AM4HauuuT09q2TN42qvdXhJ9LvtCh0e7ZyLjP-7sANs,1498
 jaxsim/math/inertia.py,sha256=T-iAjPYSD_72R0ZG8GDJhe5i3Jc3ojhlbBRSscTdCKg,1577
-jaxsim/math/joint_model.py,sha256=mV6F3SVHUxHcXMIw-hc4fG341R4rkK8MvRPn53HcEpg,6915
+jaxsim/math/joint_model.py,sha256=vBnwXSsw2LCb2Tr5wl2iCo0KvLqcibBbeKcsoH5r9tk,6990
 jaxsim/math/quaternion.py,sha256=MSaZywzJDxs2te1ZELeIcupKSFIA9q_pdXy7fDAEqM4,4539
 jaxsim/math/rotation.py,sha256=TEUtT3X2tFieNxdlccup1pfaTgCTtfX-hTNotd8-nNk,1892
 jaxsim/math/skew.py,sha256=z_9YN-NDHL3n4KXWNbzTSMkFDZ0SDpz4RUcwwYFOaao,1402
@@ -31,7 +31,7 @@ jaxsim/mujoco/__init__.py,sha256=1kAWzYOS7nP29S5FGyWPqiAnPf4yPSoaPW-WBGBjVV0,214
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
 jaxsim/mujoco/loaders.py,sha256=OCk1T11iIm3qZUibNpo_bxxLgaGSkCpLt7ae_ND0ExA,23272
 jaxsim/mujoco/model.py,sha256=bRXo1uhWDN37sP9qdejr_2vq_PXHQ7p6eyBlFff_JcE,16492
-jaxsim/mujoco/utils.py,sha256=r8BFrFAkW3lMAnEM3YUGZYLBNT2ptln4qDPm68CYrJk,8951
+jaxsim/mujoco/utils.py,sha256=q75OSjxLU2BromVUejt0DVnSbrV5D177YW6LkOdu78g,8823
 jaxsim/mujoco/visualizer.py,sha256=cmI6DhFb1XC7oEtg_wl-s-U56dWHA-F7GlBD6EDYTyA,7744
 jaxsim/parsers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxsim/parsers/kinematic_graph.py,sha256=ARq11Pv6yMDZeRRDlqrWzfVfUS7qSDwR57aWA4k54as,35758
@@ -65,8 +65,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=XzmZeIibcaOzaxpprsGSxH3UrM66PAO456rFV91sNXg,11453
 jaxsim/utils/tracing.py,sha256=Btwxdfhb7fJLk3r5PlQkGYj60Y2KbFT1gANGIA697FU,530
 jaxsim/utils/wrappers.py,sha256=3IMwydqFgmSPqeuUQ3PRmdhDc1IoT6XC23jPC_LjWXs,4175
-jaxsim-0.6.2.dev281.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.6.2.dev281.dist-info/METADATA,sha256=G9vLGI8GOu0nIih0Bf6yGAIzkiRIntqAzdISuvAN5_w,19658
-jaxsim-0.6.2.dev281.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-jaxsim-0.6.2.dev281.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.6.2.dev281.dist-info/RECORD,,
+jaxsim-0.6.2.dev296.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.6.2.dev296.dist-info/METADATA,sha256=4jgGP0Aq1uGxo0AybZcSHs9U8Xqxq-noHcs6ctLrFWM,19658
+jaxsim-0.6.2.dev296.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+jaxsim-0.6.2.dev296.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.6.2.dev296.dist-info/RECORD,,

{jaxsim-0.6.2.dev281.dist-info → jaxsim-0.6.2.dev296.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (79.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any