jaxsim 0.6.2.dev162__tar.gz → 0.6.2.dev179__tar.gz

jaxsim-0.6.2.dev179/.github/workflows/gpu_benchmark.yml

@@ -0,0 +1,104 @@
+name: GPU Benchmarks
+
+on:
+  pull_request:
+    types: [opened, reopened, synchronize]
+  workflow_dispatch:
+  schedule:
+    - cron: "0 0  * * 1" # Run At 00:00 on Monday
+
+permissions:
+  pull-requests: write
+  deployments: write
+  contents: write
+
+jobs:
+  benchmark:
+    runs-on: self-hosted
+    container:
+      image: ghcr.io/nvidia/jax:jax
+      options: --rm --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up environment
+        run: |
+          pip install .[all]
+
+      - name: Install Gazebo SDF
+        run: |
+            apt update && apt install -y lsb-release wget
+            sh -c 'echo "deb http://packages.osrfoundation.org/gazebo/ubuntu-stable `lsb_release -cs` main" > /etc/apt/sources.list.d/gazebo-stable.list'
+            wget http://packages.osrfoundation.org/gazebo.key -O - | apt-key add -
+            apt-get update
+            apt install -y libsdformat14-dev libsdformat14
+
+      - name: Get main branch SHA
+        id: get-main-branch-sha
+        run: |
+          git config --global --add safe.directory /__w/jaxsim/jaxsim
+          SHA=$(git rev-parse origin/main)
+          echo "sha=$SHA" >> $GITHUB_OUTPUT
+
+      - name: Get benchmark results from main branch
+        id: cache
+        uses: actions/cache/restore@v4
+        with:
+          path: ./cache
+          key: ${{ steps.get-main-branch-sha.outputs.sha }}-${{ runner.os }}-benchmark
+
+      - 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
+
+      - name: Compare benchmark results with main branch
+        uses: benchmark-action/github-action-benchmark@v1.20.4
+        if: steps.cache.outputs.cache-hit == 'true'
+        with:
+          tool: 'pytest'
+          output-file-path: output.json
+          external-data-json-path: ./cache/benchmark-data.json
+          save-data-file: false
+          fail-on-alert: true
+          summary-always: true
+          comment-always: true
+          alert-threshold: 150%
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Store benchmark result for main branch
+        uses: benchmark-action/github-action-benchmark@v1.20.4
+        if: ${{ github.ref_name == 'main' }}
+        with:
+          tool: 'pytest'
+          output-file-path: output.json
+          external-data-json-path: ./cache/benchmark-data.json
+          save-data-file: true
+          fail-on-alert: false
+          summary-always: true
+          comment-always: true
+          alert-threshold: 150%
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: Publish Benchmark Results to GitHub Pages
+        uses: benchmark-action/github-action-benchmark@v1.20.4
+        if: ${{ github.ref_name == 'main' }}
+        with:
+          tool: 'pytest'
+          output-file-path: output.json
+          benchmark-data-dir-path: "benchmarks"
+          fail-on-alert: false
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+          comment-on-alert: true
+          summary-always: true
+          save-data-file: true
+          alert-threshold: "150%"
+          auto-push: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
+
+      - name: Update Benchmark Results cache
+        uses: actions/cache/save@v4
+        if: ${{ github.ref_name == 'main' }}
+        with:
+          path: ./cache
+          key: ${{ steps.get-main-branch-sha.outputs.sha }}-${{ runner.os }}-benchmark

{jaxsim-0.6.2.dev162 → jaxsim-0.6.2.dev179}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jaxsim
-Version: 0.6.2.dev162
+Version: 0.6.2.dev179
 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.dev162 → jaxsim-0.6.2.dev179}/src/jaxsim/_version.py

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

{jaxsim-0.6.2.dev162 → jaxsim-0.6.2.dev179}/src/jaxsim/api/model.py

@@ -1118,6 +1118,34 @@ def forward_dynamics_crb(
     return v̇_WB, s̈
 
 
+@jax.jit
+@js.common.named_scope
+def forward_kinematics(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Matrix:
+    """
+    Compute the forward kinematics of the model.
+
+    Args:
+        model: The model to consider.
+        data: The data of the considered model.
+
+    Returns:
+        The nL x 4 x 4 array containing the stacked homogeneous transformations
+        of the links. The first axis is the link index.
+    """
+
+    W_H_LL, _ = jaxsim.rbda.forward_kinematics_model(
+        model=model,
+        base_position=data.base_position,
+        base_quaternion=data.base_quaternion,
+        joint_positions=data.joint_positions,
+        joint_velocities=data.joint_velocities,
+        base_linear_velocity_inertial=data._base_linear_velocity,
+        base_angular_velocity_inertial=data._base_angular_velocity,
+    )
+
+    return W_H_LL
+
+
 @jax.jit
 @js.common.named_scope
 def free_floating_mass_matrix(

{jaxsim-0.6.2.dev162 → jaxsim-0.6.2.dev179}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jaxsim
-Version: 0.6.2.dev162
+Version: 0.6.2.dev179
 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.dev162 → jaxsim-0.6.2.dev179}/src/jaxsim.egg-info/SOURCES.txt

@@ -14,6 +14,7 @@ setup.py
 .github/CODEOWNERS
 .github/dependabot.yml
 .github/workflows/ci_cd.yml
+.github/workflows/gpu_benchmark.yml
 .github/workflows/pixi.yml
 .github/workflows/read_the_docs.yml
 docs/Makefile

{jaxsim-0.6.2.dev162 → jaxsim-0.6.2.dev179}/tests/test_benchmark.py

@@ -104,6 +104,9 @@ def test_soft_contact_model(
 ):
     model = jaxsim_model_ergocub_reduced
 
+    with model.editable(validate=False) as model:
+        model.contact_model = jaxsim.rbda.contacts.SoftContacts()
+
     benchmark_test_function(js.ode.system_dynamics, model, benchmark, batch_size)
 
 
@@ -129,3 +132,15 @@ def test_relaxed_rigid_contact_model(
         model.contact_model = jaxsim.rbda.contacts.RelaxedRigidContacts()
 
     benchmark_test_function(js.ode.system_dynamics, model, benchmark, batch_size)
+
+
+@pytest.mark.benchmark
+def test_simulation_step(
+    jaxsim_model_ergocub_reduced: js.model.JaxSimModel, benchmark, batch_size
+):
+    model = jaxsim_model_ergocub_reduced
+
+    with model.editable(validate=False) as model:
+        model.contact_model = jaxsim.rbda.contacts.RelaxedRigidContacts()
+
+    benchmark_test_function(js.model.step, model, benchmark, batch_size)