[feat] Equality constraints from the model, cassie examples (#35)

* feat: introduce types of barrier constraints

* feat: init Model class

* fix: fix typo

* Revert "feat: introduce types of barrier constraints"

This reverts commit 954b9d8.

* feat: renewed constraints; base, object, equality (weld, connect) constraints

* fix: make second object name required

* fix: proper jax class type

* feat: joint constraint

* fix: extra inverse in difference frame

* feat: frame constraint

* feat: position constraint

* feat: joint constraint

* fix: change joint_difference to array difference, to underline requirement on scalar joints in joint task and barrier

* feat: add active flag to the constraint

* feat: support for constraints in locak IK solver

* fix: add NotImplementedError in global IK for constraints

* fix: minor fixes based on examples

* example: manipulator constrained local ik

* hotfix: remove debug pring

* example: go2 and g1 examples with constrained contact

* fix: frame constraint fix

* wip: rework equality constraints, does not work

* fix: generalize joint barrier and task to handle multiple nonscalar joints

* hotfix: joint task and barrier user mask

* example: cassie example, closed kin does not work

* example: working cassie example

* fix: joint barrier and task fix

* fix: everything works

* example: cassie squat batched

* chore: remove all constraints except for equality constraint

* wip: dead-beat ALIP example init

* feat: ALIP planar motion demo

* feat: rotational motion in ALIP, legs only

* feat: torso orientation in ALIP example

* refactor: clean up ALIP example

* hotfix: Simeon, ne bei

* feat: proper alip implementation

* minor prints in examples

* wip: sobaki layut, karavan idet

* feat: finalize example, add visualizations

* docs: add caravan to the documentation

* chore: remove redundant example

* chore: attempt to rework README gifs

* chore: attempt to rework README gifs

* chore: attempt to rework README gifs

* Update README.md

* fix: fix joint task for new format, like joint barrier

* test: update joint task and barrier tests

* test: update configuration tests

* chore: add mujoco copyright to copied functions

* chore: udpate disclamer for collision file

* test: finalize unit tests

* docs: add documentation for constraints

* ci: attempt to fix problems with mypy

* refactor: ruff format

* hotfix: pipeline

* chore: attempt to fix tests

* chore: update cast of one variable

* chore: accidentally broke the test

---------

Co-authored-by: Simeon <[email protected]>

Author: domrachev03

.DS_Store

@@ -2,29 +2,28 @@ name: mypy
 
 on:
   push:
-    branches: [ main ]
+    branches: [main]
   pull_request:
-    branches: [ main ]
+    branches: [main]
 
 jobs:
   mypy:
     runs-on: ubuntu-latest
     steps:
-    - name: Check out the repository
-      uses: actions/checkout@v3
+      - name: Check out the repository
+        uses: actions/checkout@v3
 
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.12'
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.12"
 
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install -e .
-        pip install mypy
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install mypy
 
-    - name: Run mypy
-      continue-on-error: true
-      run: |
-        mypy mjinx
+      - name: Run mypy
+        run: |
+          mypy mjinx

.github/workflows/mypy.yml

@@ -39,7 +39,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
 
     steps:
       - name: Check out the repository

.github/workflows/tests.yml

@@ -11,9 +11,10 @@
 **Mjinx** is a library for auto-differentiable numerical inverse kinematics, powered by **JAX** and **Mujoco MJX**. The library was heavily inspired by the similar Pinocchio-based tool [pink](https://github.com/stephane-caron/pink/tree/main) and Mujoco-based analogue [mink](https://github.com/kevinzakka/mink/tree/main).
 
 <p align="center">
-  <img src="img/local_ik_output.gif" style="width: 200px" />
-  <img src="img/go2_stance.gif" style="width: 200px" /> 
-  <img src="img/g1_heart.gif" style="width: 200px"/>
+  <img src="img/local_ik_output.gif" style="width: 400px" />
+  <img src="img/go2_stance.gif" style="width: 400px" /> 
+  <img src="img/g1_heart.gif" style="width: 400px"/>
+  <img src="img/cassie_caravan.gif" style="width: 400px"/>
 </p>
 
 ## Key features

README.md

@@ -54,3 +54,16 @@ MJINX provides several barrier implementations:
    :maxdepth: 1
 
    barriers
+
+***********
+Constraints
+***********
+
+Constraints represent a new type of component with a strictly enforced equality condition :math:`f(q) = 0`. Those constraints might be either treated strictly as differentiated exponentially stable equality: :math:`\nabla h(q)^T v = -\alpha h(q)`, with :math:`\alpha` controls constraint enforcement and :math:`v` is the velocity vector, or as a soft constraint -- task with high gain.
+
+Yet, only the following constraints are implemented:
+
+.. toctree::
+   :maxdepth: 1
+
+   constraints

docs/components.rst

@@ -25,6 +25,7 @@
 
 autodoc_type_aliases = {
     "ndarray": ":data:`mjinx.typing.ndarray`",
+    "ArrayLike": ":data:`mjinx.typing.ArrayLike`",
     "ArrayOrFloat": ":data:`mjinx.typing.ArrayOrFloat`",
     "ClassKFunctions": ":data:`mjinx.typing.ClassKFunctions`",
     "CollisionBody": ":data:`mjinx.typing.CollisionBody`",

docs/conf.py

@@ -0,0 +1,31 @@
+:github_url: https://github.com/based-robotics/mjinx/tree/docs/github_pages/docs/constraints.rst
+
+.. _Constraints:
+
+Constraints represent a new type of component with a strictly enforced equality, either hardly or softly.
+
+The constratin is formulated as :math:`f(q) = 0`. Those constraints might be either treated strictly as differentiated exponentially stable equality: :math:`\nabla h(q)^T v = -\alpha h(q)`, with :math:`\alpha` controls constraint enforcement and :math:`v` is the velocity vector, or as a soft constraint -- task with high gain.
+
+Base Constraint
+---------------
+The foundational class that all constraints extend. It defines the core interface and mathematical properties for constraint objectives.
+
+.. automodule:: mjinx.components.constraints._base
+    :members:
+    :member-order: bysource
+
+Model Equality Constraint
+-------------------------
+The constraints, described in MuJoCo model as equality constrainst. 
+
+.. automodule:: mjinx.components.constraints._equality_constraint
+    :members:
+    :member-order: bysource
+
+.. Base Task
+.. ---------
+.. The foundational class that all tasks extend. It defines the core interface and mathematical properties for task objectives.
+
+.. .. automodule:: mjinx.components.tasks._base
+..     :members:
+..     :member-order: bysource

docs/constraints.rst

@@ -33,6 +33,7 @@ and `MuJoCo MJX <https://mujoco.readthedocs.io/en/stable/mjx.html>`_. The librar
      <img src="https://github.com/based-robotics/mjinx/raw/main/img/local_ik_output.gif" style="width: 200px;" alt="KUKA arm example">
      <img src="https://github.com/based-robotics/mjinx/raw/main/img/go2_stance.gif" style="width: 200px;" alt="GO2 robot example">
      <img src="https://github.com/based-robotics/mjinx/raw/main/img/g1_heart.gif" style="width: 200px;" alt="Heart path example">
+     <img src="https://github.com/based-robotics/mjinx/raw/main/img/cassie_caravan.gif" style="width: 200px;" alt="Heart path example">
    </div>
 
 *************