jaxsim 0.2.dev191__py3-none-any.whl → 0.6.1.dev5__py3-none-any.whl

jaxsim/parsers/rod/meshes.py

@@ -0,0 +1,104 @@
+import numpy as np
+import trimesh
+
+VALID_AXIS = {"x": 0, "y": 1, "z": 2}
+
+
+def extract_points_vertices(mesh: trimesh.Trimesh) -> np.ndarray:
+    """
+    Extract the vertices of a mesh as points.
+    """
+    return mesh.vertices
+
+
+def extract_points_random_surface_sampling(mesh: trimesh.Trimesh, n) -> np.ndarray:
+    """
+    Extract N random points from the surface of a mesh.
+
+    Args:
+        mesh: The mesh from which to extract points.
+        n: The number of points to extract.
+
+    Returns:
+        The extracted points (N x 3 array).
+    """
+
+    return mesh.sample(n)
+
+
+def extract_points_uniform_surface_sampling(
+    mesh: trimesh.Trimesh, n: int
+) -> np.ndarray:
+    """
+    Extract N uniformly sampled points from the surface of a mesh.
+
+    Args:
+        mesh: The mesh from which to extract points.
+        n: The number of points to extract.
+
+    Returns:
+        The extracted points (N x 3 array).
+    """
+
+    return trimesh.sample.sample_surface_even(mesh=mesh, count=n)[0]
+
+
+def extract_points_select_points_over_axis(
+    mesh: trimesh.Trimesh, axis: str, direction: str, n: int
+) -> np.ndarray:
+    """
+    Extract N points from a mesh along a specified axis. The points are selected based on their position along the axis.
+
+    Args:
+        mesh: The mesh from which to extract points.
+        axis: The axis along which to extract points.
+        direction: The direction along the axis from which to extract points. Valid values are "higher" and "lower".
+        n: The number of points to extract.
+
+    Returns:
+        The extracted points (N x 3 array).
+    """
+
+    dirs = {"higher": np.s_[-n:], "lower": np.s_[:n]}
+    arr = mesh.vertices
+
+    # Sort rows lexicographically first, then columnar.
+    arr.sort(axis=0)
+    sorted_arr = arr[dirs[direction]]
+    return sorted_arr
+
+
+def extract_points_aap(
+    mesh: trimesh.Trimesh,
+    axis: str,
+    upper: float | None = None,
+    lower: float | None = None,
+) -> np.ndarray:
+    """
+    Extract points from a mesh along a specified axis within a specified range. The points are selected based on their position along the axis.
+
+    Args:
+        mesh: The mesh from which to extract points.
+        axis: The axis along which to extract points.
+        upper: The upper bound of the range.
+        lower: The lower bound of the range.
+
+    Returns:
+        The extracted points (N x 3 array).
+
+    Raises:
+        AssertionError: If the lower bound is greater than the upper bound.
+    """
+
+    # Check bounds.
+    upper = upper if upper is not None else np.inf
+    lower = lower if lower is not None else -np.inf
+    assert lower < upper, "Invalid bounds for axis-aligned plane"
+
+    # Logic.
+    points = mesh.vertices[
+        (mesh.vertices[:, VALID_AXIS[axis]] >= lower)
+        & (mesh.vertices[:, VALID_AXIS[axis]] <= upper)
+    ]
+
+    return points

jaxsim/parsers/rod/parser.py

@@ -1,13 +1,14 @@
 import dataclasses
+import os
 import pathlib
-from typing import Dict, List, NamedTuple, Optional, Union
+from typing import NamedTuple
 
 import jax.numpy as jnp
 import numpy as np
 import rod
 
 from jaxsim import logging
-from jaxsim.math.quaternion import Quaternion
+from jaxsim.math import Quaternion
 from jaxsim.parsers import descriptions, kinematic_graph
 
 from . import utils
@@ -23,16 +24,17 @@ class SDFData(NamedTuple):
     fixed_base: bool
     base_link_name: str
 
-    link_descriptions: List[descriptions.LinkDescription]
-    joint_descriptions: List[descriptions.JointDescription]
-    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 | rod.Sdf,
     model_name: str | None = None,
     is_urdf: bool | None = None,
 ) -> SDFData:
@@ -40,47 +42,53 @@ def extract_model_data(
     Extract data from an SDF/URDF resource useful to build a JaxSim model.
 
     Args:
-        model_description: A path to an SDF/URDF file, a string containing its content,
-          or a pre-parsed/pre-built rod model.
+        model_description:
+            A path to an SDF/URDF file, a string containing its content, or
+            a pre-parsed/pre-built rod model.
         model_name: The name of the model to extract from the SDF resource.
-        is_urdf: Whether the SDF resource is a URDF file. Needed only if model_description
-            is a URDF string.
+        is_urdf:
+            Whether to force parsing the resource as a URDF file. Automatically
+            detected if not provided.
 
     Returns:
         The extracted model data.
     """
 
-    if isinstance(model_description, rod.Model):
-        sdf_model = model_description
-    else:
-        # Parse the SDF resource
-        sdf_element = rod.Sdf.load(sdf=model_description, is_urdf=is_urdf)
-
-        if len(sdf_element.models()) == 0:
-            raise RuntimeError("Failed to find any model in SDF resource")
-
-        # Assume the SDF resource has only one model, or the desired model name is given
-        sdf_models = {m.name: m for m in sdf_element.models()}
-        sdf_model = (
-            sdf_element.models()[0] if len(sdf_models) == 1 else sdf_models[model_name]
-        )
+    match model_description:
+        case rod.Model():
+            sdf_model = model_description
+        case rod.Sdf() | str() | pathlib.Path():
+            sdf_element = (
+                model_description
+                if isinstance(model_description, rod.Sdf)
+                else rod.Sdf.load(sdf=model_description, is_urdf=is_urdf)
+            )
+            if not sdf_element.models():
+                raise RuntimeError("Failed to find any model in SDF resource")
+
+            # Assume the SDF resource has only one model, or the desired model name is given.
+            sdf_models = {m.name: m for m in sdf_element.models()}
+            sdf_model = (
+                sdf_element.models()[0]
+                if len(sdf_models) == 1
+                else sdf_models[model_name]
+            )
 
-    # Log model name
+    # Log model name.
     logging.debug(msg=f"Found model '{sdf_model.name}' in SDF resource")
 
     # Jaxsim supports only models compatible with URDF, i.e. those having all links
     # directly attached to their parent joint without additional roto-translations.
+    # Furthermore, the following switch also post-processes frames such that their
+    # pose is expressed wrt the parent link they are rigidly attached to.
     sdf_model.switch_frame_convention(frame_convention=rod.FrameConvention.Urdf)
 
-    # Log type of base link
+    # 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
+    # Log detected base link.
     logging.debug(msg=f"Considering '{sdf_model.get_canonical_link()}' as base link")
 
     # Pose of the model
@@ -98,11 +106,11 @@ def extract_model_data(
     # Parse links
     # ===========
 
-    # Parse the links (unconnected)
+    # Parse the links (unconnected).
     links = [
         descriptions.LinkDescription(
             name=l.name,
-            mass=jnp.float32(l.inertial.mass),
+            mass=float(l.inertial.mass),
             inertia=utils.from_sdf_inertial(inertial=l.inertial),
             pose=l.pose.transform() if l.pose is not None else np.eye(4),
         )
@@ -110,15 +118,32 @@ def extract_model_data(
         if l.inertial.mass > 0
     ]
 
-    # Create a dictionary to find easily links
-    links_dict: Dict[str, descriptions.LinkDescription] = {l.name: l for l in links}
+    # Create a dictionary to find easily links.
+    links_dict: dict[str, descriptions.LinkDescription] = {l.name: l for l in links}
+
+    # ============
+    # Parse frames
+    # ============
+
+    # Parse the frames (unconnected).
+    frames = [
+        descriptions.LinkDescription(
+            name=f.name,
+            mass=jnp.array(0.0, dtype=float),
+            inertia=jnp.zeros(shape=(3, 3)),
+            parent=links_dict[f.attached_to],
+            pose=f.pose.transform() if f.pose is not None else jnp.eye(4),
+        )
+        for f in sdf_model.frames()
+        if f.attached_to in links_dict
+    ]
 
     # =========================
     # Process fixed-base models
     # =========================
 
     # In this case, we need to get the pose of the joint that connects the base link
-    # to the world and combine their pose
+    # to the world and combine their pose.
     if sdf_model.is_fixed_base():
         # Create a massless word link
         world_link = descriptions.LinkDescription(
@@ -134,7 +159,7 @@ def extract_model_data(
                 name=j.name,
                 parent=world_link,
                 child=links_dict[j.child],
-                jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
+                jtype=utils.joint_to_joint_type(joint=j),
                 axis=(
                     np.array(j.axis.xyz.xyz)
                     if j.axis is not None
@@ -147,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}
         ]
 
@@ -159,28 +184,23 @@ def extract_model_data(
             msg = "Found more/less than one joint connecting a fixed-base model to the world"
             raise ValueError(msg + f": {[j.name for j in joints_with_world_parent]}")
 
+        base_link_name = joints_with_world_parent[0].child.name
+
         msg = "Combining the pose of base link '{}' with the pose of joint '{}'"
-        logging.info(
-            msg.format(
-                joints_with_world_parent[0].child.name, joints_with_world_parent[0].name
-            )
-        )
+        logging.info(msg.format(base_link_name, joints_with_world_parent[0].name))
 
         # Combine the pose of the base link (child of the found fixed joint)
         # with the pose of the fixed joint connecting with the world.
         # Note: we assume it's a fixed joint and ignore any joint angle.
-        links_dict[joints_with_world_parent[0].child.name].mutable(
-            validate=False
-        ).pose = (
-            joints_with_world_parent[0].pose
-            @ links_dict[joints_with_world_parent[0].child.name].pose
+        links_dict[base_link_name].mutable(validate=False).pose = (
+            joints_with_world_parent[0].pose @ links_dict[base_link_name].pose
         )
 
     # ============
     # Parse joints
     # ============
 
-    # Check that all joint poses are expressed w.r.t. their parent link
+    # Check that all joint poses are expressed w.r.t. their parent link.
     for j in sdf_model.joints():
         if j.pose is None:
             continue
@@ -195,15 +215,15 @@ def extract_model_data(
             msg = "Pose of joint '{}' is not expressed wrt its parent link '{}'"
             raise ValueError(msg.format(j.name, j.parent))
 
-    # Parse the joints
+    # Parse the joints.
     joints = [
         descriptions.JointDescription(
             name=j.name,
             parent=links_dict[j.parent],
             child=links_dict[j.child],
-            jtype=utils.axis_to_jtype(axis=j.axis, type=j.type),
+            jtype=utils.joint_to_joint_type(joint=j),
             axis=(
-                np.array(j.axis.xyz.xyz)
+                np.array(j.axis.xyz.xyz, dtype=float)
                 if j.axis is not None
                 and j.axis.xyz is not None
                 and j.axis.xyz.xyz is not None
@@ -212,56 +232,60 @@ def extract_model_data(
             pose=j.pose.transform() if j.pose is not None else np.eye(4),
             initial_position=0.0,
             position_limit=(
-                (
-                    float(j.axis.limit.lower)
-                    if j.axis is not None and j.axis.limit is not None
-                    else np.finfo(float).min
+                float(
+                    j.axis.limit.lower
+                    if j.axis is not None
+                    and j.axis.limit is not None
+                    and j.axis.limit.lower is not None
+                    else jnp.finfo(float).min
                 ),
-                (
-                    float(j.axis.limit.upper)
-                    if j.axis is not None and j.axis.limit is not None
-                    else np.finfo(float).max
+                float(
+                    j.axis.limit.upper
+                    if j.axis is not None
+                    and j.axis.limit is not None
+                    and j.axis.limit.upper is not None
+                    else jnp.finfo(float).max
                 ),
             ),
-            friction_static=(
+            friction_static=float(
                 j.axis.dynamics.friction
                 if j.axis is not None
                 and j.axis.dynamics is not None
                 and j.axis.dynamics.friction is not None
                 else 0.0
             ),
-            friction_viscous=(
+            friction_viscous=float(
                 j.axis.dynamics.damping
                 if j.axis is not None
                 and j.axis.dynamics is not None
                 and j.axis.dynamics.damping is not None
                 else 0.0
             ),
-            position_limit_damper=(
+            position_limit_damper=float(
                 j.axis.limit.dissipation
                 if j.axis is not None
                 and j.axis.limit is not None
                 and j.axis.limit.dissipation is not None
-                else 0.0
+                else os.environ.get("JAXSIM_JOINT_POSITION_LIMIT_DAMPER", 0.0)
             ),
-            position_limit_spring=(
+            position_limit_spring=float(
                 j.axis.limit.stiffness
                 if j.axis is not None
                 and j.axis.limit is not None
                 and j.axis.limit.stiffness is not None
-                else 0.0
+                else os.environ.get("JAXSIM_JOINT_POSITION_LIMIT_SPRING", 0.0)
             ),
         )
         for j in sdf_model.joints()
-        if j.type in {"revolute", "prismatic", "fixed"}
+        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
+    # Create a dictionary to find the parent joint of the links.
     joint_dict = {j.child.name: j.name for j in joints}
 
-    # Check that all the link poses are expressed wrt their parent joint
+    # Check that all the link poses are expressed wrt their parent joint.
     for l in sdf_model.links():
         if l.name not in links_dict:
             continue
@@ -284,7 +308,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():
@@ -305,10 +329,23 @@ def extract_model_data(
 
                 collisions.append(sphere_collision)
 
+            if collision.geometry.mesh is not None and int(
+                os.environ.get("JAXSIM_COLLISION_MESH_ENABLED", "0")
+            ):
+                logging.warning("Mesh collision support is still experimental.")
+                mesh_collision = utils.create_mesh_collision(
+                    collision=collision,
+                    link_description=links_dict[link.name],
+                    method=utils.meshes.extract_points_vertices,
+                )
+
+                collisions.append(mesh_collision)
+
     return SDFData(
         model_name=sdf_model.name,
         link_descriptions=links,
         joint_descriptions=joints,
+        frame_descriptions=frames,
         collision_shapes=collisions,
         fixed_base=sdf_model.is_fixed_base(),
         base_link_name=sdf_model.get_canonical_link(),
@@ -318,33 +355,39 @@ 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 = None,
 ) -> descriptions.ModelDescription:
     """
-    Builds a model description from an SDF/URDF resource.
+    Build a model description from an SDF/URDF resource.
 
     Args:
         model_description: A path to an SDF/URDF file, a string containing its content,
           or a pre-parsed/pre-built rod model.
-        is_urdf: Whether the SDF resource is a URDF file. Needed only if model_description
-            is a URDF string.
+        is_urdf: Whether the force parsing the resource as a URDF file. Automatically
+            detected if not provided.
+
     Returns:
         The parsed model description.
     """
 
-    # Parse data from the SDF assuming it contains a single model
+    # Parse data from the SDF assuming it contains a single model.
     sdf_data = extract_model_data(
         model_description=model_description, model_name=None, is_urdf=is_urdf
     )
 
-    # Build the model description.
+    # Build the intermediate representation used for building a JaxSim model.
+    # This process, beyond other operations, removes the fixed joints.
     # Note: if the model is fixed-base, the fixed joint between world and the first
     #       link is removed and the pose of the first link is updated.
-    model = descriptions.ModelDescription.build_model_from(
+    #
+    # The whole process is:
+    # URDF/SDF ⟶ rod.Model ⟶ ModelDescription ⟶ JaxSimModel.
+    graph = descriptions.ModelDescription.build_model_from(
         name=sdf_data.model_name,
         links=sdf_data.link_descriptions,
         joints=sdf_data.joint_descriptions,
+        frames=sdf_data.frame_descriptions,
         collisions=sdf_data.collision_shapes,
         fixed_base=sdf_data.fixed_base,
         base_link_name=sdf_data.base_link_name,
@@ -352,11 +395,11 @@ def build_model_description(
         considered_joints=[
             j.name
             for j in sdf_data.joint_descriptions
-            if j.jtype is not descriptions.JointType.F
+            if j.jtype is not descriptions.JointType.Fixed
         ],
     )
 
     # Store the parsed SDF tree as extra info
-    model = dataclasses.replace(model, extra_info={"sdf_model": sdf_data.sdf_model})
+    graph = dataclasses.replace(graph, _extra_info={"sdf_model": sdf_data.sdf_model})
 
-    return model
+    return graph