jaxsim 0.6.2.dev182__tar.gz → 0.6.2.dev225__tar.gz

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/.github/workflows/gpu_benchmark.yml

@@ -51,7 +51,7 @@ jobs:
 
       - name: Run benchmark and store result
         run: |
-            pytest tests/test_benchmark.py -k 'not test_rigid_contact_model and not test_soft_contact_model' --gpu-only --batch-size 128 --benchmark-only --benchmark-json output.json
+            pytest tests/test_benchmark.py --gpu-only --batch-size 128 --benchmark-only --benchmark-json output.json
 
       - name: Compare benchmark results with main branch
         uses: benchmark-action/github-action-benchmark@v1.20.4

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.6.2.dev182
+Version: 0.6.2.dev225
 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>
@@ -89,6 +89,7 @@ Requires-Dist: mujoco>=3.0.0; extra == "viz"
 Requires-Dist: scipy>=1.14.0; extra == "viz"
 Provides-Extra: all
 Requires-Dist: jaxsim[style,testing,viz]; extra == "all"
+Dynamic: license-file
 
 # JaxSim
 

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/examples/jaxsim_as_physics_engine_advanced.ipynb

@@ -60,6 +60,8 @@
     "\n",
     "import os\n",
     "\n",
+    "os.environ[\"MUJOCO_GL\"] = \"osmesa\"\n",
+    "\n",
     "import jax\n",
     "\n",
     "import jax.numpy as jnp\n",
@@ -178,7 +180,6 @@
    },
    "outputs": [],
    "source": [
-    "\n",
     "# Initialize the simulated time.\n",
     "T = jnp.arange(start=0, stop=1.0, step=model.time_step)"
    ]
@@ -212,7 +213,9 @@
     "key = jax.random.PRNGKey(seed=0)\n",
     "\n",
     "# Split subkeys for sampling random initial data.\n",
-    "batch_size = 10\n",
+    "batch_size = 16\n",
+    "row_length = int(jnp.sqrt(batch_size))\n",
+    "row_dist = 0.3 * row_length\n",
     "key, *subkeys = jax.random.split(key=key, num=batch_size + 1)\n",
     "\n",
     "# Create the batched data by sampling the height from [0.5, 0.6] meters.\n",
@@ -220,12 +223,24 @@
     "    lambda key: js.data.random_model_data(\n",
     "        model=model,\n",
     "        key=key,\n",
-    "        base_pos_bounds=([0, 0, 0.3], [0, 0, 0.6]),\n",
+    "        base_pos_bounds=([0, 0, 0.3], [0, 0, 1.2]),\n",
     "        base_vel_lin_bounds=(0, 0),\n",
     "        base_vel_ang_bounds=(0, 0),\n",
     "    )\n",
     ")(jnp.vstack(subkeys))\n",
     "\n",
+    "x, y = jnp.meshgrid(\n",
+    "    jnp.linspace(-row_dist, row_dist, num=row_length),\n",
+    "    jnp.linspace(-row_dist, row_dist, num=row_length),\n",
+    ")\n",
+    "xy_coordinate = jnp.stack([x.flatten(), y.flatten()], axis=-1)\n",
+    "\n",
+    "# Reset the x and y position to a grid.\n",
+    "data_batch_t0 = data_batch_t0.replace(\n",
+    "    model=model,\n",
+    "    base_position=data_batch_t0.base_position.at[:, :2].set(xy_coordinate),\n",
+    ")\n",
+    "\n",
     "print(\"W_p_B(t0)=\\n\", data_batch_t0.base_position[0:10])"
    ]
   },
@@ -339,7 +354,7 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "\n",
-    "plt.plot(T, data_trajectory.base_position[:, 0:5, 2])\n",
+    "plt.plot(T, data_trajectory.base_position[:, :, 2])\n",
     "plt.grid(True)\n",
     "plt.xlabel(\"Time [s]\")\n",
     "plt.ylabel(\"Height [m]\")\n",
@@ -347,6 +362,74 @@
     "plt.show()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jaxsim.mujoco\n",
+    "\n",
+    "mjcf_string, assets = jaxsim.mujoco.ModelToMjcf.convert(\n",
+    "    model.built_from,\n",
+    "    cameras=jaxsim.mujoco.loaders.MujocoCamera.build_from_target_view(\n",
+    "        camera_name=\"sphere_cam\",\n",
+    "        lookat=[0, 0, 0.3],\n",
+    "        distance=4,\n",
+    "        azimuth=150,\n",
+    "        elevation=-10,\n",
+    "    ),\n",
+    ")\n",
+    "\n",
+    "# Create a helper for each parallel instance.\n",
+    "mj_model_helpers = [\n",
+    "    jaxsim.mujoco.MujocoModelHelper.build_from_xml(\n",
+    "        mjcf_description=mjcf_string, assets=assets\n",
+    "    )\n",
+    "    for _ in range(batch_size)\n",
+    "]\n",
+    "\n",
+    "# Create the video recorder.\n",
+    "recorder = jaxsim.mujoco.MujocoVideoRecorder(\n",
+    "    model=mj_model_helpers[0].model,\n",
+    "    data=[helper.data for helper in mj_model_helpers],\n",
+    "    fps=int(1 / model.time_step),\n",
+    "    width=320 * 2,\n",
+    "    height=240 * 2,\n",
+    ")\n",
+    "\n",
+    "for data_t in data_trajectory_list:\n",
+    "\n",
+    "    for helper, base_position, base_quaternion, joint_position in zip(\n",
+    "        mj_model_helpers,\n",
+    "        data_t.base_position,\n",
+    "        data_t.base_orientation,\n",
+    "        data_t.joint_positions,\n",
+    "        strict=True,\n",
+    "    ):\n",
+    "        helper.set_base_position(position=base_position)\n",
+    "        helper.set_base_orientation(orientation=base_quaternion)\n",
+    "\n",
+    "        if model.dofs() > 0:\n",
+    "            helper.set_joint_positions(\n",
+    "                positions=joint_position, joint_names=model.joint_names()\n",
+    "            )\n",
+    "\n",
+    "    # Record a new video frame.\n",
+    "    recorder.record_frame(camera_name=\"sphere_cam\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import mediapy as media\n",
+    "\n",
+    "media.show_video(recorder.frames, fps=recorder.fps)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -362,7 +445,7 @@
    "toc_visible": true
   },
   "kernelspec": {
-   "display_name": "comodo_jaxsim",
+   "display_name": "jaxpypi",
    "language": "python",
    "name": "python3"
   },
@@ -376,7 +459,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.13.1"
   }
  },
  "nbformat": 4,

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/src/jaxsim/_version.py

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

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/src/jaxsim/api/__init__.py

@@ -4,7 +4,6 @@ from . import (
     actuation_model,
     com,
     contact,
-    contact_model,
     frame,
     integrators,
     joint,

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/src/jaxsim/api/com.py

@@ -301,9 +301,7 @@ def bias_acceleration(
             C_v̇_WL = W_v̇_bias_WL = v̇_bias_WL  # noqa: F841
             C_v_WC = W_v_WW = jnp.zeros(6)  # noqa: F841
 
-            L_H_C = L_H_W = jax.vmap(  # noqa: F841
-                lambda W_H_L: jaxsim.math.Transform.inverse(W_H_L)
-            )(W_H_L)
+            L_H_C = L_H_W = jax.vmap(jaxsim.math.Transform.inverse)(W_H_L)  # noqa: F841
 
             L_v_LC = L_v_LW = jax.vmap(  # noqa: F841
                 lambda i: -js.link.velocity(

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/src/jaxsim/api/common.py

@@ -121,14 +121,8 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             The 6D quantity in the other representation.
         """
 
-        W_array = array.squeeze()
-        W_H_O = transform.squeeze()
-
-        if W_array.size != 6:
-            raise ValueError(W_array.size, 6)
-
-        if W_H_O.shape != (4, 4):
-            raise ValueError(W_H_O.shape, (4, 4))
+        W_array = array
+        W_H_O = transform
 
         match other_representation:
 
@@ -139,25 +133,24 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
 
                 if not is_force:
                     O_Xv_W = Adjoint.from_transform(transform=W_H_O, inverse=True)
-                    O_array = O_Xv_W @ W_array
+                    O_array = jnp.einsum("...ij,...j->...i", O_Xv_W, W_array)
 
                 else:
-                    O_Xf_W = Adjoint.from_transform(transform=W_H_O).T
-                    O_array = O_Xf_W @ W_array
+                    O_Xf_W = Adjoint.from_transform(transform=W_H_O).swapaxes(-1, -2)
+                    O_array = jnp.einsum("...ij,...j->...i", O_Xf_W, W_array)
 
                 return O_array
 
             case VelRepr.Mixed:
-                W_p_O = W_H_O[0:3, 3]
-                W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
+                W_H_OW = W_H_O.at[..., 0:3, 0:3].set(jnp.eye(3))
 
                 if not is_force:
                     OW_Xv_W = Adjoint.from_transform(transform=W_H_OW, inverse=True)
-                    OW_array = OW_Xv_W @ W_array
+                    OW_array = jnp.einsum("...ij,...j->...i", OW_Xv_W, W_array)
 
                 else:
-                    OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).T
-                    OW_array = OW_Xf_W @ W_array
+                    OW_Xf_W = Adjoint.from_transform(transform=W_H_OW).swapaxes(-1, -2)
+                    OW_array = jnp.einsum("...ij,...j->...i", OW_Xf_W, W_array)
 
                 return OW_array
 
@@ -188,45 +181,40 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
             The 6D quantity in the inertial-fixed representation.
         """
 
-        W_array = array.squeeze()
-        W_H_O = transform.squeeze()
-
-        if W_array.size != 6:
-            raise ValueError(W_array.size, 6)
-
-        if W_H_O.shape != (4, 4):
-            raise ValueError(W_H_O.shape, (4, 4))
+        O_array = array
+        W_H_O = transform
 
         match other_representation:
             case VelRepr.Inertial:
-                W_array = array
-                return W_array
+                return O_array
 
             case VelRepr.Body:
-                O_array = array
 
                 if not is_force:
-                    W_Xv_O: jtp.Array = Adjoint.from_transform(W_H_O)
-                    W_array = W_Xv_O @ O_array
+                    W_Xv_O = Adjoint.from_transform(W_H_O)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xv_O, O_array)
 
                 else:
-                    W_Xf_O = Adjoint.from_transform(transform=W_H_O, inverse=True).T
-                    W_array = W_Xf_O @ O_array
+                    W_Xf_O = Adjoint.from_transform(
+                        transform=W_H_O, inverse=True
+                    ).swapaxes(-1, -2)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xf_O, O_array)
 
                 return W_array
 
             case VelRepr.Mixed:
-                BW_array = array
-                W_p_O = W_H_O[0:3, 3]
-                W_H_OW = jnp.eye(4).at[0:3, 3].set(W_p_O)
+
+                W_H_OW = W_H_O.at[..., 0:3, 0:3].set(jnp.eye(3))
 
                 if not is_force:
-                    W_Xv_BW: jtp.Array = Adjoint.from_transform(W_H_OW)
-                    W_array = W_Xv_BW @ BW_array
+                    W_Xv_BW = Adjoint.from_transform(W_H_OW)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xv_BW, O_array)
 
                 else:
-                    W_Xf_BW = Adjoint.from_transform(transform=W_H_OW, inverse=True).T
-                    W_array = W_Xf_BW @ BW_array
+                    W_Xf_BW = Adjoint.from_transform(
+                        transform=W_H_OW, inverse=True
+                    ).swapaxes(-1, -2)
+                    W_array = jnp.einsum("...ij,...j->...i", W_Xf_BW, O_array)
 
                 return W_array
 

{jaxsim-0.6.2.dev182 → jaxsim-0.6.2.dev225}/src/jaxsim/api/contact.py

@@ -11,7 +11,7 @@ import jaxsim.terrain
 import jaxsim.typing as jtp
 from jaxsim import logging
 from jaxsim.math import Adjoint, Cross, Transform
-from jaxsim.rbda import contacts
+from jaxsim.rbda.contacts import SoftContacts
 
 from .common import VelRepr
 
@@ -37,14 +37,11 @@ def collidable_point_kinematics(
         the linear component of the mixed 6D frame velocity.
     """
 
-    # Switch to inertial-fixed since the RBDAs expect velocities in this representation.
-    with data.switch_velocity_representation(VelRepr.Inertial):
-
-        W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
-            model=model,
-            link_transforms=data._link_transforms,
-            link_velocities=data._link_velocities,
-        )
+    W_p_Ci, W_ṗ_Ci = jaxsim.rbda.collidable_points.collidable_points_pos_vel(
+        model=model,
+        link_transforms=data._link_transforms,
+        link_velocities=data._link_velocities,
+    )
 
     return W_p_Ci, W_ṗ_Ci
 
@@ -164,18 +161,23 @@ def estimate_good_soft_contacts_parameters(
 def estimate_good_contact_parameters(
     model: js.model.JaxSimModel,
     *,
+    standard_gravity: jtp.FloatLike = jaxsim.math.STANDARD_GRAVITY,
     static_friction_coefficient: jtp.FloatLike = 0.5,
-    **kwargs,
+    number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
+    damping_ratio: jtp.FloatLike = 1.0,
+    max_penetration: jtp.FloatLike | None = None,
 ) -> jaxsim.rbda.contacts.ContactParamsTypes:
     """
     Estimate good contact parameters.
 
     Args:
         model: The model to consider.
+        standard_gravity: The standard gravity acceleration.
         static_friction_coefficient: The static friction coefficient.
-        kwargs:
-            Additional model-specific parameters passed to the builder method of
-            the parameters class.
+        number_of_active_collidable_points_steady_state:
+            The number of active collidable points in steady state.
+        damping_ratio: The damping ratio.
+        max_penetration: The maximum penetration allowed.
 
     Returns:
         The estimated good contacts parameters.
@@ -190,20 +192,41 @@ def estimate_good_contact_parameters(
         specific application.
     """
 
-    match model.contact_model:
+    def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
+        """
+        Displacement between the CoM and the lowest collidable point using zero
+        joint positions.
+        """
+
+        zero_data = js.data.JaxSimModelData.build(
+            model=model,
+        )
+
+        W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
 
-        case contacts.RelaxedRigidContacts():
-            assert isinstance(model.contact_model, contacts.RelaxedRigidContacts)
+        if model.floating_base():
+            W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
+            return 2 * (W_pz_CoM - W_pz_C.min())
 
-            parameters = contacts.RelaxedRigidContactsParams.build(
-                mu=static_friction_coefficient,
-                **kwargs,
-            )
+        return 2 * W_pz_CoM
 
-        case _:
-            raise ValueError(f"Invalid contact model: {model.contact_model}")
+    max_δ = (
+        max_penetration
+        if max_penetration is not None
+        # Consider as default a 0.5% of the model height.
+        else 0.005 * estimate_model_height(model=model)
+    )
 
-    return parameters
+    nc = number_of_active_collidable_points_steady_state
+
+    return model.contact_model._parameters_class().build_default_from_jaxsim_model(
+        model=model,
+        standard_gravity=standard_gravity,
+        static_friction_coefficient=static_friction_coefficient,
+        max_penetration=max_δ,
+        number_of_active_collidable_points_steady_state=nc,
+        damping_ratio=damping_ratio,
+    )
 
 
 @jax.jit
@@ -244,7 +267,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
 
     # Build the link-to-point transform from the displacement between the link frame L
     # and the implicit contact frame C.
-    L_H_C = jax.vmap(lambda L_p_C: jnp.eye(4).at[0:3, 3].set(L_p_C))(L_p_Ci)
+    L_H_C = jax.vmap(jnp.eye(4).at[0:3, 3].set)(L_p_Ci)
 
     # Compose the work-to-link and link-to-point transforms.
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
@@ -504,3 +527,96 @@ def jacobian_derivative(
     )
 
     return O_J̇_WC
+
+
+@jax.jit
+@js.common.named_scope
+def link_contact_forces(
+    model: js.model.JaxSimModel,
+    data: js.data.JaxSimModelData,
+    *,
+    link_forces: jtp.MatrixLike | None = None,
+    joint_torques: jtp.VectorLike | None = None,
+) -> tuple[jtp.Matrix, dict[str, jtp.Matrix]]:
+    """
+    Compute the 6D contact forces of all links of the model in inertial representation.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+        link_forces:
+            The 6D external forces to apply to the links expressed in inertial representation
+        joint_torques:
+            The joint torques acting on the joints.
+
+    Returns:
+        A `(nL, 6)` array containing the stacked 6D contact forces of the links,
+        expressed in inertial representation.
+    """
+
+    # Compute the contact forces for each collidable point with the active contact model.
+    W_f_C, aux_dict = model.contact_model.compute_contact_forces(
+        model=model,
+        data=data,
+        **(
+            dict(link_forces=link_forces, joint_force_references=joint_torques)
+            if not isinstance(model.contact_model, SoftContacts)
+            else {}
+        ),
+    )
+
+    # Compute the 6D forces applied to the links equivalent to the forces applied
+    # to the frames associated to the collidable points.
+    W_f_L = link_forces_from_contact_forces(model=model, contact_forces=W_f_C)
+
+    return W_f_L, aux_dict
+
+
+@staticmethod
+def link_forces_from_contact_forces(
+    model: js.model.JaxSimModel,
+    *,
+    contact_forces: jtp.MatrixLike,
+) -> jtp.Matrix:
+    """
+    Compute the link forces from the contact forces.
+
+    Args:
+        model: The robot model considered by the contact model.
+        contact_forces: The contact forces computed by the contact model.
+
+    Returns:
+        The 6D contact forces applied to the links and expressed in the frame of
+        the velocity representation of data.
+    """
+
+    # Get the object storing the contact parameters of the model.
+    contact_parameters = model.kin_dyn_parameters.contact_parameters
+
+    # Extract the indices corresponding to the enabled collidable points.
+    indices_of_enabled_collidable_points = (
+        contact_parameters.indices_of_enabled_collidable_points
+    )
+
+    # Convert the contact forces to a JAX array.
+    W_f_C = jnp.atleast_2d(jnp.array(contact_forces, dtype=float).squeeze())
+
+    # Construct the vector defining the parent link index of each collidable point.
+    # We use this vector to sum the 6D forces of all collidable points rigidly
+    # attached to the same link.
+    parent_link_index_of_collidable_points = jnp.array(
+        contact_parameters.body, dtype=int
+    )[indices_of_enabled_collidable_points]
+
+    # Create the mask that associate each collidable point to their parent link.
+    # We use this mask to sum the collidable points to the right link.
+    mask = parent_link_index_of_collidable_points[:, jnp.newaxis] == jnp.arange(
+        model.number_of_links()
+    )
+
+    # Sum the forces of all collidable points rigidly attached to a body.
+    # Since the contact forces W_f_C are expressed in the world frame,
+    # we don't need any coordinate transformation.
+    W_f_L = mask.T @ W_f_C
+
+    return W_f_L