├── .github └── workflows │ └── build.yaml ├── .gitignore ├── .pre-commit-config.yaml ├── .vscode └── settings.json ├── README.md ├── _notebooks ├── README.md ├── dynamics_a1.ipynb ├── match_acceleration.ipynb ├── match_energy.ipynb ├── match_object_vel.ipynb ├── match_rne.ipynb └── skydio.ipynb ├── data ├── ltv_lqr_traj.json ├── models │ └── cart_pole │ │ ├── model.xml │ │ └── scene.xml └── trajectories │ ├── skydio.npz │ └── skydio_easy.npz ├── examples ├── mujoco_sysid_demo.ipynb └── utils.py ├── mujoco_sysid ├── __init__.py ├── convert.py ├── mjx │ ├── __init__.py │ ├── convert.py │ ├── parameters.py │ └── regressors.py ├── parameters.py ├── regressors.py └── utils.py ├── pyproject.toml └── tests ├── mjx ├── test_mjx_dynamics.py ├── test_mjx_energy.py └── test_mjx_parameters.py ├── test_dynamics.py ├── test_energy.py └── test_parameters.py /.github/workflows/build.yaml: -------------------------------------------------------------------------------- 1 | name: Upload Python Package 2 | 3 | on: 4 | push: 5 | 6 | permissions: 7 | contents: read 8 | 9 | jobs: 10 | ruff-check: 11 | runs-on: ubuntu-latest 12 | steps: 13 | - uses: actions/checkout@v3 14 | - name: Install Python 15 | uses: actions/setup-python@v4 16 | with: 17 | python-version: "3.11" 18 | - name: Install dependencies 19 | run: | 20 | python -m pip install --upgrade pip 21 | pip install ruff==0.1.11 22 | # Update output format to enable automatic inline annotations. 23 | - name: Run Ruff 24 | # run: ruff check --output-format=github . 25 | run: ruff format --diff . 26 | 27 | tests: 28 | runs-on: ubuntu-latest 29 | needs: ruff-check 30 | steps: 31 | - uses: actions/checkout@v3 32 | - name: Set up Python 33 | uses: actions/setup-python@v3 34 | with: 35 | python-version: "3.10" 36 | - name: Install dependencies 37 | run: | 38 | python -m pip install --upgrade pip 39 | pip install .[mjx_cpu] 40 | pip install pytest robot_descriptions pin 41 | - name: Run tests 42 | run: pytest 43 | 44 | deploy: 45 | runs-on: ubuntu-latest 46 | needs: ruff-check 47 | # run on tag only 48 | if: startsWith(github.ref, 'refs/tags/') 49 | steps: 50 | - uses: actions/checkout@v3 51 | - name: Set up Python 52 | uses: actions/setup-python@v3 53 | with: 54 | python-version: "3.10" 55 | - name: Install dependencies 56 | run: | 57 | python -m pip install --upgrade pip 58 | pip install build 59 | - name: Build package 60 | run: python -m build --wheel 61 | - name: Publish package 62 | uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 63 | with: 64 | user: __token__ 65 | password: ${{ secrets.PYPI_TOKEN }} 66 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | #.idea/ -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/astral-sh/ruff-pre-commit 3 | # Ruff version. 4 | rev: v0.1.11 5 | hooks: 6 | # Run the formatter. 7 | - id: ruff-format 8 | types_or: [python, pyi, jupyter] 9 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "workbench.colorCustomizations": { 3 | "activityBar.activeBackground": "#545890", 4 | "activityBar.background": "#545890", 5 | "activityBar.foreground": "#e7e7e7", 6 | "activityBar.inactiveForeground": "#e7e7e799", 7 | "activityBarBadge.background": "#b9868a", 8 | "activityBarBadge.foreground": "#15202b", 9 | "commandCenter.border": "#e7e7e799", 10 | "sash.hoverBorder": "#545890", 11 | "statusBar.background": "#414470", 12 | "statusBar.foreground": "#e7e7e7", 13 | "statusBarItem.hoverBackground": "#545890", 14 | "statusBarItem.remoteBackground": "#414470", 15 | "statusBarItem.remoteForeground": "#e7e7e7", 16 | "titleBar.activeBackground": "#414470", 17 | "titleBar.activeForeground": "#e7e7e7", 18 | "titleBar.inactiveBackground": "#41447099", 19 | "titleBar.inactiveForeground": "#e7e7e799" 20 | }, 21 | "peacock.color": "#414470" 22 | } -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # MuJoCo SysId 2 | 3 | MuJoCo System Identification (mujoco-sysid) is a Python module designed to perform system identification using the MuJoCo physics engine. This module facilitates the estimation of model parameters to match the simulated dynamics with observed data, enhancing the accuracy and reliability of robotic simulations. 4 | 5 | # Features 6 | 7 | - **Dynamic parameters**: getters and setters for dynamic parameters of MuJoCo and MJX models. 8 | - **Regressor models**: energy and dynamics-based regressors for robotic systems. 9 | - **Parameters representation**: utilities for converting between dynamic parameters, Log-Cholesky parametrization and pseudo-inertia. 10 | 11 | # Installation 12 | 13 | To install the module, run the following command: 14 | 15 | ```bash 16 | pip install mujoco-sysid 17 | ``` 18 | 19 | In order to use with MJX use `[mjx]` or `[mjx_cpu]` for CPU-only version: 20 | 21 | ```bash 22 | pip install mujoco-sysid[mjx] 23 | ``` 24 | 25 | # Documentation 26 | 27 | To be generated. 28 | 29 | # Examples 30 | 31 | 32 | 33 | The demo is inspired by recent advancements in MuJoCo utilities, specifically the [Levenberg-Marquardt nonlinear least squares method](https://colab.research.google.com/github/google-deepmind/mujoco/blob/main/python/least_squares.ipynb). 34 | 35 | Our primary focus is on mechanical systems where **the model structure is known**, including the number of state variables and the configuration of the kinematic tree. While the dynamics of these systems can be inherently complex, the general forms of the equations are known and have already been implemented in MuJoCo. In this context, the task of identification is essentially to estimate the parameters within a structured model. 36 | 37 | This repository includes the following [examples](examples/mujoco_sysid_demo.ipynb): 38 | 39 | 40 | - Estimation of cart-pole inertial parameters through random forcing and LQR stabilization of the identified system 41 | - Identification of end-effector load for the Franka Emika Panda and compensation using inverse dynamics 42 | - Determination of mass, center of mass, and spatial inertia for a Skydio X2 Quadrotor following LTV LQR tracking attempts. 43 | 44 | We hope these examples and utilities will be useful for all MuJoCo users and assist in resolving their system identification challenges. 45 | 46 | For further questions and suggestions, please do not hesitate to create an issue or start a discussion [on GitHub](https://github.com/lvjonok/mujoco-sysid/issues/new/choose). 47 | -------------------------------------------------------------------------------- /_notebooks/README.md: -------------------------------------------------------------------------------- 1 | # Notes 2 | 3 | - Notebooks in this directory were used to compare something and test the behavior 4 | - It is okay if some of them do not run or are not complete 5 | -------------------------------------------------------------------------------- /_notebooks/dynamics_a1.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import mujoco\n", 10 | "import numpy as np\n", 11 | "import pinocchio as pin\n", 12 | "from robot_descriptions.h1_description import URDF_PATH\n", 13 | "from robot_descriptions.h1_mj_description import MJCF_PATH\n", 14 | "\n", 15 | "from mujoco_sysid import parameters\n", 16 | "from mujoco_sysid.utils import muj2pin\n", 17 | "\n", 18 | "mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH)\n", 19 | "mjdata = mujoco.MjData(mjmodel)\n", 20 | "\n", 21 | "pinmodel = pin.buildModelFromUrdf(URDF_PATH, pin.JointModelFreeFlyer())" 22 | ] 23 | }, 24 | { 25 | "cell_type": "markdown", 26 | "metadata": {}, 27 | "source": [ 28 | "## Setup the model to have the same dynamic parameters" 29 | ] 30 | }, 31 | { 32 | "cell_type": "code", 33 | "execution_count": 2, 34 | "metadata": {}, 35 | "outputs": [ 36 | { 37 | "name": "stdout", 38 | "output_type": "stream", 39 | "text": [ 40 | "[ 5.39 -0.00108 0.00022 -0.24374 0.0556 0.00009 0.01927 -0.00007 0.00001 0.04902]\n", 41 | "[ 5.39 -0.00108 0.00022 -0.24374 0.0556 0.00009 0.01927 -0.00007 0.00001 0.04902]\n", 42 | "[ 2.244 -0.11047 0.00022 0.01616 0.00269 0.00002 0.0086 0.00028 0. 0.00773]\n", 43 | "[ 2.244 -0.11047 0.00022 0.01616 0.00269 0.00002 0.0086 0.00028 0. 0.00773]\n", 44 | "[ 2.232 -0.01295 -0.00712 -0.0002 0.00208 -0.00001 0.00232 0. -0.00001 0.00253]\n", 45 | "[ 2.232 -0.01295 -0.00712 -0.0002 0.00208 -0.00001 0.00232 0. -0.00001 0.00253]\n", 46 | "[ 4.152 0.03097 -0.09741 -0.34017 0.11277 0.00006 0.10968 0.00661 -0.00078 0.00852]\n", 47 | "[ 4.152 0.03097 -0.09741 -0.34017 0.11277 0.00006 0.10968 0.00661 -0.00078 0.00852]\n", 48 | "[ 1.721 -0.00234 -0.00881 -0.23819 0.04522 -0.00008 0.04548 0.00076 -0.00099 0.00211]\n", 49 | "[ 1.721 -0.00234 -0.00881 -0.23819 0.04522 -0.00008 0.04548 0.00076 -0.00099 0.00211]\n", 50 | "[ 0.446 0.02998 0.00007 -0.02006 0.00116 -0. 0.00513 0.00163 0. 0.00416]\n", 51 | "[ 0.474 0.02018 -0. -0.02117 0.00111 0. 0.00471 0.00104 -0. 0.00366]\n", 52 | "changed to [ 0.474 0.02018 -0. -0.02117 0.00111 0. 0.00471 0.00104 -0. 0.00366]\n", 53 | "[ 2.244 -0.11047 -0.00022 0.01616 0.00269 -0.00002 0.0086 0.00028 -0. 0.00773]\n", 54 | "[ 2.244 -0.11047 -0.00022 0.01616 0.00269 -0.00002 0.0086 0.00028 -0. 0.00773]\n", 55 | "[ 2.232 -0.01295 0.00712 -0.0002 0.00208 0.00001 0.00232 0. 0.00001 0.00253]\n", 56 | "[ 2.232 -0.01295 0.00712 -0.0002 0.00208 0.00001 0.00232 0. 0.00001 0.00253]\n", 57 | "[ 4.152 0.03097 0.09741 -0.34017 0.11277 -0.00006 0.10968 0.00661 0.00078 0.00852]\n", 58 | "[ 4.152 0.03097 0.09741 -0.34017 0.11277 -0.00006 0.10968 0.00661 0.00078 0.00852]\n", 59 | "[ 1.721 -0.00234 0.00881 -0.23819 0.04522 0.00008 0.04548 0.00076 0.00099 0.00211]\n", 60 | "[ 1.721 -0.00234 0.00881 -0.23819 0.04522 0.00008 0.04548 0.00076 0.00099 0.00211]\n", 61 | "[ 0.446 0.02998 -0.00007 -0.02006 0.00116 0. 0.00513 0.00163 -0. 0.00416]\n", 62 | "[ 0.474 0.02018 0. -0.02117 0.00111 -0. 0.00471 0.00104 0. 0.00366]\n", 63 | "changed to [ 0.474 0.02018 0. -0.02117 0.00111 -0. 0.00471 0.00104 0. 0.00366]\n", 64 | "[17.789 0.0087 0.04976 3.6439 1.23386 -0.00056 1.15605 0.00025 -0.01094 0.12799]\n", 65 | "[17.789 0.0087 0.04976 3.6439 1.23386 -0.00056 1.15605 0.00025 -0.01094 0.12799]\n", 66 | "[ 1.033 0.00521 0.05543 -0.01623 0.00453 -0.0003 0.00115 0.00009 0.00091 0.00397]\n", 67 | "[ 1.033 0.00521 0.05543 -0.01623 0.00453 -0.0003 0.00115 0.00009 0.00091 0.00397]\n", 68 | "[ 0.793 0.00054 0.00091 -0.0746 0.00859 0. 0.00872 -0.00002 0.00002 0.00102]\n", 69 | "[ 0.793 0.00054 0.00091 -0.0746 0.00859 0. 0.00872 -0.00002 0.00002 0.00102]\n", 70 | "[ 0.839 0.01145 0.00232 -0.13647 0.02587 -0.00004 0.02643 0.00221 0.00045 0.00083]\n", 71 | "[ 0.839 0.01145 0.00232 -0.13647 0.02587 -0.00004 0.02643 0.00221 0.00045 0.00083]\n", 72 | "[ 0.669 0.10642 -0.0001 -0.01055 0.00059 -0.00002 0.0231 0.00197 0. 0.02293]\n", 73 | "[ 0.723 0.1192 0.00009 -0.01138 0.0006 -0.00005 0.02584 0.00217 0.00001 0.02565]\n", 74 | "changed to [ 0.723 0.1192 0.00009 -0.01138 0.0006 -0.00005 0.02584 0.00217 0.00001 0.02565]\n", 75 | "[ 1.033 0.00521 -0.05543 -0.01623 0.00453 0.0003 0.00115 0.00009 -0.00091 0.00397]\n", 76 | "[ 1.033 0.00521 -0.05543 -0.01623 0.00453 0.0003 0.00115 0.00009 -0.00091 0.00397]\n", 77 | "[ 0.793 0.00054 -0.00091 -0.0746 0.00859 -0. 0.00872 -0.00002 -0.00002 0.00102]\n", 78 | "[ 0.793 0.00054 -0.00091 -0.0746 0.00859 -0. 0.00872 -0.00002 -0.00002 0.00102]\n", 79 | "[ 0.839 0.01145 -0.00232 -0.13647 0.02587 0.00004 0.02643 0.00221 -0.00045 0.00083]\n", 80 | "[ 0.839 0.01145 -0.00232 -0.13647 0.02587 0.00004 0.02643 0.00221 -0.00045 0.00083]\n", 81 | "[ 0.669 0.10642 0.0001 -0.01055 0.00059 0.00002 0.0231 0.00197 -0. 0.02293]\n", 82 | "[ 0.723 0.1192 -0.00009 -0.01138 0.0006 0.00005 0.02584 0.00217 -0.00001 0.02565]\n", 83 | "changed to [ 0.723 0.1192 -0.00009 -0.01138 0.0006 0.00005 0.02584 0.00217 -0.00001 0.02565]\n" 84 | ] 85 | } 86 | ], 87 | "source": [ 88 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 89 | " for body_id in mjmodel.jnt_bodyid:\n", 90 | " mjparams = parameters.get_dynamic_parameters(mjmodel, body_id)\n", 91 | " pinparams = pinmodel.inertias[int(body_id)].toDynamicParameters()\n", 92 | "\n", 93 | " print(mjparams)\n", 94 | " print(pinparams)\n", 95 | " if not np.allclose(mjparams, pinparams, atol=1e-6):\n", 96 | " parameters.set_dynamic_parameters(mjmodel, body_id, pinparams)\n", 97 | "\n", 98 | " mjparams = parameters.get_dynamic_parameters(mjmodel, body_id)\n", 99 | " print(\"changed to\", mjparams)" 100 | ] 101 | }, 102 | { 103 | "cell_type": "code", 104 | "execution_count": 3, 105 | "metadata": {}, 106 | "outputs": [], 107 | "source": [ 108 | "q = np.random.randn(mjmodel.nq)\n", 109 | "v = np.random.randn(mjmodel.nv)\n", 110 | "dv = np.random.randn(mjmodel.nv)\n", 111 | "tau = np.random.randn(mjmodel.nu)\n", 112 | "\n", 113 | "mjdata.qpos[:] = q\n", 114 | "mjdata.qvel[:] = v\n", 115 | "mjdata.qacc[:] = dv\n", 116 | "mjdata.ctrl[:] = tau\n", 117 | "\n", 118 | "mujoco.mj_step(mjmodel, mjdata)\n", 119 | "mujoco.mj_inverse(mjmodel, mjdata)" 120 | ] 121 | }, 122 | { 123 | "cell_type": "code", 124 | "execution_count": 4, 125 | "metadata": {}, 126 | "outputs": [], 127 | "source": [ 128 | "# setup same data in pinocchio\n", 129 | "pinq, pinv = muj2pin(mjdata.qpos, mjdata.qvel)" 130 | ] 131 | }, 132 | { 133 | "cell_type": "code", 134 | "execution_count": 5, 135 | "metadata": {}, 136 | "outputs": [ 137 | { 138 | "name": "stdout", 139 | "output_type": "stream", 140 | "text": [ 141 | "[[ -26.04972 -2.26343 951.46826 -99.99949 0. 0. 0. 0. 0. 0. ]\n", 142 | " [ 5.94784 -950.59184 -3.00354 68.78174 0. 0. 0. 0. 0. 0. ]\n", 143 | " [ 28.77278 102.78788 -67.48682 -1.14711 0. 0. 0. 0. 0. 0. ]\n", 144 | " [ 0. 0. 28.77278 -5.94784 -68.13428 -102.78788 -0.64746 -950.59184 -1.85643 0.64746]\n", 145 | " [ 0. -28.77278 0. -26.04972 1.39419 -67.48682 -101.39368 1.11632 -951.46826 -1.39419]\n", 146 | " [ 0. 5.94784 26.04972 0. -0.43821 0.74011 0.43821 -68.78174 -99.99949 -951.03005]]\n" 147 | ] 148 | } 149 | ], 150 | "source": [ 151 | "from mujoco_sysid.regressors import joint_body_regressor\n", 152 | "\n", 153 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 154 | " print(joint_body_regressor(mjmodel, mjdata, 2))" 155 | ] 156 | }, 157 | { 158 | "cell_type": "code", 159 | "execution_count": 6, 160 | "metadata": {}, 161 | "outputs": [ 162 | { 163 | "name": "stdout", 164 | "output_type": "stream", 165 | "text": [ 166 | "[[ -27.22494 -2.26343 951.46826 -99.99949 0. 0. 0. 0. 0. 0. ]\n", 167 | " [ 19.5201 -950.59184 -3.00354 68.78174 0. 0. 0. 0. 0. 0. ]\n", 168 | " [ 28.80256 102.78788 -67.48682 -1.14711 0. 0. 0. 0. 0. 0. ]\n", 169 | " [ 0. 0. 28.80256 -19.5201 -68.13428 -102.78788 -0.64746 -950.59184 -1.85643 0.64746]\n", 170 | " [ 0. -28.80256 0. -27.22494 1.39419 -67.48682 -101.39368 1.11632 -951.46826 -1.39419]\n", 171 | " [ 0. 19.5201 27.22494 0. -0.43821 0.74011 0.43821 -68.78174 -99.99949 -951.03005]]\n" 172 | ] 173 | } 174 | ], 175 | "source": [ 176 | "pindata = pin.Data(pinmodel)\n", 177 | "tau = pin.rnea(pinmodel, pindata, pinq, pinv, mjdata.qacc.copy())\n", 178 | "\n", 179 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 180 | " print(pin.jointBodyRegressor(pinmodel, pindata, 2))" 181 | ] 182 | }, 183 | { 184 | "cell_type": "code", 185 | "execution_count": null, 186 | "metadata": {}, 187 | "outputs": [], 188 | "source": [] 189 | } 190 | ], 191 | "metadata": { 192 | "kernelspec": { 193 | "display_name": "venv", 194 | "language": "python", 195 | "name": "python3" 196 | }, 197 | "language_info": { 198 | "codemirror_mode": { 199 | "name": "ipython", 200 | "version": 3 201 | }, 202 | "file_extension": ".py", 203 | "mimetype": "text/x-python", 204 | "name": "python", 205 | "nbconvert_exporter": "python", 206 | "pygments_lexer": "ipython3", 207 | "version": "3.10.12" 208 | } 209 | }, 210 | "nbformat": 4, 211 | "nbformat_minor": 2 212 | } 213 | -------------------------------------------------------------------------------- /_notebooks/match_acceleration.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import numpy as np\n", 10 | "import mujoco\n", 11 | "from mujoco_sysid import regressors\n", 12 | "from mujoco_sysid.utils import muj2pin\n", 13 | "import pinocchio as pin\n", 14 | "from mujoco_sysid.parameters import get_dynamic_parameters\n", 15 | "\n", 16 | "np.random.seed(0)" 17 | ] 18 | }, 19 | { 20 | "cell_type": "code", 21 | "execution_count": 2, 22 | "metadata": {}, 23 | "outputs": [], 24 | "source": [ 25 | "from robot_descriptions.skydio_x2_description import URDF_PATH\n", 26 | "from robot_descriptions.skydio_x2_mj_description import MJCF_PATH" 27 | ] 28 | }, 29 | { 30 | "cell_type": "code", 31 | "execution_count": 3, 32 | "metadata": {}, 33 | "outputs": [], 34 | "source": [ 35 | "pinmodel = pin.buildModelFromUrdf(URDF_PATH)\n", 36 | "pindata = pinmodel.createData()\n", 37 | "\n", 38 | "mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH)\n", 39 | "mjmodel.opt.gravity = np.zeros(3)\n", 40 | "for act_id in range(4):\n", 41 | " mjmodel.actuator(act_id).ctrlrange = np.array([-1e4, 1e4])\n", 42 | "mjdata = mujoco.MjData(mjmodel)\n", 43 | "theta = get_dynamic_parameters(mjmodel, 1)" 44 | ] 45 | }, 46 | { 47 | "cell_type": "code", 48 | "execution_count": 4, 49 | "metadata": {}, 50 | "outputs": [ 51 | { 52 | "name": "stdout", 53 | "output_type": "stream", 54 | "text": [ 55 | "Configuration\n", 56 | "Pinocchio: [0.5488135 0.71518937 0.60276338 0.40671359 0.62006595 0.42008887\n", 57 | " 0.52309427], [0.68314008 0.72550602 0.93695947 0.79172504 0.52889492 0.56804456], [0.92559664 0.07103606 0.0871293 0.0202184 0.83261985 0.77815675]\n", 58 | "Mujoco: [0.5488135 0.71518937 0.60276338 0.52309427 0.40671359 0.62006595\n", 59 | " 0.42008887], [0.891773 0.96366276 0.38344152 0.79172504 0.52889492 0.56804456], [0.92559664 0.07103606 0.0871293 0.0202184 0.83261985 0.77815675]\n" 60 | ] 61 | } 62 | ], 63 | "source": [ 64 | "q = np.random.rand(pinmodel.nq)\n", 65 | "# normalize the quaternion\n", 66 | "q[3:7] /= np.linalg.norm(q[3:7])\n", 67 | "# q = np.array([0, 0, 1, 1, 0, 0, 0])\n", 68 | "v, dv = np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv)\n", 69 | "\n", 70 | "# v[:3] *= 0 # FIXME: occasionally when we set the linear velocity to zero, the test works\n", 71 | "# otherwise the inverse dynamics does not match\n", 72 | "\n", 73 | "pinq, pinv = muj2pin(q, v)\n", 74 | "\n", 75 | "print(\"Configuration\")\n", 76 | "print(f\"Pinocchio: {pinq}, {pinv}, {dv}\")\n", 77 | "print(f\"Mujoco: {q}, {v}, {dv}\")\n", 78 | "\n", 79 | "# Selector matrix for actuators\n", 80 | "selector = np.array(\n", 81 | " [\n", 82 | " [0.0, 0.0, 0.0, 0.0],\n", 83 | " [0.0, 0.0, 0.0, 0.0],\n", 84 | " [1.0, 1.0, 1.0, 1.0],\n", 85 | " [-0.18, 0.18, 0.18, -0.18],\n", 86 | " [0.14, 0.14, -0.14, -0.14],\n", 87 | " [-0.0201, 0.0201, 0.0201, -0.0201],\n", 88 | " ]\n", 89 | ")\n", 90 | "\n", 91 | "# Random control input\n", 92 | "ctrl = np.zeros(4) # np.random.randn(4)\n", 93 | "\n", 94 | "# compute with pinocchio all terms\n", 95 | "tau = selector @ ctrl\n", 96 | "dv = pin.aba(pinmodel, pindata, pinq, pinv, tau)\n", 97 | "result = pin.rnea(pinmodel, pindata, pinq, pinv, dv)\n", 98 | "\n", 99 | "# compute with mujoco all terms\n", 100 | "mjdata.qpos[:] = q\n", 101 | "mjdata.qvel[:] = v\n", 102 | "mjdata.qacc[:] = dv\n", 103 | "mjdata.ctrl[:] = ctrl\n", 104 | "\n", 105 | "mujoco.mj_step(mjmodel, mjdata)\n", 106 | "# mujoco.mj_inverse(mjmodel, mjdata)" 107 | ] 108 | }, 109 | { 110 | "cell_type": "code", 111 | "execution_count": 5, 112 | "metadata": {}, 113 | "outputs": [ 114 | { 115 | "name": "stdout", 116 | "output_type": "stream", 117 | "text": [ 118 | "[[-0.04571 -0.6024 0.36044 0.84716 0. 0. 0. 0. 0. 0. ]\n", 119 | " [-0.03027 0.47704 -0.9495 0.56095 0. 0. 0. 0. 0. 0. ]\n", 120 | " [ 0.04892 0.05231 0.03992 -0.90656 0. 0. 0. 0. 0. 0. ]\n", 121 | " [ 0. 0. 0.04892 0.03027 -0.26052 -0.05231 -0.30044 0.47704 -0.04294 0.30044]\n", 122 | " [ 0. -0.04892 0. -0.04571 0.44974 0.03992 0.39743 -0.30415 -0.36044 -0.44974]\n", 123 | " [ 0. -0.03027 0.04571 0. -0.41874 0.3471 0.41874 -0.56095 0.84716 0.0583 ]]\n" 124 | ] 125 | } 126 | ], 127 | "source": [ 128 | "pinY = pin.jointBodyRegressor(pinmodel, pindata, 1)\n", 129 | "\n", 130 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 131 | " print(pinY)" 132 | ] 133 | }, 134 | { 135 | "cell_type": "code", 136 | "execution_count": 6, 137 | "metadata": {}, 138 | "outputs": [ 139 | { 140 | "data": { 141 | "text/plain": [ 142 | " v = -0.129148 0.32349 -0.164172\n", 143 | " w = -0.260518 0.397427 0.0582974" 144 | ] 145 | }, 146 | "execution_count": 6, 147 | "metadata": {}, 148 | "output_type": "execute_result" 149 | } 150 | ], 151 | "source": [ 152 | "# pin.forwardKinematics(pinmodel, pindata, pinq, pinv)\n", 153 | "# pin.computeAllTerms(pinmodel, pindata, pinq, pinv)\n", 154 | "# pin.aba(pinmodel, pindata, pinq, pinv, tau)\n", 155 | "# pin.getVelocity(pinmodel, pindata, 1), pin.getClassicalAcceleration(pinmodel, pindata, 1)\n", 156 | "\n", 157 | "pindata.a_gf[1]" 158 | ] 159 | }, 160 | { 161 | "cell_type": "code", 162 | "execution_count": 7, 163 | "metadata": {}, 164 | "outputs": [ 165 | { 166 | "data": { 167 | "text/plain": [ 168 | " v = -0.129148 0.32349 -0.164172\n", 169 | " w = -0.260518 0.397427 0.0582974" 170 | ] 171 | }, 172 | "execution_count": 7, 173 | "metadata": {}, 174 | "output_type": "execute_result" 175 | } 176 | ], 177 | "source": [ 178 | "pindata.a_gf[1]" 179 | ] 180 | }, 181 | { 182 | "cell_type": "code", 183 | "execution_count": 8, 184 | "metadata": {}, 185 | "outputs": [ 186 | { 187 | "data": { 188 | "text/plain": [ 189 | " v = -3.01161 9.28482 -0.978993\n", 190 | " w = 0 -0 0" 191 | ] 192 | }, 193 | "execution_count": 8, 194 | "metadata": {}, 195 | "output_type": "execute_result" 196 | } 197 | ], 198 | "source": [ 199 | "pindata.liMi[1].actInv(pindata.a_gf[0])" 200 | ] 201 | }, 202 | { 203 | "cell_type": "code", 204 | "execution_count": 9, 205 | "metadata": {}, 206 | "outputs": [ 207 | { 208 | "data": { 209 | "text/plain": [ 210 | "array([0.5488135 , 0.71518937, 0.60276338, 0.40671359, 0.62006595,\n", 211 | " 0.42008887, 0.52309427])" 212 | ] 213 | }, 214 | "execution_count": 9, 215 | "metadata": {}, 216 | "output_type": "execute_result" 217 | } 218 | ], 219 | "source": [ 220 | "pinq" 221 | ] 222 | }, 223 | { 224 | "cell_type": "code", 225 | "execution_count": 10, 226 | "metadata": {}, 227 | "outputs": [ 228 | { 229 | "data": { 230 | "text/plain": [ 231 | "array([[-0.12191288, 0.06488634, 0.99041759],\n", 232 | " [ 0.94387066, 0.31621879, 0.09546651],\n", 233 | " [-0.30699418, 0.94646471, -0.09979545]])" 234 | ] 235 | }, 236 | "execution_count": 10, 237 | "metadata": {}, 238 | "output_type": "execute_result" 239 | } 240 | ], 241 | "source": [ 242 | "mjdata.xmat[1].reshape(3, 3)" 243 | ] 244 | }, 245 | { 246 | "cell_type": "code", 247 | "execution_count": 11, 248 | "metadata": {}, 249 | "outputs": [ 250 | { 251 | "data": { 252 | "text/plain": [ 253 | " R =\n", 254 | "1 0 0\n", 255 | "0 1 0\n", 256 | "0 0 1\n", 257 | " p = 0 0 0" 258 | ] 259 | }, 260 | "execution_count": 11, 261 | "metadata": {}, 262 | "output_type": "execute_result" 263 | } 264 | ], 265 | "source": [ 266 | "pindata.oMi[1]" 267 | ] 268 | }, 269 | { 270 | "cell_type": "code", 271 | "execution_count": 12, 272 | "metadata": {}, 273 | "outputs": [ 274 | { 275 | "data": { 276 | "text/plain": [ 277 | " R =\n", 278 | " -0.121913 0.0648863 0.990418\n", 279 | " 0.943871 0.316219 0.0954665\n", 280 | " -0.306994 0.946465 -0.0997955\n", 281 | " p = 0.548814 0.715189 0.602763" 282 | ] 283 | }, 284 | "execution_count": 12, 285 | "metadata": {}, 286 | "output_type": "execute_result" 287 | } 288 | ], 289 | "source": [ 290 | "pindata.liMi[1]" 291 | ] 292 | }, 293 | { 294 | "cell_type": "code", 295 | "execution_count": 13, 296 | "metadata": {}, 297 | "outputs": [ 298 | { 299 | "data": { 300 | "text/plain": [ 301 | "array([-0.04280117, -0.06304948, -0.01692062, -0.27890477, 0.39128987,\n", 302 | " 0.05578023])" 303 | ] 304 | }, 305 | "execution_count": 13, 306 | "metadata": {}, 307 | "output_type": "execute_result" 308 | } 309 | ], 310 | "source": [ 311 | "# compute with mujoco all terms\n", 312 | "mjdata.qpos[:] = q\n", 313 | "mjdata.qvel[:] = v\n", 314 | "mjdata.qacc[:] = dv\n", 315 | "mjdata.ctrl[:] = ctrl\n", 316 | "\n", 317 | "# position\n", 318 | "mujoco.mj_kinematics(mjmodel, mjdata)\n", 319 | "mujoco.mj_comPos(mjmodel, mjdata)\n", 320 | "mujoco.mj_crb(mjmodel, mjdata)\n", 321 | "mujoco.mj_factorM(mjmodel, mjdata)\n", 322 | "\n", 323 | "# velocity\n", 324 | "mujoco.mj_fwdVelocity(mjmodel, mjdata)\n", 325 | "mujoco.mj_comVel(mjmodel, mjdata)\n", 326 | "mujoco.mj_referenceConstraint(mjmodel, mjdata)\n", 327 | "# mujoco.mj_rne(mjmodel, mjdata)\n", 328 | "\n", 329 | "mujoco.mj_fwdActuation(mjmodel, mjdata)\n", 330 | "mujoco.mj_fwdAcceleration(mjmodel, mjdata)\n", 331 | "\n", 332 | "mjdata.qacc_smooth" 333 | ] 334 | }, 335 | { 336 | "cell_type": "code", 337 | "execution_count": 14, 338 | "metadata": {}, 339 | "outputs": [ 340 | { 341 | "data": { 342 | "text/plain": [ 343 | "array([-0.2605178 , 0.39742672, 0.05829744, -9.05989038, -1.87550827,\n", 344 | " 1.91802153])" 345 | ] 346 | }, 347 | "execution_count": 14, 348 | "metadata": {}, 349 | "output_type": "execute_result" 350 | } 351 | ], 352 | "source": [ 353 | "accel = np.zeros(6)\n", 354 | "\n", 355 | "mujoco.mj_rnePostConstraint(mjmodel, mjdata)\n", 356 | "mujoco.mj_objectAcceleration(mjmodel, mjdata, 2, 1, accel, 1)\n", 357 | "\n", 358 | "accel" 359 | ] 360 | }, 361 | { 362 | "cell_type": "code", 363 | "execution_count": 15, 364 | "metadata": {}, 365 | "outputs": [ 366 | { 367 | "data": { 368 | "text/plain": [ 369 | " v = -0.129148 0.32349 -0.164172\n", 370 | " w = -0.260518 0.397427 0.0582974" 371 | ] 372 | }, 373 | "execution_count": 15, 374 | "metadata": {}, 375 | "output_type": "execute_result" 376 | } 377 | ], 378 | "source": [ 379 | "pindata.a_gf[1]" 380 | ] 381 | }, 382 | { 383 | "cell_type": "code", 384 | "execution_count": 16, 385 | "metadata": {}, 386 | "outputs": [ 387 | { 388 | "data": { 389 | "text/plain": [ 390 | "array([-0.2605178 , 0.39742672, 0.05829744, -9.05989038, -1.87550827,\n", 391 | " 1.91802153])" 392 | ] 393 | }, 394 | "execution_count": 16, 395 | "metadata": {}, 396 | "output_type": "execute_result" 397 | } 398 | ], 399 | "source": [ 400 | "velocity = np.zeros(6)\n", 401 | "\n", 402 | "mujoco.mj_objectVelocity(mjmodel, mjdata, 2, 1, velocity, 1)\n", 403 | "\n", 404 | "v, w = velocity[3:], velocity[:3]\n", 405 | "\n", 406 | "# accel[3:] -= np.cross(w, v)\n", 407 | "\n", 408 | "accel" 409 | ] 410 | }, 411 | { 412 | "cell_type": "code", 413 | "execution_count": 17, 414 | "metadata": {}, 415 | "outputs": [ 416 | { 417 | "data": { 418 | "text/plain": [ 419 | " v = 2.88246 -8.96133 0.814821\n", 420 | " w = -0.260518 0.397427 0.0582974" 421 | ] 422 | }, 423 | "execution_count": 17, 424 | "metadata": {}, 425 | "output_type": "execute_result" 426 | } 427 | ], 428 | "source": [ 429 | "# try to reverse all the values\n", 430 | "pindata.a_gf[1] - pindata.liMi[1].actInv(pindata.a_gf[0])" 431 | ] 432 | }, 433 | { 434 | "cell_type": "code", 435 | "execution_count": 18, 436 | "metadata": {}, 437 | "outputs": [ 438 | { 439 | "data": { 440 | "text/plain": [ 441 | " v = 0.68314 0.725506 0.936959\n", 442 | " w = 0.791725 0.528895 0.568045" 443 | ] 444 | }, 445 | "execution_count": 18, 446 | "metadata": {}, 447 | "output_type": "execute_result" 448 | } 449 | ], 450 | "source": [ 451 | "pindata.v[1]" 452 | ] 453 | }, 454 | { 455 | "cell_type": "code", 456 | "execution_count": 19, 457 | "metadata": {}, 458 | "outputs": [ 459 | { 460 | "data": { 461 | "text/plain": [ 462 | "(array([-0.2605178 , 0.39742672, 0.05829744]),\n", 463 | " array([-0.2605178 , 0.39742672, 0.05829744]))" 464 | ] 465 | }, 466 | "execution_count": 19, 467 | "metadata": {}, 468 | "output_type": "execute_result" 469 | } 470 | ], 471 | "source": [ 472 | "# check that angular part matches\n", 473 | "accel[:3], pindata.a_gf[1].angular" 474 | ] 475 | }, 476 | { 477 | "cell_type": "code", 478 | "execution_count": 20, 479 | "metadata": {}, 480 | "outputs": [ 481 | { 482 | "name": "stdout", 483 | "output_type": "stream", 484 | "text": [ 485 | "[-9.05989 -1.87551 1.91802]\n", 486 | "[-0.12915 0.32349 -0.16417]\n", 487 | "[-0.0428 -0.06305 -0.01692 -0.2789 0.39129 0.05578]\n", 488 | "[-0.05883 -0.05083 -0.01213 -0.00865 0.00794 0.00396]\n" 489 | ] 490 | } 491 | ], 492 | "source": [ 493 | "# check that linear part matches\n", 494 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 495 | " print(accel[3:])\n", 496 | " print(pindata.a_gf[1].linear)\n", 497 | " print(mjdata.qacc_smooth)\n", 498 | " print(mjdata.qfrc_smooth)" 499 | ] 500 | }, 501 | { 502 | "cell_type": "code", 503 | "execution_count": 21, 504 | "metadata": {}, 505 | "outputs": [ 506 | { 507 | "name": "stdout", 508 | "output_type": "stream", 509 | "text": [ 510 | "[[-9.05989 -0.6024 0.36044 0.84716 0. 0. 0. 0. 0. 0. ]\n", 511 | " [-1.87551 0.47704 -0.9495 0.56095 0. 0. 0. 0. 0. 0. ]\n", 512 | " [ 1.91802 0.05231 0.03992 -0.90656 0. 0. 0. 0. 0. 0. ]\n", 513 | " [ 0. 0. 1.91802 1.87551 -0.26052 -0.05231 -0.30044 0.47704 -0.04294 0.30044]\n", 514 | " [ 0. -1.91802 0. -9.05989 0.44974 0.03992 0.39743 -0.30415 -0.36044 -0.44974]\n", 515 | " [ 0. -1.87551 9.05989 0. -0.41874 0.3471 0.41874 -0.56095 0.84716 0.0583 ]]\n" 516 | ] 517 | } 518 | ], 519 | "source": [ 520 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 521 | " print(regressors.joint_body_regressor(mjmodel, mjdata, 1))" 522 | ] 523 | }, 524 | { 525 | "cell_type": "code", 526 | "execution_count": null, 527 | "metadata": {}, 528 | "outputs": [], 529 | "source": [] 530 | }, 531 | { 532 | "cell_type": "code", 533 | "execution_count": 22, 534 | "metadata": {}, 535 | "outputs": [ 536 | { 537 | "data": { 538 | "text/plain": [ 539 | "(False, 16.262269739443305)" 540 | ] 541 | }, 542 | "execution_count": 22, 543 | "metadata": {}, 544 | "output_type": "execute_result" 545 | } 546 | ], 547 | "source": [ 548 | "(\n", 549 | " np.allclose(regressors.joint_body_regressor(mjmodel, mjdata, 1), pinY),\n", 550 | " np.linalg.norm(regressors.joint_body_regressor(mjmodel, mjdata, 1) - pinY),\n", 551 | ")" 552 | ] 553 | }, 554 | { 555 | "cell_type": "code", 556 | "execution_count": 23, 557 | "metadata": {}, 558 | "outputs": [ 559 | { 560 | "data": { 561 | "text/plain": [ 562 | "(array([0.68314008, 0.72550602, 0.93695947]),\n", 563 | " array([0.79172504, 0.52889492, 0.56804456]),\n", 564 | " array([-9.05989038, -1.87550827, 1.91802153]),\n", 565 | " array([-0.2605178 , 0.39742672, 0.05829744]))" 566 | ] 567 | }, 568 | "execution_count": 23, 569 | "metadata": {}, 570 | "output_type": "execute_result" 571 | } 572 | ], 573 | "source": [ 574 | "velocity = np.zeros(6)\n", 575 | "accel = np.zeros(6)\n", 576 | "_cross = np.zeros(3)\n", 577 | "\n", 578 | "mujoco.mj_objectVelocity(mjmodel, mjdata, 2, 1, velocity, 1)\n", 579 | "mujoco.mj_rnePostConstraint(mjmodel, mjdata)\n", 580 | "mujoco.mj_objectAcceleration(mjmodel, mjdata, 2, 1, accel, 1)\n", 581 | "\n", 582 | "v, w = velocity[3:], velocity[:3]\n", 583 | "# dv - classical acceleration, already contains g\n", 584 | "dv, dw = accel[3:], accel[:3]\n", 585 | "mujoco.mju_cross(_cross, w, v)\n", 586 | "\n", 587 | "# dv -= _cross\n", 588 | "\n", 589 | "v, w, dv, dw" 590 | ] 591 | }, 592 | { 593 | "cell_type": "code", 594 | "execution_count": 24, 595 | "metadata": {}, 596 | "outputs": [ 597 | { 598 | "name": "stdout", 599 | "output_type": "stream", 600 | "text": [ 601 | "Joint space forces\n", 602 | "Pinocchio: f = 2.77556e-16 6.66134e-16 -8.32667e-17\n", 603 | "tau = -3.46945e-17 -8.67362e-18 -4.0766e-17\n", 604 | " [ 2.77555756e-16 6.66133815e-16 -8.32667268e-17 -3.46944695e-17\n", 605 | " -8.67361738e-18 -4.07660017e-17]\n", 606 | "rotation [[-0.12191288 0.06488634 0.99041759]\n", 607 | " [ 0.94387066 0.31621879 0.09546651]\n", 608 | " [-0.30699418 0.94646471 -0.09979545]]\n", 609 | "Mujoco: [0. 0. 0. 0. 0. 0.] [0. 0. 0. 0. 0. 0.] [0. 0. 0. 0. 0. 0.]\n" 610 | ] 611 | } 612 | ], 613 | "source": [ 614 | "print(\"Joint space forces\")\n", 615 | "print(f\"Pinocchio: {pindata.f[1]} {result}\")\n", 616 | "\n", 617 | "rot = mjdata.xmat[1].reshape(3, 3)\n", 618 | "print(\"rotation\", rot)\n", 619 | "mujtau = mjdata.qfrc_actuator.copy()\n", 620 | "\n", 621 | "mujtau[:3] = rot.T @ mujtau[:3]\n", 622 | "print(f\"Mujoco: {mjdata.qfrc_actuator} {mujtau} {mjdata.qfrc_inverse}\")" 623 | ] 624 | }, 625 | { 626 | "cell_type": "code", 627 | "execution_count": null, 628 | "metadata": {}, 629 | "outputs": [], 630 | "source": [] 631 | } 632 | ], 633 | "metadata": { 634 | "kernelspec": { 635 | "display_name": "venv", 636 | "language": "python", 637 | "name": "python3" 638 | }, 639 | "language_info": { 640 | "codemirror_mode": { 641 | "name": "ipython", 642 | "version": 3 643 | }, 644 | "file_extension": ".py", 645 | "mimetype": "text/x-python", 646 | "name": "python", 647 | "nbconvert_exporter": "python", 648 | "pygments_lexer": "ipython3", 649 | "version": "3.10.12" 650 | } 651 | }, 652 | "nbformat": 4, 653 | "nbformat_minor": 2 654 | } 655 | -------------------------------------------------------------------------------- /_notebooks/match_energy.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import mujoco\n", 10 | "import numpy as np\n", 11 | "import pinocchio as pin\n", 12 | "\n", 13 | "# from robot_descriptions.skydio_x2_description import URDF_PATH\n", 14 | "# from robot_descriptions.skydio_x2_mj_description import MJCF_PATH\n", 15 | "from robot_descriptions.z1_description import URDF_PATH\n", 16 | "from robot_descriptions.z1_mj_description import MJCF_PATH\n", 17 | "\n", 18 | "from mujoco_sysid import parameters, regressors\n", 19 | "from mujoco_sysid.utils import muj2pin" 20 | ] 21 | }, 22 | { 23 | "cell_type": "code", 24 | "execution_count": null, 25 | "metadata": {}, 26 | "outputs": [], 27 | "source": [ 28 | "mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH)\n", 29 | "# enable energy\n", 30 | "mjmodel.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_ENERGY\n", 31 | "# disable friction, contact and limits\n", 32 | "mjmodel.opt.disableflags |= (\n", 33 | " mujoco.mjtDisableBit.mjDSBL_CONTACT | mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS | mujoco.mjtDisableBit.mjDSBL_LIMIT\n", 34 | ")\n", 35 | "mjdata = mujoco.MjData(mjmodel)\n", 36 | "\n", 37 | "pinmodel = pin.buildModelFromUrdf(URDF_PATH)\n", 38 | "pindata = pin.Data(pinmodel)" 39 | ] 40 | }, 41 | { 42 | "cell_type": "code", 43 | "execution_count": null, 44 | "metadata": {}, 45 | "outputs": [], 46 | "source": [ 47 | "# np.random.seed(10)" 48 | ] 49 | }, 50 | { 51 | "cell_type": "code", 52 | "execution_count": null, 53 | "metadata": {}, 54 | "outputs": [], 55 | "source": [ 56 | "q, v = np.random.randn(mjmodel.nq), np.zeros(mjmodel.nv)\n", 57 | "\n", 58 | "\n", 59 | "pinq, pinv = muj2pin(q, v)" 60 | ] 61 | }, 62 | { 63 | "cell_type": "code", 64 | "execution_count": null, 65 | "metadata": {}, 66 | "outputs": [], 67 | "source": [ 68 | "mjdata.qpos[:] = q.copy()\n", 69 | "mjdata.qvel[:] = v.copy()\n", 70 | "\n", 71 | "mujoco.mj_step(mjmodel, mjdata)\n", 72 | "\n", 73 | "mj_en = mjdata.energy.copy()\n", 74 | "mj_en[0] += regressors.potential_energy_bias(mjmodel)\n", 75 | "\n", 76 | "mj_en, np.sum(mj_en)" 77 | ] 78 | }, 79 | { 80 | "cell_type": "code", 81 | "execution_count": null, 82 | "metadata": {}, 83 | "outputs": [], 84 | "source": [ 85 | "(\n", 86 | " np.sum(\n", 87 | " [\n", 88 | " pin.computePotentialEnergy(pinmodel, pindata, pinq),\n", 89 | " pin.computeKineticEnergy(pinmodel, pindata, pinq, pinv),\n", 90 | " ]\n", 91 | " ),\n", 92 | " [\n", 93 | " pin.computePotentialEnergy(pinmodel, pindata, pinq),\n", 94 | " pin.computeKineticEnergy(pinmodel, pindata, pinq, pinv),\n", 95 | " ],\n", 96 | ")" 97 | ] 98 | }, 99 | { 100 | "cell_type": "code", 101 | "execution_count": null, 102 | "metadata": {}, 103 | "outputs": [], 104 | "source": [ 105 | "theta = np.concatenate([parameters.get_dynamic_parameters(mjmodel, i) for i in mjmodel.jnt_bodyid])\n", 106 | "\n", 107 | "theta.shape, theta[:10]" 108 | ] 109 | }, 110 | { 111 | "cell_type": "code", 112 | "execution_count": null, 113 | "metadata": {}, 114 | "outputs": [], 115 | "source": [ 116 | "params = []\n", 117 | "\n", 118 | "for i in range(len(pinmodel.inertias) - 1):\n", 119 | " params.extend(pinmodel.inertias[i + 1].toDynamicParameters())\n", 120 | "\n", 121 | " last_params = np.array(params[-10:])\n", 122 | " last_theta = theta[i * 10 : (i + 1) * 10]\n", 123 | "\n", 124 | " # mass should match\n", 125 | " assert np.isclose(last_params[0], last_theta[0])\n", 126 | "\n", 127 | " # lever arm should match\n", 128 | " assert np.allclose(last_params[1:4], last_theta[1:4])\n", 129 | "\n", 130 | " print(f\"for body {i} norm of difference is {np.linalg.norm(last_params - last_theta)}\")\n", 131 | "\n", 132 | "params = np.array(params)\n", 133 | "\n", 134 | "params.shape, np.linalg.norm(params - theta)" 135 | ] 136 | }, 137 | { 138 | "cell_type": "code", 139 | "execution_count": null, 140 | "metadata": {}, 141 | "outputs": [], 142 | "source": [ 143 | "reg_en = regressors.mj_energyRegressor(mjmodel, mjdata)[2]\n", 144 | "\n", 145 | "reg_en @ theta" 146 | ] 147 | } 148 | ], 149 | "metadata": { 150 | "kernelspec": { 151 | "display_name": "venv", 152 | "language": "python", 153 | "name": "python3" 154 | }, 155 | "language_info": { 156 | "codemirror_mode": { 157 | "name": "ipython", 158 | "version": 3 159 | }, 160 | "file_extension": ".py", 161 | "mimetype": "text/x-python", 162 | "name": "python", 163 | "nbconvert_exporter": "python", 164 | "pygments_lexer": "ipython3", 165 | "version": "3.10.12" 166 | } 167 | }, 168 | "nbformat": 4, 169 | "nbformat_minor": 2 170 | } 171 | -------------------------------------------------------------------------------- /_notebooks/match_object_vel.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [ 8 | { 9 | "name": "stderr", 10 | "output_type": "stream", 11 | "text": [ 12 | "2024-07-03 16:08:21.341305: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.4 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n", 13 | "/home/lvjonok/github.com/lvjonok/mujoco-sysid/venv/lib/python3.10/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast\n", 14 | " return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))\n" 15 | ] 16 | } 17 | ], 18 | "source": [ 19 | "import jax\n", 20 | "import jax.numpy as jnp\n", 21 | "import mujoco\n", 22 | "import numpy as np\n", 23 | "from mujoco import mjx\n", 24 | "from robot_descriptions.z1_mj_description import MJCF_PATH\n", 25 | "\n", 26 | "key = jax.random.PRNGKey(0)\n", 27 | "\n", 28 | "mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH)\n", 29 | "mjdata = mujoco.MjData(mjmodel)\n", 30 | "\n", 31 | "# alter the model so it becomes mjx compatible\n", 32 | "mjmodel.dof_frictionloss = 0\n", 33 | "mjmodel.opt.integrator = 0\n", 34 | "\n", 35 | "mjxmodel = mjx.put_model(mjmodel)\n", 36 | "mjxdata = mjx.put_data(mjmodel, mjdata)" 37 | ] 38 | }, 39 | { 40 | "cell_type": "code", 41 | "execution_count": 2, 42 | "metadata": {}, 43 | "outputs": [ 44 | { 45 | "data": { 46 | "text/plain": [ 47 | "'/home/lvjonok/.cache/robot_descriptions/mujoco_menagerie/unitree_z1/z1.xml'" 48 | ] 49 | }, 50 | "execution_count": 2, 51 | "metadata": {}, 52 | "output_type": "execute_result" 53 | } 54 | ], 55 | "source": [ 56 | "MJCF_PATH" 57 | ] 58 | }, 59 | { 60 | "cell_type": "code", 61 | "execution_count": 3, 62 | "metadata": {}, 63 | "outputs": [ 64 | { 65 | "data": { 66 | "text/plain": [ 67 | "(Array([-0.87181777, -0.13588488, -1.765594 , 0.22557412, -1.3263583 ,\n", 68 | " -0.6441258 ], dtype=float32),\n", 69 | " Array([ 0.95800805, -0.03642227, -1.8372649 , -1.254374 , -1.2776358 ,\n", 70 | " -0.23177852], dtype=float32),\n", 71 | " Array([-0.2456366 , 0.3126293 , 1.0859151 , 0.05576393, -0.42242417,\n", 72 | " -1.1261328 ], dtype=float32))" 73 | ] 74 | }, 75 | "execution_count": 3, 76 | "metadata": {}, 77 | "output_type": "execute_result" 78 | } 79 | ], 80 | "source": [ 81 | "q, v, dv = jax.random.normal(key, (3, mjmodel.nq))\n", 82 | "\n", 83 | "q, v, dv" 84 | ] 85 | }, 86 | { 87 | "cell_type": "code", 88 | "execution_count": 4, 89 | "metadata": {}, 90 | "outputs": [ 91 | { 92 | "name": "stdout", 93 | "output_type": "stream", 94 | "text": [ 95 | "bodyid: 2, v: [0. 0. 0.], w: [0. 0. 0.95800805]\n", 96 | "bodyid: 2, a_lin: [0. 0. 9.81], a_ang: [ 0. 0. 133.51350487]\n", 97 | "bodyid: 3, v: [0. 0. 0.], w: [ 0.12977856 -0.03642227 0.94917699]\n", 98 | "bodyid: 3, a_lin: [1.32893213 0. 9.71956992], a_ang: [ 17.55922354 264.69706853 132.3548755 ]\n", 99 | "bodyid: 4, v: [-0.01250669 -0.33221195 0.00246757], w: [ 0.90610408 -1.87368716 -0.31105437]\n", 100 | "bodyid: 4, a_lin: [100.0634086 -46.32255204 -21.43708861], a_ang: [126.31661808 137.17039253 -43.24327451]\n", 101 | "bodyid: 5, v: [-0.20819569 -0.45166973 0.37383589], w: [ 0.95272099 -3.12806119 -0.10050942]\n", 102 | "bodyid: 5, a_lin: [120.10952403 -63.28652317 -24.05620638], a_ang: [ 131.53577472 -955.73404544 -25.7699149 ]\n", 103 | "bodyid: 6, v: [ 0.39468402 -0.31301854 0.59280017], w: [ 3.26564387 0.16737446 -1.37814524]\n", 104 | "bodyid: 6, a_lin: [88.94263641 95.88811455 42.83847378], a_ang: [ 960.71640941 -134.10758379 9345.2909092 ]\n", 105 | "bodyid: 7, v: [ 0.39468402 -0.65554712 0.23874882], w: [ 3.03386535 0.96141254 -1.00149111]\n", 106 | "bodyid: 7, a_lin: [ 88.84781332 421.20436758 390.63651541], a_ang: [163211.18571526 -5745.70976954 7366.63249779]\n" 107 | ] 108 | } 109 | ], 110 | "source": [ 111 | "mjdata.qpos[:] = q\n", 112 | "mjdata.qvel[:] = v\n", 113 | "mjdata.qacc[:] = dv\n", 114 | "\n", 115 | "mujoco.mj_step(mjmodel, mjdata)\n", 116 | "mujoco.mj_rnePostConstraint(mjmodel, mjdata)\n", 117 | "\n", 118 | "velocity = np.zeros(6)\n", 119 | "for bodyid in mjmodel.jnt_bodyid:\n", 120 | " mujoco.mj_objectVelocity(mjmodel, mjdata, 2, bodyid, velocity, 1)\n", 121 | "\n", 122 | " print(f\"bodyid: {bodyid}, v: {velocity[3:]}, w: {velocity[:3]}\")\n", 123 | " mujoco.mj_objectAcceleration(mjmodel, mjdata, 2, bodyid, velocity, 1)\n", 124 | " print(f\"bodyid: {bodyid}, a_lin: {velocity[3:]}, a_ang: {velocity[:3]}\")" 125 | ] 126 | }, 127 | { 128 | "cell_type": "code", 129 | "execution_count": 5, 130 | "metadata": {}, 131 | "outputs": [ 132 | { 133 | "name": "stdout", 134 | "output_type": "stream", 135 | "text": [ 136 | "[[ 0. 0. 0. 0. 0. 0. ]\n", 137 | " [ 0. 0. 0. 0. 0. 0. ]\n", 138 | " [ 0. 0. 0.958 -0.1736 -0.1554 0. ]\n", 139 | " [-0.0279 -0.0234 0.958 -0.1748 -0.154 -0.0089]\n", 140 | " [-1.4343 -1.2056 0.958 -0.2899 -0.017 0.1817]\n", 141 | " [-2.3945 -2.0127 0.958 -0.217 -0.1037 0.4682]\n", 142 | " [-1.577 -2.9853 1.0921 -0.0366 0.0465 0.4585]\n", 143 | " [-1.4011 -2.8451 1.0363 -0.0694 0.0665 0.4055]]\n" 144 | ] 145 | } 146 | ], 147 | "source": [ 148 | "with np.printoptions(precision=4, suppress=True):\n", 149 | " print(mjdata.cvel)" 150 | ] 151 | }, 152 | { 153 | "cell_type": "code", 154 | "execution_count": 6, 155 | "metadata": {}, 156 | "outputs": [ 157 | { 158 | "name": "stdout", 159 | "output_type": "stream", 160 | "text": [ 161 | "[[ 0. 0. 0. -0. -0. 9.81 ]\n", 162 | " [ 0. 0. 0. 0. 0. 9.81 ]\n", 163 | " [ 0. 0. 133.5135 -24.1929 -21.6563 9.81 ]\n", 164 | " [ 202.283 170.7232 133.5135 -15.662 -31.7642 74.1558]\n", 165 | " [ 104.9311 88.347 133.5135 -23.6876 -22.2796 87.376 ]\n", 166 | " [ -724.0039 -624.0078 133.5135 38.6625 -94.7821 341.256 ]\n", 167 | " [ -6725.2352 6501.8936 -878.6433 -1286.669 -1201.2067 409.665 ]\n", 168 | " [-129888.0899 -91650.9038 38136.1708 21636.3522 -15194.9754 37568.2828]]\n" 169 | ] 170 | } 171 | ], 172 | "source": [ 173 | "with np.printoptions(precision=4, suppress=True, linewidth=500):\n", 174 | " print(mjdata.cacc)" 175 | ] 176 | }, 177 | { 178 | "cell_type": "code", 179 | "execution_count": 11, 180 | "metadata": {}, 181 | "outputs": [], 182 | "source": [ 183 | "from mujoco_sysid.mjx.regressors import object_velocity, com_acc\n", 184 | "\n", 185 | "mjxdata = mjxdata.replace(qpos=q, qvel=v, qacc=dv)\n", 186 | "mjxdata = mjx.step(mjxmodel, mjxdata)\n", 187 | "mjxdata = mjx.rne(mjxmodel, mjxdata)" 188 | ] 189 | }, 190 | { 191 | "cell_type": "code", 192 | "execution_count": 12, 193 | "metadata": {}, 194 | "outputs": [ 195 | { 196 | "name": "stdout", 197 | "output_type": "stream", 198 | "text": [ 199 | "[[ 0. 0. 0. 0. 0. 0. ]\n", 200 | " [ 0. 0. 0. 0. 0. 0. ]\n", 201 | " [ 0. 0. 0.958 -0.1736 -0.1554 0. ]\n", 202 | " [-0.0279 -0.0234 0.958 -0.1748 -0.154 -0.0089]\n", 203 | " [-1.4343 -1.2056 0.958 -0.2899 -0.017 0.1817]\n", 204 | " [-2.3945 -2.0127 0.958 -0.217 -0.1037 0.4682]\n", 205 | " [-1.577 -2.9853 1.0921 -0.0366 0.0465 0.4585]\n", 206 | " [-1.4011 -2.8451 1.0363 -0.0694 0.0665 0.4055]]\n" 207 | ] 208 | } 209 | ], 210 | "source": [ 211 | "with np.printoptions(precision=4, suppress=True):\n", 212 | " print(mjxdata.cvel)" 213 | ] 214 | }, 215 | { 216 | "cell_type": "code", 217 | "execution_count": 13, 218 | "metadata": {}, 219 | "outputs": [ 220 | { 221 | "name": "stdout", 222 | "output_type": "stream", 223 | "text": [ 224 | "[[ 0. 0. 0. 0. 0. 9.81 ]\n", 225 | " [ 0. 0. 0. 0. 0. 9.81 ]\n", 226 | " [ 0. 0. 0. 0. 0. 9.81 ]\n", 227 | " [ -0.3425 0.4075 0. 0.0204 0.0171 9.814 ]\n", 228 | " [ -0.0728 0.0866 0. -0.0144 -0.0049 9.8101]\n", 229 | " [ 7.6103 -9.054 -0. -1.1641 3.0386 14.0812]\n", 230 | " [ 2.7813 -17.1085 -28.992 -3.4752 -1.0579 11.8998]\n", 231 | " [ 1.2405 -29.0586 -63.8849 -11.4691 3.6247 28.2652]]\n" 232 | ] 233 | } 234 | ], 235 | "source": [ 236 | "with np.printoptions(precision=4, suppress=True, linewidth=500):\n", 237 | " print(com_acc(mjxmodel, mjxdata))" 238 | ] 239 | }, 240 | { 241 | "cell_type": "code", 242 | "execution_count": null, 243 | "metadata": {}, 244 | "outputs": [ 245 | { 246 | "name": "stdout", 247 | "output_type": "stream", 248 | "text": [ 249 | "bodyid: 2, v: [0. 0. 0.], w: [0. 0. 0.95800805]\n" 250 | ] 251 | }, 252 | { 253 | "ename": "NameError", 254 | "evalue": "name 'object_acceleration' is not defined", 255 | "output_type": "error", 256 | "traceback": [ 257 | "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", 258 | "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", 259 | "Cell \u001b[0;32mIn[9], line 6\u001b[0m\n\u001b[1;32m 2\u001b[0m velocity \u001b[38;5;241m=\u001b[39m object_velocity(mjxmodel, mjxdata, bodyid)\n\u001b[1;32m 4\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbodyid: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mbodyid\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, v: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mvelocity[\u001b[38;5;241m3\u001b[39m:]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, w: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mvelocity[:\u001b[38;5;241m3\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 6\u001b[0m acceleration \u001b[38;5;241m=\u001b[39m \u001b[43mobject_acceleration\u001b[49m(mjxmodel, mjxdata, bodyid)\n\u001b[1;32m 8\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbodyid: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mbodyid\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, a_lin: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00macceleration[\u001b[38;5;241m3\u001b[39m:]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, a_ang: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00macceleration[:\u001b[38;5;241m3\u001b[39m]\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n", 260 | "\u001b[0;31mNameError\u001b[0m: name 'object_acceleration' is not defined" 261 | ] 262 | } 263 | ], 264 | "source": [ 265 | "for bodyid in mjmodel.jnt_bodyid:\n", 266 | " velocity = object_velocity(mjxmodel, mjxdata, bodyid)\n", 267 | "\n", 268 | " print(f\"bodyid: {bodyid}, v: {velocity[3:]}, w: {velocity[:3]}\")\n", 269 | "\n", 270 | " acceleration = object_acceleration(mjxmodel, mjxdata, bodyid)\n", 271 | "\n", 272 | " print(f\"bodyid: {bodyid}, a_lin: {acceleration[3:]}, a_ang: {acceleration[:3]}\")" 273 | ] 274 | } 275 | ], 276 | "metadata": { 277 | "kernelspec": { 278 | "display_name": "venv", 279 | "language": "python", 280 | "name": "python3" 281 | }, 282 | "language_info": { 283 | "codemirror_mode": { 284 | "name": "ipython", 285 | "version": 3 286 | }, 287 | "file_extension": ".py", 288 | "mimetype": "text/x-python", 289 | "name": "python", 290 | "nbconvert_exporter": "python", 291 | "pygments_lexer": "ipython3", 292 | "version": "3.10.12" 293 | } 294 | }, 295 | "nbformat": 4, 296 | "nbformat_minor": 2 297 | } 298 | -------------------------------------------------------------------------------- /_notebooks/match_rne.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [ 8 | { 9 | "name": "stderr", 10 | "output_type": "stream", 11 | "text": [ 12 | "2024-07-04 10:15:46.022376: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.4 which is older than the ptxas CUDA version (12.5.40). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.\n" 13 | ] 14 | } 15 | ], 16 | "source": [ 17 | "import jax\n", 18 | "import jax.numpy as jnp\n", 19 | "import jax.typing as jpt\n", 20 | "import mujoco as mj\n", 21 | "import mujoco.mjx as mjx\n", 22 | "from mujoco.mjx._src import scan\n", 23 | "from mujoco.mjx._src.math import transform_motion\n", 24 | "from mujoco.mjx._src.types import DisableBit\n", 25 | "from robot_descriptions.z1_mj_description import MJCF_PATH\n", 26 | "\n", 27 | "jnp.set_printoptions(precision=5, suppress=True, linewidth=500)\n", 28 | "\n", 29 | "key = jax.random.PRNGKey(0)" 30 | ] 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": 2, 35 | "metadata": {}, 36 | "outputs": [ 37 | { 38 | "name": "stderr", 39 | "output_type": "stream", 40 | "text": [ 41 | "/home/lvjonok/github.com/lvjonok/mujoco-sysid/venv/lib/python3.10/site-packages/jax/_src/interpreters/xla.py:155: RuntimeWarning: overflow encountered in cast\n", 42 | " return np.asarray(x, dtypes.canonicalize_dtype(x.dtype))\n" 43 | ] 44 | } 45 | ], 46 | "source": [ 47 | "mjmodel = mj.MjModel.from_xml_path(MJCF_PATH)\n", 48 | "mjdata = mj.MjData(mjmodel)\n", 49 | "\n", 50 | "# alter the model so it becomes mjx compatible\n", 51 | "mjmodel.dof_frictionloss = 0\n", 52 | "mjmodel.opt.integrator = 1\n", 53 | "\n", 54 | "mjxmodel = mjx.put_model(mjmodel)\n", 55 | "mjxdata = mjx.put_data(mjmodel, mjdata)" 56 | ] 57 | }, 58 | { 59 | "cell_type": "code", 60 | "execution_count": 3, 61 | "metadata": {}, 62 | "outputs": [], 63 | "source": [ 64 | "q, v, dv = jax.random.normal(key, (3, mjmodel.nq))" 65 | ] 66 | }, 67 | { 68 | "cell_type": "markdown", 69 | "metadata": {}, 70 | "source": [ 71 | "mj_inverse + mj_rnePostConstraint" 72 | ] 73 | }, 74 | { 75 | "cell_type": "code", 76 | "execution_count": 4, 77 | "metadata": {}, 78 | "outputs": [ 79 | { 80 | "data": { 81 | "text/plain": [ 82 | "array([[0., 0., 0., 0., 0., 0.],\n", 83 | " [0., 0., 0., 0., 0., 0.],\n", 84 | " [0., 0., 0., 0., 0., 0.],\n", 85 | " [0., 0., 0., 0., 0., 0.],\n", 86 | " [0., 0., 0., 0., 0., 0.],\n", 87 | " [0., 0., 0., 0., 0., 0.],\n", 88 | " [0., 0., 0., 0., 0., 0.],\n", 89 | " [0., 0., 0., 0., 0., 0.]])" 90 | ] 91 | }, 92 | "execution_count": 4, 93 | "metadata": {}, 94 | "output_type": "execute_result" 95 | } 96 | ], 97 | "source": [ 98 | "mjdata.qpos = q\n", 99 | "mjdata.qvel = v\n", 100 | "mjdata.qacc = dv\n", 101 | "\n", 102 | "mj.mj_inverse(mjmodel, mjdata)\n", 103 | "mjdata.cacc" 104 | ] 105 | }, 106 | { 107 | "cell_type": "code", 108 | "execution_count": 5, 109 | "metadata": {}, 110 | "outputs": [ 111 | { 112 | "data": { 113 | "text/plain": [ 114 | "array([[ 0. , 0. , 0. , -0. , -0. , 9.81 ],\n", 115 | " [ 0. , 0. , 0. , 0. , 0. , 9.81 ],\n", 116 | " [ 0. , 0. , -0.24564, 0.04451, 0.03984, 9.81 ],\n", 117 | " [ 0.26177, 0.17445, -0.24564, 0.05323, 0.02676, 9.88573],\n", 118 | " [ 2.22556, -0.4742 , -0.24564, -0.02487, -0.14671, 9.75662],\n", 119 | " [ 3.04146, -1.35822, -0.24564, -0.14381, 0.16344, 10.17372],\n", 120 | " [ 3.97377, -0.57563, 3.77315, 0.23265, 0.7747 , 10.4696 ],\n", 121 | " [ 4.84202, 0.20971, 3.80627, 0.14318, 0.83105, 10.06915]])" 122 | ] 123 | }, 124 | "execution_count": 5, 125 | "metadata": {}, 126 | "output_type": "execute_result" 127 | } 128 | ], 129 | "source": [ 130 | "mj.mj_rnePostConstraint(mjmodel, mjdata)\n", 131 | "\n", 132 | "mjdata.cacc" 133 | ] 134 | }, 135 | { 136 | "cell_type": "code", 137 | "execution_count": 6, 138 | "metadata": {}, 139 | "outputs": [ 140 | { 141 | "data": { 142 | "text/plain": [ 143 | "(array([[ 0. , 0. , 1. , -0.1812 , -0.1622 , 0. ],\n", 144 | " [ 0.7655 , 0.64344, 0. , 0.03215, -0.03825, 0.24308],\n", 145 | " [ 0.7655 , 0.64344, 0. , 0.06266, -0.07455, -0.1037 ],\n", 146 | " [ 0.7655 , 0.64344, 0. , -0.0581 , 0.06912, -0.22839],\n", 147 | " [-0.63989, 0.76128, -0.10492, -0.14118, -0.11763, 0.00753],\n", 148 | " [-0.75908, -0.60487, 0.24068, 0.14133, -0.08628, 0.22892]]),\n", 149 | " array([[ 0. , 0. , 0. , 0. , 0. , 0. ],\n", 150 | " [-0.61642, 0.73335, 0. , 0.03664, 0.0308 , 0.00726],\n", 151 | " [-0.61642, 0.73335, 0. , 0.07954, 0.05036, 0.00898],\n", 152 | " [-0.61642, 0.73335, 0. , 0.09224, -0.24418, -0.34267],\n", 153 | " [-0.51814, -0.86424, -3.1108 , -0.24798, -0.43955, -0.23407],\n", 154 | " [-0.05795, -0.44941, -1.31224, -0.30063, 0.1761 , 0.61546]]))" 155 | ] 156 | }, 157 | "execution_count": 6, 158 | "metadata": {}, 159 | "output_type": "execute_result" 160 | } 161 | ], 162 | "source": [ 163 | "mjdata.cdof, mjdata.cdof_dot" 164 | ] 165 | }, 166 | { 167 | "cell_type": "markdown", 168 | "metadata": {}, 169 | "source": [ 170 | "mjx_rne + mjx_rnePostConstraint" 171 | ] 172 | }, 173 | { 174 | "cell_type": "code", 175 | "execution_count": 7, 176 | "metadata": {}, 177 | "outputs": [], 178 | "source": [ 179 | "mjxdata = mjxdata.replace(qpos=q, qvel=v, qacc=dv)\n", 180 | "\n", 181 | "\n", 182 | "def mjx_invPosition(m: mjx.Model, d: mjx.Data) -> mjx.Data:\n", 183 | " d = mjx.kinematics(m, d)\n", 184 | " d = mjx.com_pos(m, d)\n", 185 | " d = mjx.camlight(m, d)\n", 186 | " # flex is missing\n", 187 | " # tendon is missing\n", 188 | "\n", 189 | " d = mjx.crb(m, d)\n", 190 | " d = mjx.factor_m(m, d)\n", 191 | " d = mjx.collision(m, d)\n", 192 | " d = mjx.make_constraint(m, d)\n", 193 | " d = mjx.transmission(m, d)\n", 194 | "\n", 195 | " return d\n", 196 | "\n", 197 | "\n", 198 | "def mjx_invVelocity(m: mjx.Model, d: mjx.Data) -> mjx.Data:\n", 199 | " d = mjx.fwd_velocity(m, d)\n", 200 | "\n", 201 | " return d\n", 202 | "\n", 203 | "\n", 204 | "def mjx_invConstraint(m: mjx.Model, d: mjx.Data) -> mjx.Data:\n", 205 | " # return data if there are no constraints\n", 206 | " if d.nefc == 0:\n", 207 | " return d\n", 208 | "\n", 209 | " # jar = Jac*qacc - aref\n", 210 | " jar = d.efc_J @ d.qacc - d.efc_aref" 211 | ] 212 | }, 213 | { 214 | "cell_type": "code", 215 | "execution_count": null, 216 | "metadata": {}, 217 | "outputs": [], 218 | "source": [ 219 | "def mjx_inverse(m: mjx.Model, d: mjx.Data) -> mjx.Data:\n", 220 | " d = mjx_invPosition(m, d)\n", 221 | " d = mjx_invVelocity(m, d)\n", 222 | "\n", 223 | " # acceleration dependent\n", 224 | " mjx\n", 225 | "\n", 226 | " d = mjx.rne(m, d)\n", 227 | "\n", 228 | " return d\n", 229 | "\n", 230 | "\n", 231 | "mjxdata = mjx.rne(mjxmodel, mjxdata)" 232 | ] 233 | }, 234 | { 235 | "cell_type": "code", 236 | "execution_count": 8, 237 | "metadata": {}, 238 | "outputs": [ 239 | { 240 | "data": { 241 | "text/plain": [ 242 | "Array([[0. , 0. , 0. , 0. , 0. , 9.81],\n", 243 | " [0. , 0. , 0. , 0. , 0. , 9.81],\n", 244 | " [0. , 0. , 0. , 0. , 0. , 9.81],\n", 245 | " [0. , 0. , 0. , 0. , 0. , 9.81],\n", 246 | " [0. , 0. , 0. , 0. , 0. , 9.81],\n", 247 | " [0. , 0. , 0. , 0. , 0. , 9.81],\n", 248 | " [0. , 0. , 0. , 0. , 0. , 9.81],\n", 249 | " [0. , 0. , 0. , 0. , 0. , 9.81]], dtype=float32)" 250 | ] 251 | }, 252 | "execution_count": 8, 253 | "metadata": {}, 254 | "output_type": "execute_result" 255 | } 256 | ], 257 | "source": [ 258 | "def com_acc(m: mjx.Model, d: mjx.Data) -> jpt.ArrayLike:\n", 259 | " # forward scan over tree: accumulate link center of mass acceleration\n", 260 | " def cacc_fn(cacc, cdof_dot, qvel):\n", 261 | " if cacc is None:\n", 262 | " if m.opt.disableflags & DisableBit.GRAVITY:\n", 263 | " cacc = jnp.zeros((6,))\n", 264 | " else:\n", 265 | " cacc = jnp.concatenate((jnp.zeros((3,)), -m.opt.gravity))\n", 266 | "\n", 267 | " cacc += jnp.sum(jax.vmap(jnp.multiply)(cdof_dot, qvel), axis=0)\n", 268 | "\n", 269 | " return cacc\n", 270 | "\n", 271 | " return scan.body_tree(m, cacc_fn, \"vv\", \"b\", d.cdof_dot, d.qvel)\n", 272 | "\n", 273 | "\n", 274 | "com_acc(mjxmodel, mjxdata)" 275 | ] 276 | }, 277 | { 278 | "cell_type": "code", 279 | "execution_count": 9, 280 | "metadata": {}, 281 | "outputs": [ 282 | { 283 | "data": { 284 | "text/plain": [ 285 | "(Array([[0., 0., 0., 0., 0., 0.],\n", 286 | " [0., 0., 0., 0., 0., 0.],\n", 287 | " [0., 0., 0., 0., 0., 0.],\n", 288 | " [0., 0., 0., 0., 0., 0.],\n", 289 | " [0., 0., 0., 0., 0., 0.],\n", 290 | " [0., 0., 0., 0., 0., 0.]], dtype=float32),\n", 291 | " Array([[0., 0., 0., 0., 0., 0.],\n", 292 | " [0., 0., 0., 0., 0., 0.],\n", 293 | " [0., 0., 0., 0., 0., 0.],\n", 294 | " [0., 0., 0., 0., 0., 0.],\n", 295 | " [0., 0., 0., 0., 0., 0.],\n", 296 | " [0., 0., 0., 0., 0., 0.]], dtype=float32))" 297 | ] 298 | }, 299 | "execution_count": 9, 300 | "metadata": {}, 301 | "output_type": "execute_result" 302 | } 303 | ], 304 | "source": [ 305 | "mjxdata.cdof, mjxdata.cdof_dot" 306 | ] 307 | }, 308 | { 309 | "cell_type": "code", 310 | "execution_count": 10, 311 | "metadata": {}, 312 | "outputs": [ 313 | { 314 | "data": { 315 | "text/plain": [ 316 | "Array([[ 0. , 0. , 0. , -0. , -0. , 9.81],\n", 317 | " [ 0. , 0. , 0. , -0. , -0. , 9.81],\n", 318 | " [ 0. , 0. , 0. , 0. , 0. , 9.81],\n", 319 | " [ 0. , 0. , 0. , 0. , 0. , 9.81],\n", 320 | " [ 0. , 0. , 0. , 0. , 0. , 9.81],\n", 321 | " [ 0. , 0. , 0. , 0. , 0. , 9.81],\n", 322 | " [ 0. , 0. , 0. , 0. , 0. , 9.81],\n", 323 | " [ 0. , 0. , 0. , 0. , 0. , 9.81]], dtype=float32)" 324 | ] 325 | }, 326 | "execution_count": 10, 327 | "metadata": {}, 328 | "output_type": "execute_result" 329 | } 330 | ], 331 | "source": [ 332 | "def mjx_rnePostConstraint(m: mjx.Model, d: mjx.Data):\n", 333 | " nbody = m.nbody\n", 334 | "\n", 335 | " all_cacc = jnp.zeros((nbody, 6))\n", 336 | "\n", 337 | " # clear cacc, set world acceleration to -gravity\n", 338 | " if not m.opt.disableflags & DisableBit.GRAVITY:\n", 339 | " all_cacc = all_cacc.at[0, 3:].set(-m.opt.gravity)\n", 340 | "\n", 341 | " # FIXME: assumption that xfrc_applied is zero\n", 342 | " # FIXME: assumption that contacts are zero\n", 343 | " # FIXME: assumption that connect and weld constraints are zero\n", 344 | "\n", 345 | " # forward pass over bodies: compute acc\n", 346 | " for j in range(nbody):\n", 347 | " bda = m.body_dofadr[j]\n", 348 | "\n", 349 | " # cacc = cacc_parent + cdofdot * qvel + cdof * qacc\n", 350 | " cacc_j = all_cacc[m.body_parentid[j]] + d.cdof_dot[bda] * d.qvel[bda] + d.cdof[bda] * d.qacc[bda]\n", 351 | " all_cacc = all_cacc.at[j].set(cacc_j)\n", 352 | "\n", 353 | " return all_cacc\n", 354 | "\n", 355 | "\n", 356 | "mjx_rnePostConstraint(mjxmodel, mjxdata)" 357 | ] 358 | }, 359 | { 360 | "cell_type": "code", 361 | "execution_count": null, 362 | "metadata": {}, 363 | "outputs": [], 364 | "source": [] 365 | } 366 | ], 367 | "metadata": { 368 | "kernelspec": { 369 | "display_name": "venv", 370 | "language": "python", 371 | "name": "python3" 372 | }, 373 | "language_info": { 374 | "codemirror_mode": { 375 | "name": "ipython", 376 | "version": 3 377 | }, 378 | "file_extension": ".py", 379 | "mimetype": "text/x-python", 380 | "name": "python", 381 | "nbconvert_exporter": "python", 382 | "pygments_lexer": "ipython3", 383 | "version": "3.10.12" 384 | } 385 | }, 386 | "nbformat": 4, 387 | "nbformat_minor": 2 388 | } 389 | -------------------------------------------------------------------------------- /_notebooks/skydio.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 1, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "import mujoco\n", 10 | "import numpy as np\n", 11 | "import pinocchio as pin\n", 12 | "from robot_descriptions.skydio_x2_description import URDF_PATH\n", 13 | "from robot_descriptions.skydio_x2_mj_description import MJCF_PATH\n", 14 | "\n", 15 | "from mujoco_sysid import parameters\n", 16 | "from mujoco_sysid.utils import muj2pin\n", 17 | "\n", 18 | "model = mujoco.MjModel.from_xml_path(MJCF_PATH)\n", 19 | "mjdata = mujoco.MjData(model)\n", 20 | "\n", 21 | "for act_id in range(4):\n", 22 | " model.actuator(act_id).ctrlrange = np.array([-1e4, 1e4])\n", 23 | "\n", 24 | "model.opt.timestep = 1e-3" 25 | ] 26 | }, 27 | { 28 | "cell_type": "code", 29 | "execution_count": 2, 30 | "metadata": {}, 31 | "outputs": [ 32 | { 33 | "name": "stdout", 34 | "output_type": "stream", 35 | "text": [ 36 | "[ 1.325 0. 0. 0.0715 0.04051 0. 0.02927 -0.0021\n", 37 | " 0. 0.06053]\n" 38 | ] 39 | } 40 | ], 41 | "source": [ 42 | "sparceM = np.zeros((6, 6))\n", 43 | "\n", 44 | "mujoco.mj_step(model, mjdata)\n", 45 | "mujoco.mj_fullM(model, sparceM, mjdata.qM)\n", 46 | "\n", 47 | "mSrc = sparceM[3:6, :3]\n", 48 | "mrc = np.array([mSrc[2, 1], mSrc[0, 2], mSrc[1, 0]])\n", 49 | "Ib = sparceM[3:6, 3:6]\n", 50 | "vechIb = np.array([Ib[0, 0], Ib[0, 1], Ib[1, 1], Ib[0, 2], Ib[1, 2], Ib[2, 2]])\n", 51 | "skydio_parameters = np.array([sparceM[0, 0], *mrc, *vechIb])\n", 52 | "\n", 53 | "with np.printoptions(precision=5, suppress=True):\n", 54 | " print(skydio_parameters)" 55 | ] 56 | }, 57 | { 58 | "cell_type": "code", 59 | "execution_count": 3, 60 | "metadata": {}, 61 | "outputs": [ 62 | { 63 | "name": "stdout", 64 | "output_type": "stream", 65 | "text": [ 66 | "[ 1.325 0. 0. 0.0715 0.04051 0. 0.02927 -0.0021\n", 67 | " 0. 0.06053]\n" 68 | ] 69 | } 70 | ], 71 | "source": [ 72 | "with np.printoptions(precision=5, suppress=True):\n", 73 | " print(parameters.get_dynamic_parameters(model, 1))" 74 | ] 75 | }, 76 | { 77 | "cell_type": "code", 78 | "execution_count": 4, 79 | "metadata": {}, 80 | "outputs": [ 81 | { 82 | "name": "stdout", 83 | "output_type": "stream", 84 | "text": [ 85 | "[ 1.325 0. 0. 0.0715 0.04051 0. 0.02927 -0.0021\n", 86 | " 0. 0.06053]\n" 87 | ] 88 | } 89 | ], 90 | "source": [ 91 | "pinmodel = pin.buildModelFromUrdf(URDF_PATH)\n", 92 | "\n", 93 | "with np.printoptions(precision=5, suppress=True):\n", 94 | " print(pinmodel.inertias[1].toDynamicParameters())" 95 | ] 96 | }, 97 | { 98 | "cell_type": "code", 99 | "execution_count": 5, 100 | "metadata": {}, 101 | "outputs": [], 102 | "source": [ 103 | "def mj_bodyRegressor(mj_model: mujoco.MjModel, mj_data: mujoco.MjData, body_id: int) -> np.ndarray:\n", 104 | " \"\"\"\n", 105 | " Computes the body regressor matrix for a specific body in a MuJoCo model.\n", 106 | "\n", 107 | " Args:\n", 108 | " mj_model (mujoco.MjModel): MuJoCo model.\n", 109 | " mj_data (mujoco.MjData): MuJoCo data.\n", 110 | " body_id (int): ID of the body for which to compute the regressor matrix.\n", 111 | "\n", 112 | " Returns:\n", 113 | " np.ndarray: Body regressor matrix of shape (6, 10).\n", 114 | " \"\"\"\n", 115 | " velocity = np.zeros(6)\n", 116 | " accel = np.zeros(6)\n", 117 | " _cross = np.zeros(3)\n", 118 | "\n", 119 | " mujoco.mj_objectVelocity(mj_model, mj_data, 2, body_id, velocity, 1)\n", 120 | " mujoco.mj_rnePostConstraint(mj_model, mj_data)\n", 121 | " mujoco.mj_objectAcceleration(mj_model, mj_data, 2, body_id, accel, 1)\n", 122 | "\n", 123 | " v, w = velocity[3:], velocity[:3]\n", 124 | " dv, dw = accel[3:], accel[:3] # dv - classical acceleration, already contains g\n", 125 | " mujoco.mju_cross(_cross, w, v)\n", 126 | " # dv -= _cross\n", 127 | " print(f\"regressor dv: {dv} dw: {dw}\")\n", 128 | "\n", 129 | " v1, v2, v3 = v\n", 130 | " v4, v5, v6 = w\n", 131 | "\n", 132 | " a1, a2, a3 = dv\n", 133 | " a4, a5, a6 = dw\n", 134 | "\n", 135 | " # fmt: off\n", 136 | " Y = np.array(\n", 137 | " [\n", 138 | " [a1 - v2 * v6 + v3 * v5, -(v5**2) - v6**2, -a6 + v4 * v5, a5 + v4 * v6, 0, 0, 0, 0, 0, 0],\n", 139 | " [a2 + v1 * v6 - v3 * v4, a6 + v4 * v5, -(v4**2) - v6**2, -a4 + v5 * v6, 0, 0, 0, 0, 0, 0],\n", 140 | " [a3 - v1 * v5 + v2 * v4, -a5 + v4 * v6, a4 + v5 * v6, -(v4**2) - v5**2, 0, 0, 0, 0, 0, 0],\n", 141 | " [0, 0, a3 - v1 * v5 + v2 * v4, -a2 - v1 * v6 + v3 * v4, a4, a5 - v4 * v6, -v5 * v6, a6 + v4 * v5, v5**2 - v6**2, v5 * v6],\n", 142 | " [0, -a3 + v1 * v5 - v2 * v4, 0, a1 - v2 * v6 + v3 * v5, v4 * v6, a4 + v5 * v6, a5, -(v4**2) + v6**2, a6 - v4 * v5, -v4 * v6],\n", 143 | " [0, a2 + v1 * v6 - v3 * v4, -a1 + v2 * v6 - v3 * v5, 0, -v4 * v5, v4**2 - v5**2, v4 * v5, a4 - v5 * v6, a5 + v4 * v6, a6],\n", 144 | " ]\n", 145 | " )\n", 146 | " # fmt: on\n", 147 | "\n", 148 | " return Y" 149 | ] 150 | }, 151 | { 152 | "cell_type": "code", 153 | "execution_count": 6, 154 | "metadata": {}, 155 | "outputs": [], 156 | "source": [ 157 | "import json\n", 158 | "\n", 159 | "with open(\"../data/ltv_lqr_traj.json\") as f:\n", 160 | " data = json.load(f)" 161 | ] 162 | }, 163 | { 164 | "cell_type": "code", 165 | "execution_count": 7, 166 | "metadata": {}, 167 | "outputs": [], 168 | "source": [ 169 | "np.random.seed(0)\n", 170 | "\n", 171 | "# traj_idx = 2\n", 172 | "\n", 173 | "# q = data[\"q\"][traj_idx]\n", 174 | "# v = data[\"v\"][traj_idx]\n", 175 | "# a = data[\"dv\"][traj_idx]\n", 176 | "# ctrl = data[\"u\"][traj_idx]\n", 177 | "\n", 178 | "q = np.random.randn(model.nq)\n", 179 | "q[3:7] = q[3:7] / np.linalg.norm(q[3:7])\n", 180 | "v = np.random.randn(model.nv)\n", 181 | "a = np.random.randn(model.nv)\n", 182 | "ctrl = np.random.randn(model.nu)\n", 183 | "\n", 184 | "mjdata = mujoco.MjData(model)\n", 185 | "mjdata.qpos[:] = q\n", 186 | "mjdata.qvel[:] = v\n", 187 | "mjdata.qacc[:] = a\n", 188 | "# mjdata.ctrl[:] = ctrl\n", 189 | "\n", 190 | "# mujoco.mj_step(model, mjdata)\n", 191 | "# mujoco.mj_forward(model, mjdata)\n", 192 | "mujoco.mj_inverse(model, mjdata)" 193 | ] 194 | }, 195 | { 196 | "cell_type": "code", 197 | "execution_count": 8, 198 | "metadata": {}, 199 | "outputs": [ 200 | { 201 | "name": "stdout", 202 | "output_type": "stream", 203 | "text": [ 204 | "regressor dv: [7.86184879 6.34771817 1.00187616] dw: [ 1.49407907 -0.20515826 0.3130677 ]\n", 205 | "[[ 7.84074 -2.69409 -0.10359 -0.09554 0. 0. 0. 0. 0. 0. ]\n", 206 | " [ 6.48332 0.52255 -0.59993 -0.38732 0. 0. 0. 0. 0. 0. ]\n", 207 | " [ 0.74675 0.31478 2.60084 -2.13566 0. 0. 0. 0. 0. 0. ]\n", 208 | " [ 0. 0. 0.74675 -6.48332 1.49408 -0.31478 -1.10676 0.52255 1.53573 1.10676]\n", 209 | " [ 0. -0.74675 0. 7.84074 0.10962 2.60084 -0.20516 0.55843 0.10359 -0.10962]\n", 210 | " [ 0. 6.48332 -7.84074 0. -0.20948 -2.09416 0.20948 0.38732 -0.09554 0.31307]]\n" 211 | ] 212 | } 213 | ], 214 | "source": [ 215 | "regressor = mj_bodyRegressor(model, mjdata, 1)\n", 216 | "\n", 217 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 218 | " print(regressor)" 219 | ] 220 | }, 221 | { 222 | "cell_type": "code", 223 | "execution_count": 9, 224 | "metadata": {}, 225 | "outputs": [ 226 | { 227 | "name": "stdout", 228 | "output_type": "stream", 229 | "text": [ 230 | "[[ 7.86185 -2.69409 -0.10359 -0.09554 0. 0. 0. 0. 0. 0. ]\n", 231 | " [ 6.34772 0.52255 -0.59993 -0.38732 0. 0. 0. 0. 0. 0. ]\n", 232 | " [ 1.00188 0.31478 2.60084 -2.13566 0. 0. 0. 0. 0. 0. ]\n", 233 | " [ 0. 0. 1.00188 -6.34772 1.49408 -0.31478 -1.10676 0.52255 1.53573 1.10676]\n", 234 | " [ 0. -1.00188 0. 7.86185 0.10962 2.60084 -0.20516 0.55843 0.10359 -0.10962]\n", 235 | " [ 0. 6.34772 -7.86185 0. -0.20948 -2.09416 0.20948 0.38732 -0.09554 0.31307]]\n" 236 | ] 237 | } 238 | ], 239 | "source": [ 240 | "from mujoco_sysid.regressors import joint_body_regressor\n", 241 | "\n", 242 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 243 | " print(joint_body_regressor(model, mjdata, 1))" 244 | ] 245 | }, 246 | { 247 | "cell_type": "code", 248 | "execution_count": 10, 249 | "metadata": {}, 250 | "outputs": [ 251 | { 252 | "name": "stdout", 253 | "output_type": "stream", 254 | "text": [ 255 | "[0. 0. 0. 0. 0. 0.]\n" 256 | ] 257 | } 258 | ], 259 | "source": [ 260 | "force = mjdata.qfrc_actuator.copy()\n", 261 | "force[:3] = mjdata.xmat[1].reshape(3, 3).T @ force[:3]\n", 262 | "\n", 263 | "with np.printoptions(precision=5, suppress=True):\n", 264 | " print(force)" 265 | ] 266 | }, 267 | { 268 | "cell_type": "code", 269 | "execution_count": 11, 270 | "metadata": {}, 271 | "outputs": [ 272 | { 273 | "name": "stdout", 274 | "output_type": "stream", 275 | "text": [ 276 | "[10.38215 8.56271 0.83674 -0.36953 0.55124 0.01578]\n" 277 | ] 278 | } 279 | ], 280 | "source": [ 281 | "with np.printoptions(precision=5, suppress=True):\n", 282 | " print(regressor @ skydio_parameters)" 283 | ] 284 | }, 285 | { 286 | "cell_type": "code", 287 | "execution_count": 12, 288 | "metadata": {}, 289 | "outputs": [ 290 | { 291 | "data": { 292 | "text/plain": [ 293 | "0.0" 294 | ] 295 | }, 296 | "execution_count": 12, 297 | "metadata": {}, 298 | "output_type": "execute_result" 299 | } 300 | ], 301 | "source": [ 302 | "np.linalg.norm(force)" 303 | ] 304 | }, 305 | { 306 | "cell_type": "markdown", 307 | "metadata": {}, 308 | "source": [ 309 | "### Set same values for pinocchio and compare" 310 | ] 311 | }, 312 | { 313 | "cell_type": "code", 314 | "execution_count": 13, 315 | "metadata": {}, 316 | "outputs": [ 317 | { 318 | "data": { 319 | "text/plain": [ 320 | "(array([ 1.76405235, 0.40015721, 0.97873798, 0.58002289, -0.30352125,\n", 321 | " 0.2950768 , 0.69597269]),\n", 322 | " array([0.21085653, 0.35761936, 0.17263309, 0.14404357, 1.45427351,\n", 323 | " 0.76103773]),\n", 324 | " array([ 0.35928247, 0.18472281, -0.3999394 , 1.49407907, -0.20515826,\n", 325 | " 0.3130677 ]))" 326 | ] 327 | }, 328 | "execution_count": 13, 329 | "metadata": {}, 330 | "output_type": "execute_result" 331 | } 332 | ], 333 | "source": [ 334 | "pinq, pinv = muj2pin(mjdata.qpos, mjdata.qvel)\n", 335 | "\n", 336 | "_, pindv = muj2pin(mjdata.qpos, mjdata.qacc)\n", 337 | "\n", 338 | "pinq, pinv, pindv" 339 | ] 340 | }, 341 | { 342 | "cell_type": "code", 343 | "execution_count": 14, 344 | "metadata": {}, 345 | "outputs": [ 346 | { 347 | "name": "stdout", 348 | "output_type": "stream", 349 | "text": [ 350 | "[[ 7.84074 -2.69409 -0.10359 -0.09554 0. 0. 0. 0. 0. 0. ]\n", 351 | " [ 6.48332 0.52255 -0.59993 -0.38732 0. 0. 0. 0. 0. 0. ]\n", 352 | " [ 0.74675 0.31478 2.60084 -2.13566 0. 0. 0. 0. 0. 0. ]\n", 353 | " [ 0. 0. 0.74675 -6.48332 1.49408 -0.31478 -1.10676 0.52255 1.53573 1.10676]\n", 354 | " [ 0. -0.74675 0. 7.84074 0.10962 2.60084 -0.20516 0.55843 0.10359 -0.10962]\n", 355 | " [ 0. 6.48332 -7.84074 0. -0.20948 -2.09416 0.20948 0.38732 -0.09554 0.31307]]\n" 356 | ] 357 | } 358 | ], 359 | "source": [ 360 | "pindata = pin.Data(pinmodel)\n", 361 | "tau = pin.rnea(pinmodel, pindata, pinq, pinv, pindv)\n", 362 | "\n", 363 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 364 | " print(pin.jointBodyRegressor(pinmodel, pindata, 1))" 365 | ] 366 | }, 367 | { 368 | "cell_type": "code", 369 | "execution_count": 15, 370 | "metadata": {}, 371 | "outputs": [ 372 | { 373 | "data": { 374 | "text/plain": [ 375 | "( v = 7.86185 6.34772 1.00188\n", 376 | " w = 1.49408 -0.205158 0.313068,\n", 377 | " v = 0.210857 0.357619 0.172633\n", 378 | " w = 0.144044 1.45427 0.761038)" 379 | ] 380 | }, 381 | "execution_count": 15, 382 | "metadata": {}, 383 | "output_type": "execute_result" 384 | } 385 | ], 386 | "source": [ 387 | "pindata.a_gf[1], pindata.v[1]" 388 | ] 389 | }, 390 | { 391 | "cell_type": "code", 392 | "execution_count": 16, 393 | "metadata": {}, 394 | "outputs": [ 395 | { 396 | "name": "stdout", 397 | "output_type": "stream", 398 | "text": [ 399 | "regressor dv: [7.86184879 6.34771817 1.00187616] dw: [ 1.49407907 -0.20515826 0.3130677 ]\n", 400 | "[[ 7.84074 -2.69409 -0.10359 -0.09554 0. 0. 0. 0. 0. 0. ]\n", 401 | " [ 6.48332 0.52255 -0.59993 -0.38732 0. 0. 0. 0. 0. 0. ]\n", 402 | " [ 0.74675 0.31478 2.60084 -2.13566 0. 0. 0. 0. 0. 0. ]\n", 403 | " [ 0. 0. 0.74675 -6.48332 1.49408 -0.31478 -1.10676 0.52255 1.53573 1.10676]\n", 404 | " [ 0. -0.74675 0. 7.84074 0.10962 2.60084 -0.20516 0.55843 0.10359 -0.10962]\n", 405 | " [ 0. 6.48332 -7.84074 0. -0.20948 -2.09416 0.20948 0.38732 -0.09554 0.31307]]\n" 406 | ] 407 | } 408 | ], 409 | "source": [ 410 | "regressor = mj_bodyRegressor(model, mjdata, 1)\n", 411 | "\n", 412 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 413 | " print(regressor)" 414 | ] 415 | }, 416 | { 417 | "cell_type": "code", 418 | "execution_count": 17, 419 | "metadata": {}, 420 | "outputs": [ 421 | { 422 | "name": "stdout", 423 | "output_type": "stream", 424 | "text": [ 425 | "regressor [10.38215 8.56271 0.83674 -0.36953 0.55124 0.01578]\n", 426 | "rnea [10.38215 8.56271 0.83674 -0.36953 0.55124 0.01578]\n" 427 | ] 428 | } 429 | ], 430 | "source": [ 431 | "with np.printoptions(precision=5, suppress=True, linewidth=400):\n", 432 | " print(\"regressor\", pin.jointBodyRegressor(pinmodel, pindata, 1) @ pinmodel.inertias[1].toDynamicParameters())\n", 433 | " print(\"rnea\", tau)" 434 | ] 435 | }, 436 | { 437 | "cell_type": "code", 438 | "execution_count": 18, 439 | "metadata": {}, 440 | "outputs": [ 441 | { 442 | "data": { 443 | "text/plain": [ 444 | "13.499994938181835" 445 | ] 446 | }, 447 | "execution_count": 18, 448 | "metadata": {}, 449 | "output_type": "execute_result" 450 | } 451 | ], 452 | "source": [ 453 | "np.linalg.norm(tau)" 454 | ] 455 | } 456 | ], 457 | "metadata": { 458 | "kernelspec": { 459 | "display_name": "venv", 460 | "language": "python", 461 | "name": "python3" 462 | }, 463 | "language_info": { 464 | "codemirror_mode": { 465 | "name": "ipython", 466 | "version": 3 467 | }, 468 | "file_extension": ".py", 469 | "mimetype": "text/x-python", 470 | "name": "python", 471 | "nbconvert_exporter": "python", 472 | "pygments_lexer": "ipython3", 473 | "version": "3.10.12" 474 | } 475 | }, 476 | "nbformat": 4, 477 | "nbformat_minor": 2 478 | } 479 | -------------------------------------------------------------------------------- /data/models/cart_pole/model.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 16 | 17 | 18 | 19 | 20 | 21 | -------------------------------------------------------------------------------- /data/models/cart_pole/scene.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 24 | 25 | -------------------------------------------------------------------------------- /data/trajectories/skydio.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/based-robotics/mujoco-sysid/2c64beb5eb7c0026eb69837287587a6f37e283f8/data/trajectories/skydio.npz -------------------------------------------------------------------------------- /data/trajectories/skydio_easy.npz: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/based-robotics/mujoco-sysid/2c64beb5eb7c0026eb69837287587a6f37e283f8/data/trajectories/skydio_easy.npz -------------------------------------------------------------------------------- /examples/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | class ActuatorMotor: 5 | def __init__(self) -> None: 6 | self.dyn = np.array([1, 0, 0]) 7 | self.gain = np.array([1, 0, 0]) 8 | self.bias = np.array([0, 0, 0]) 9 | 10 | def __repr__(self) -> str: 11 | return f"ActuatorMotor(dyn={self.dyn}, gain={self.gain}, bias={self.bias})" 12 | 13 | 14 | def update_actuator(model, actuator_id, actuator): 15 | """ 16 | Update actuator in model 17 | model - mujoco.MjModel 18 | actuator_id - int or str (name) (for reference see, named access to model elements) 19 | actuator - ActuatorMotor, ActuatorPosition, ActuatorVelocity 20 | """ 21 | 22 | model.actuator(actuator_id).dynprm[:3] = actuator.dyn 23 | model.actuator(actuator_id).gainprm[:3] = actuator.gain 24 | model.actuator(actuator_id).biasprm[:3] = actuator.bias 25 | model.actuator(actuator_id).ctrlrange = None 26 | -------------------------------------------------------------------------------- /mujoco_sysid/__init__.py: -------------------------------------------------------------------------------- 1 | from .regressors import body_regressor, joint_body_regressor, joint_torque_regressor 2 | -------------------------------------------------------------------------------- /mujoco_sysid/convert.py: -------------------------------------------------------------------------------- 1 | """ 2 | This module provides functions for converting between different representations of pseudo inertia matrices, 3 | including theta parameters, logarithmic Cholesky parameters, and Cholesky decomposition. 4 | 5 | The main representations used are: 6 | - Theta parameters: A 10-dimensional vector containing mass, first moments, and inertia tensor components. 7 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 8 | 9 | - Pseudo inertia matrix: A 4x4 symmetric matrix representing the pseudo inertia of a rigid body. 10 | pseudo = [[-0.5*I_xx + 0.5*I_yy + 0.5*I_zz, -I_xy, -I_xz, mr_x], 11 | [-I_xy, 0.5*I_xx - 0.5*I_yy + 0.5*I_zz, -I_yz, mr_y], 12 | [-I_xz, -I_yz, 0.5*I_xx + 0.5*I_yy - 0.5*I_zz, mr_z], 13 | [mr_x, mr_y, mr_z, m]] 14 | 15 | - Cholesky decomposition: An upper triangular matrix U such that the pseudo inertia matrix J = U * U^T. 16 | U = [[exp(alpha)*exp(d_1), s_12*exp(alpha), s_13*exp(alpha), t_1*exp(alpha)], 17 | [0, exp(alpha)*exp(d_2), s_23*exp(alpha), t_2*exp(alpha)], 18 | [0, 0, exp(alpha)*exp(d_3), t_3*exp(alpha)], 19 | [0, 0, 0, exp(alpha)]] 20 | 21 | - Logarithmic Cholesky parameters: A 10-dimensional vector containing logarithmic parameters derived from the Cholesky decomposition. 22 | logchol = [alpha, d_1, d_2, d_3, s_12, s_23, s_13, t_1, t_2, t_3] 23 | 24 | Functions: 25 | - `theta2pseudo(theta: np.ndarray) -> np.ndarray` 26 | Converts theta parameters to a pseudo inertia matrix. 27 | - `pseudo2theta(pseudo_inertia: np.ndarray) -> np.ndarray` 28 | Converts a pseudo inertia matrix to theta parameters. 29 | - `logchol2chol(log_cholesky: np.ndarray) -> np.ndarray` 30 | Converts logarithmic Cholesky parameters to the Cholesky matrix, factoring out exp(alpha). 31 | - `chol2logchol(U: np.ndarray) -> np.ndarray` 32 | Converts the upper triangular matrix U to logarithmic Cholesky parameters. 33 | - `pseudo2cholesky(pseudo_inertia: np.ndarray) -> np.ndarray` 34 | Computes the Cholesky decomposition of a pseudo inertia matrix. 35 | - `cholesky2pseudo(U: np.ndarray) -> np.ndarray` 36 | Converts the upper triangular Cholesky matrix U back into a pseudo inertia matrix. 37 | - `pseudo2logchol(pseudo_inertia: np.ndarray) -> np.ndarray` 38 | Converts a pseudo inertia matrix to logarithmic Cholesky parameters. 39 | - `theta2logchol(theta: np.ndarray) -> np.ndarray` 40 | Converts theta parameters directly to logarithmic Cholesky parameters. 41 | 42 | For more information and derivations, please consider reviewing the following references: 43 | - Rucker C, Wensing PM. Smooth parameterization of rigid-body inertia. IEEE Robotics and Automation Letters. 2022 Jan 21;7(2):2771-8. 44 | - Wensing PM, Kim S, Slotine JJ. Linear matrix inequalities for physically consistent inertial parameter identification: 45 | A statistical perspective on the mass distribution. IEEE Robotics and Automation Letters. 2017 Jul 20;3(1):60-7. 46 | 47 | Additional discussions and derivations are available at: 48 | https://colab.research.google.com/drive/1xFte2FT0nQ0ePs02BoOx4CmLLw5U-OUZ#scrollTo=Xt86l6AtZBhI 49 | """ # noqa: E501 50 | 51 | import numpy as np 52 | 53 | 54 | def theta2pseudo(theta: np.ndarray) -> np.ndarray: 55 | """ 56 | Converts theta parameters to a pseudo inertia matrix. 57 | 58 | Args: 59 | theta (np.ndarray): A 10-dimensional vector containing mass, first moments, and inertia tensor components. 60 | - theta[0]: Mass (m) 61 | - theta[1:4]: First moments (mc_x, mc_y, mc_z) 62 | - theta[4:10]: Inertia tensor components (I_xx, I_yy, I_zz, I_xy, I_xz, I_yz) 63 | 64 | Returns: 65 | np.ndarray: A 4x4 pseudo inertia matrix. 66 | """ 67 | m = theta[0] 68 | h = theta[1:4] 69 | I_xx, I_xy, I_yy, I_xz, I_yz, I_zz = theta[4:] 70 | 71 | I_bar = np.array([[I_xx, I_xy, I_xz], [I_xy, I_yy, I_yz], [I_xz, I_yz, I_zz]]) 72 | 73 | Sigma = 0.5 * np.trace(I_bar) * np.eye(3) - I_bar 74 | 75 | pseudo_inertia = np.zeros((4, 4)) 76 | pseudo_inertia[:3, :3] = Sigma 77 | pseudo_inertia[:3, 3] = h 78 | pseudo_inertia[3, :3] = h 79 | pseudo_inertia[3, 3] = m 80 | 81 | return pseudo_inertia 82 | 83 | 84 | def pseudo2theta(pseudo_inertia: np.ndarray) -> np.ndarray: 85 | """ 86 | Converts a pseudo inertia matrix to theta parameters. 87 | 88 | Args: 89 | pseudo_inertia (np.ndarray): A 4x4 symmetric matrix representing the pseudo inertia of a rigid body. 90 | 91 | Returns: 92 | np.ndarray: A 10-dimensional vector containing mass, first moments, and inertia tensor components. 93 | - theta[0]: Mass (m) 94 | - theta[1:4]: First moments (mc_x, mc_y, mc_z) 95 | - theta[4:10]: Inertia tensor components (I_xx, I_yy, I_zz, I_xy, I_xz, I_yz) 96 | """ 97 | m = pseudo_inertia[3, 3] 98 | h = pseudo_inertia[:3, 3] 99 | Sigma = pseudo_inertia[:3, :3] 100 | 101 | I_bar = np.trace(Sigma) * np.eye(3) - Sigma 102 | 103 | I_xx = I_bar[0, 0] 104 | I_xy = I_bar[0, 1] 105 | I_yy = I_bar[1, 1] 106 | I_xz = I_bar[0, 2] 107 | I_yz = I_bar[1, 2] 108 | I_zz = I_bar[2, 2] 109 | 110 | theta = np.array([m, h[0], h[1], h[2], I_xx, I_xy, I_yy, I_xz, I_yz, I_zz]) 111 | 112 | return theta 113 | 114 | 115 | def logchol2chol(log_cholesky): 116 | """ 117 | Converts logarithmic Cholesky parameters to the Cholesky matrix, factoring out exp(alpha). 118 | 119 | Args: 120 | log_cholesky (np.ndarray): A 10-dimensional vector containing logarithmic Cholesky parameters. 121 | - log_cholesky[0]: alpha 122 | - log_cholesky[1:4]: (d1, d2, d3) 123 | - log_cholesky[4:7]: (s12, s23, s13) 124 | - log_cholesky[7:10]: (t1, t2, t3) 125 | 126 | Returns: 127 | np.ndarray: A 4x4 upper triangular Cholesky matrix U. 128 | """ 129 | alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = log_cholesky 130 | 131 | # Compute the exponential terms 132 | exp_alpha = np.exp(alpha) 133 | exp_d1 = np.exp(d1) 134 | exp_d2 = np.exp(d2) 135 | exp_d3 = np.exp(d3) 136 | 137 | # Construct the scaled Cholesky matrix U without exp_alpha 138 | U = np.zeros((4, 4)) 139 | U[0, 0] = exp_d1 140 | U[0, 1] = s12 141 | U[0, 2] = s13 142 | U[0, 3] = t1 143 | U[1, 1] = exp_d2 144 | U[1, 2] = s23 145 | U[1, 3] = t2 146 | U[2, 2] = exp_d3 147 | U[2, 3] = t3 148 | U[3, 3] = 1 149 | 150 | # Multiply the entire matrix by exp_alpha 151 | U *= exp_alpha 152 | 153 | return U 154 | 155 | 156 | def chol2logchol(U: np.ndarray) -> np.ndarray: 157 | """ 158 | Converts the upper triangular matrix U to logarithmic Cholesky parameters. 159 | 160 | Args: 161 | U (np.ndarray): A 4x4 upper triangular matrix decomposition of the pseudo inertia matrix. 162 | 163 | Returns: 164 | np.ndarray: A 10-dimensional vector containing logarithmic Cholesky parameters. 165 | - log_cholesky[0]: alpha 166 | - log_cholesky[1:4]: (d_1, d_2, d_3) 167 | - log_cholesky[4:7]: (s12, s23, s13) 168 | - log_cholesky[7:10]: (t_1, t_2, t_3) 169 | """ 170 | 171 | alpha = np.log(U[3, 3]) 172 | d1 = np.log(U[0, 0] / U[3, 3]) 173 | d2 = np.log(U[1, 1] / U[3, 3]) 174 | d3 = np.log(U[2, 2] / U[3, 3]) 175 | s12 = U[0, 1] / U[3, 3] 176 | s23 = U[1, 2] / U[3, 3] 177 | s13 = U[0, 2] / U[3, 3] 178 | t1 = U[0, 3] / U[3, 3] 179 | t2 = U[1, 3] / U[3, 3] 180 | t3 = U[2, 3] / U[3, 3] 181 | return np.array([alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3]) 182 | 183 | 184 | def logchol2theta(log_cholesky: np.ndarray) -> np.ndarray: 185 | """ 186 | Converts logarithmic Cholesky parameters directly to theta parameters. 187 | 188 | Args: 189 | log_cholesky (np.ndarray): A 10-dimensional vector containing logarithmic Cholesky parameters. 190 | - log_cholesky[0]: alpha 191 | - log_cholesky[1:4]: (d_1, d_2, d_3) 192 | - log_cholesky[4:7]: (s12, s23, s13) 193 | - log_cholesky[7:10]: (t_1, t_2, t_3) 194 | 195 | Returns: 196 | np.ndarray: A 10-dimensional vector containing mass, first moments, and inertia tensor components. 197 | - theta[0]: Mass (m) 198 | - theta[1:4]: First moments (mc_x, mc_y, mc_z) 199 | - theta[4:10]: Inertia tensor components (I_xx, I_yy, I_zz, I_xy, I_xz, I_yz) 200 | """ 201 | alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = log_cholesky 202 | 203 | # Calculate exponential terms without applying exp_2_alpha 204 | exp_d1 = np.exp(d1) 205 | exp_d2 = np.exp(d2) 206 | exp_d3 = np.exp(d3) 207 | 208 | # Compute the elements of the output vector without exp_2_alpha 209 | theta = np.zeros(10) 210 | theta[0] = 1 211 | theta[1] = t1 212 | theta[2] = t2 213 | theta[3] = t3 214 | theta[4] = s23**2 + t2**2 + t3**2 + exp_d2**2 + exp_d3**2 215 | theta[5] = -s12 * exp_d2 - s13 * s23 - t1 * t2 216 | theta[6] = s12**2 + s13**2 + t1**2 + t3**2 + exp_d1**2 + exp_d3**2 217 | theta[7] = -s13 * exp_d3 - t1 * t3 218 | theta[8] = -s23 * exp_d3 - t2 * t3 219 | theta[9] = s12**2 + s13**2 + s23**2 + t1**2 + t2**2 + exp_d1**2 + exp_d2**2 220 | 221 | # Calculate exp_2_alpha and scale the theta vector 222 | exp_2_alpha = np.exp(2 * alpha) 223 | theta *= exp_2_alpha 224 | 225 | return theta 226 | 227 | 228 | def pseudo2cholesky(pseudo_inertia: np.ndarray) -> np.ndarray: 229 | """ 230 | Computes the Cholesky decomposition of a pseudo inertia matrix. 231 | Note that this is UPPER triangular decomposition in form J = U*U^T, which is not the usual way to calculate the Cholesky decomposition. 232 | However, in this form, the associated logarithmic Cholesky parameters have a geometrical meaning. 233 | 234 | Args: 235 | pseudo_inertia (np.ndarray): A 4x4 symmetric matrix representing the pseudo inertia of a rigid body. 236 | 237 | Returns: 238 | np.ndarray: A 4x4 upper triangular Cholesky matrix U. 239 | """ # noqa: E501 240 | 241 | n = pseudo_inertia.shape[0] 242 | indices = np.arange(n - 1, -1, -1) # Indices to reverse the order 243 | 244 | # Apply the inversion using indices for rows and columns 245 | reversed_inertia = pseudo_inertia[indices][:, indices] 246 | 247 | # Perform Cholesky decomposition on the permuted matrix A' 248 | L_prime = np.linalg.cholesky(reversed_inertia) 249 | 250 | # Apply the reverse permutation to L_prime and transpose it to form U 251 | U = L_prime[indices][:, indices] 252 | 253 | return U 254 | 255 | 256 | def cholesky2pseudo(U: np.ndarray) -> np.ndarray: 257 | """ 258 | Converts the upper triangular Cholesky matrix U back into a pseudo inertia matrix. 259 | 260 | Args: 261 | U (np.ndarray): A 4x4 upper triangular Cholesky matrix. 262 | 263 | Returns: 264 | np.ndarray: A 4x4 pseudo inertia matrix. 265 | """ 266 | return U @ U.T 267 | 268 | 269 | def pseudo2logchol(pseudo_inertia: np.ndarray) -> np.ndarray: 270 | """ 271 | Converts a pseudo inertia matrix to logarithmic Cholesky parameters. 272 | 273 | Args: 274 | pseudo_inertia (np.ndarray): A 4x4 symmetric matrix representing the pseudo inertia of a rigid body. 275 | 276 | Returns: 277 | np.ndarray: A 10-dimensional vector containing logarithmic Cholesky parameters. 278 | - log_cholesky[0]: alpha 279 | - log_cholesky[1:4]: (d_1, d_2, d_3) 280 | - log_cholesky[4:7]: (s_12, s_23, s_13) 281 | - log_cholesky[7:10]: (t_1, t_2, t_3) 282 | """ 283 | # theta = pseudo2theta(pseudo_inertia) 284 | U = pseudo2cholesky(pseudo_inertia) 285 | 286 | logchol = chol2logchol(U) 287 | 288 | return logchol 289 | 290 | 291 | def theta2logchol(theta: np.ndarray) -> np.ndarray: 292 | """ 293 | Converts theta parameters directly to logarithmic Cholesky parameters. 294 | 295 | Args: 296 | theta (np.ndarray): A 10-dimensional vector containing mass, first moments, and inertia tensor components. 297 | - theta[0]: Mass (m) 298 | - theta[1:4]: First moments (mc_x, mc_y, mc_z) 299 | - theta[4:10]: Inertia tensor components (I_xx, I_yy, I_zz, I_xy, I_xz, I_yz) 300 | 301 | Returns: 302 | np.ndarray: A 10-dimensional vector containing logarithmic Cholesky parameters. 303 | - log_cholesky[0]: alpha 304 | - log_cholesky[1:4]: (d_1, d_2, d_3) 305 | - log_cholesky[4:7]: (s_12, s_23, s_13) 306 | - log_cholesky[7:10]: (t_1, t_2, t_3) 307 | """ 308 | pseudo_inertia = theta2pseudo(theta) 309 | return pseudo2logchol(pseudo_inertia) 310 | -------------------------------------------------------------------------------- /mujoco_sysid/mjx/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/based-robotics/mujoco-sysid/2c64beb5eb7c0026eb69837287587a6f37e283f8/mujoco_sysid/mjx/__init__.py -------------------------------------------------------------------------------- /mujoco_sysid/mjx/convert.py: -------------------------------------------------------------------------------- 1 | import jax.numpy as np 2 | 3 | 4 | def theta2pseudo(theta: np.ndarray) -> np.ndarray: 5 | m = theta[0] 6 | h = theta[1:4] 7 | I_xx, I_xy, I_yy, I_xz, I_yz, I_zz = theta[4:] 8 | 9 | I_bar = np.array([[I_xx, I_xy, I_xz], [I_xy, I_yy, I_yz], [I_xz, I_yz, I_zz]]) 10 | 11 | Sigma = 0.5 * np.trace(I_bar) * np.eye(3) - I_bar 12 | 13 | pseudo_inertia = np.zeros((4, 4)) 14 | pseudo_inertia = pseudo_inertia.at[:3, :3].set(Sigma) 15 | pseudo_inertia = pseudo_inertia.at[:3, 3].set(h) 16 | pseudo_inertia = pseudo_inertia.at[3, :3].set(h) 17 | pseudo_inertia = pseudo_inertia.at[3, 3].set(m) 18 | 19 | return pseudo_inertia 20 | 21 | 22 | def pseudo2theta(pseudo_inertia: np.ndarray) -> np.ndarray: 23 | m = pseudo_inertia[3, 3] 24 | h = pseudo_inertia[:3, 3] 25 | Sigma = pseudo_inertia[:3, :3] 26 | 27 | I_bar = np.trace(Sigma) * np.eye(3) - Sigma 28 | 29 | I_xx = I_bar[0, 0] 30 | I_xy = I_bar[0, 1] 31 | I_yy = I_bar[1, 1] 32 | I_xz = I_bar[0, 2] 33 | I_yz = I_bar[1, 2] 34 | I_zz = I_bar[2, 2] 35 | 36 | theta = np.array([m, h[0], h[1], h[2], I_xx, I_xy, I_yy, I_xz, I_yz, I_zz]) 37 | 38 | return theta 39 | 40 | 41 | def logchol2chol(log_cholesky): 42 | alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = log_cholesky 43 | 44 | exp_alpha = np.exp(alpha) 45 | exp_d1 = np.exp(d1) 46 | exp_d2 = np.exp(d2) 47 | exp_d3 = np.exp(d3) 48 | 49 | U = np.zeros((4, 4)) 50 | U = U.at[0, 0].set(exp_d1) 51 | U = U.at[0, 1].set(s12) 52 | U = U.at[0, 2].set(s13) 53 | U = U.at[0, 3].set(t1) 54 | U = U.at[1, 1].set(exp_d2) 55 | U = U.at[1, 2].set(s23) 56 | U = U.at[1, 3].set(t2) 57 | U = U.at[2, 2].set(exp_d3) 58 | U = U.at[2, 3].set(t3) 59 | U = U.at[3, 3].set(1) 60 | 61 | U *= exp_alpha 62 | 63 | return U 64 | 65 | 66 | def chol2logchol(U: np.ndarray) -> np.ndarray: 67 | alpha = np.log(U[3, 3]) 68 | d1 = np.log(U[0, 0] / U[3, 3]) 69 | d2 = np.log(U[1, 1] / U[3, 3]) 70 | d3 = np.log(U[2, 2] / U[3, 3]) 71 | s12 = U[0, 1] / U[3, 3] 72 | s23 = U[1, 2] / U[3, 3] 73 | s13 = U[0, 2] / U[3, 3] 74 | t1 = U[0, 3] / U[3, 3] 75 | t2 = U[1, 3] / U[3, 3] 76 | t3 = U[2, 3] / U[3, 3] 77 | return np.array([alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3]) 78 | 79 | 80 | def logchol2theta(log_cholesky: np.ndarray) -> np.ndarray: 81 | alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = log_cholesky 82 | 83 | exp_d1 = np.exp(d1) 84 | exp_d2 = np.exp(d2) 85 | exp_d3 = np.exp(d3) 86 | 87 | theta = np.array( 88 | [ 89 | 1, 90 | t1, 91 | t2, 92 | t3, 93 | s23**2 + t2**2 + t3**2 + exp_d2**2 + exp_d3**2, 94 | -s12 * exp_d2 - s13 * s23 - t1 * t2, 95 | s12**2 + s13**2 + t1**2 + t3**2 + exp_d1**2 + exp_d3**2, 96 | -s13 * exp_d3 - t1 * t3, 97 | -s23 * exp_d3 - t2 * t3, 98 | s12**2 + s13**2 + s23**2 + t1**2 + t2**2 + exp_d1**2 + exp_d2**2, 99 | ] 100 | ) 101 | 102 | exp_2_alpha = np.exp(2 * alpha) 103 | theta *= exp_2_alpha 104 | 105 | return theta 106 | 107 | 108 | def pseudo2cholesky(pseudo_inertia: np.ndarray) -> np.ndarray: 109 | n = pseudo_inertia.shape[0] 110 | indices = np.arange(n - 1, -1, -1) 111 | 112 | reversed_inertia = pseudo_inertia[indices][:, indices] 113 | 114 | L_prime = np.linalg.cholesky(reversed_inertia) 115 | 116 | U = L_prime[indices][:, indices] 117 | 118 | return U 119 | 120 | 121 | def cholesky2pseudo(U: np.ndarray) -> np.ndarray: 122 | return U @ U.T 123 | 124 | 125 | def pseudo2logchol(pseudo_inertia: np.ndarray) -> np.ndarray: 126 | U = pseudo2cholesky(pseudo_inertia) 127 | logchol = chol2logchol(U) 128 | return logchol 129 | 130 | 131 | def theta2logchol(theta: np.ndarray) -> np.ndarray: 132 | pseudo_inertia = theta2pseudo(theta) 133 | return pseudo2logchol(pseudo_inertia) 134 | -------------------------------------------------------------------------------- /mujoco_sysid/mjx/parameters.py: -------------------------------------------------------------------------------- 1 | import jax.numpy as jnp 2 | import jax.typing as jpt 3 | from jax.scipy.spatial.transform import Rotation 4 | from mujoco import mjx 5 | from mujoco.mjx._src import math 6 | 7 | 8 | def skew(vector: jpt.ArrayLike) -> jpt.ArrayLike: 9 | return jnp.cross(jnp.eye(vector.size), vector.reshape(-1)) 10 | 11 | 12 | def get_dynamic_parameters(mjx_model: mjx.Model, bodyid) -> jpt.ArrayLike: 13 | """Get the dynamic parameters \theta of a body 14 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 15 | 16 | Args: 17 | mjx_model (mujoco.mjx.MjModel): The mujoco model 18 | bodyid (int): The id of the body 19 | 20 | Returns: 21 | jpt.ArrayLike: theta of the body 22 | """ 23 | mass = mjx_model.body_mass[bodyid] 24 | rc = mjx_model.body_ipos[bodyid] 25 | diag_inertia = mjx_model.body_inertia[bodyid] 26 | 27 | # get the orientation of the body 28 | R = math.quat_to_mat(mjx_model.body_iquat[bodyid]) 29 | 30 | shift = mass * skew(rc) @ skew(rc) 31 | mjinertia = R @ jnp.diag(diag_inertia) @ R.T - shift 32 | 33 | upper_triangular = jnp.array( 34 | [ 35 | mjinertia[0, 0], 36 | mjinertia[0, 1], 37 | mjinertia[1, 1], 38 | mjinertia[0, 2], 39 | mjinertia[1, 2], 40 | mjinertia[2, 2], 41 | ] 42 | ) 43 | 44 | return jnp.concatenate([jnp.array([mass]), mass * rc, upper_triangular]) 45 | 46 | 47 | def set_dynamic_parameters(mjx_model: mjx.Model, bodyid, theta: jpt.ArrayLike) -> mjx.Model: 48 | """Set the dynamic parameters to a body 49 | 50 | Args: 51 | mjx_model (mujoco.mjx.MjModel): The mujoco model 52 | bodyid (int): The id of the body 53 | theta (jpt.ArrayLike): The dynamic parameters of the body 54 | 55 | Returns: 56 | mjx.Model: The updated model 57 | """ 58 | mass = theta[0] 59 | rc = theta[1:4] / mass 60 | inertia = theta[4:] 61 | inertia_full = jnp.array( 62 | [ 63 | [inertia[0], inertia[1], inertia[3]], 64 | [inertia[1], inertia[2], inertia[4]], 65 | [inertia[3], inertia[4], inertia[5]], 66 | ] 67 | ) 68 | 69 | # shift the inertia matrix 70 | inertia_full += mass * skew(rc) @ skew(rc) 71 | 72 | # eigen decomposition 73 | diag_inertia, R = jnp.linalg.eigh(inertia_full) 74 | 75 | # check if any of the eigenvalues are negative 76 | if jnp.any(diag_inertia < 0): 77 | raise ValueError("Inertia matrix is not positive definite") 78 | 79 | # convert the rotation matrix to quaternion 80 | quat = Rotation.from_matrix(R).as_quat() 81 | 82 | # update the body parameters 83 | return mjx_model.tree_replace( 84 | { 85 | "body_inertia": mjx_model.body_inertia.at[bodyid].set(diag_inertia), 86 | "body_iquat": mjx_model.body_iquat.at[bodyid].set(quat), 87 | "body_mass": mjx_model.body_mass.at[bodyid].set(mass), 88 | } 89 | ) 90 | -------------------------------------------------------------------------------- /mujoco_sysid/mjx/regressors.py: -------------------------------------------------------------------------------- 1 | import jax 2 | import jax.numpy as jnp 3 | import jax.typing as jpt 4 | import mujoco.mjx as mjx 5 | from mujoco.mjx._src import scan 6 | from mujoco.mjx._src.math import transform_motion 7 | from mujoco.mjx._src.types import DisableBit 8 | 9 | 10 | def object_velocity(mjx_model: mjx.Model, mjx_data: mjx.Data, bodyid) -> jpt.ArrayLike: 11 | pos = mjx_data.xpos[bodyid] 12 | rot = mjx_data.xmat[bodyid] # TODO: maybe reshape is required here 13 | 14 | # transform spatial 15 | vel = mjx_data.cvel[bodyid] 16 | oldpos = mjx_data.subtree_com[mjx_model.body_rootid[bodyid]] 17 | 18 | return transform_motion(vel, pos - oldpos, rot) 19 | 20 | 21 | # TODO: implement me in future 22 | # def mjx_rnePostConstraint(m: mjx.Model, d: mjx.Data) -> jpt.ArrayLike: 23 | # nbody = m.nbody 24 | # cfrc_com = jnp.zeros(6) 25 | # cfrc = jnp.zeros(6) 26 | # lfrc = jnp.zeros(6) 27 | 28 | # all_cacc = jnp.zeros((nbody, 6)) 29 | 30 | # # clear cacc, set world acceleration to -gravity 31 | # if not m.opt.disableflags & DisableBit.GRAVITY: 32 | # cacc = jnp.concatenate((jnp.zeros((nbody, 3)), -m.opt.gravity), axis=1) 33 | 34 | # # FIXME: assumption that xfrc_applied is zero 35 | # # FIXME: assumption that contacts are zero 36 | # # FIXME: assumption that connect and weld constraints are zero 37 | 38 | # # forward pass over bodies: compute acc 39 | # cacc = jnp.zeros(6) 40 | # for j in range(nbody): 41 | # bda = m.body_dofadr[j] 42 | 43 | # # cacc = cacc_parent + cdofdot * qvel + cdof * qacc 44 | # cacc = all_cacc[m.body_parentid[j]] + d.cdof_dot[bda] * d.qvel[bda] + d.cdof[bda] * d.qacc[bda] 45 | 46 | 47 | # def com_acc(m: mjx.Model, d: mjx.Data) -> jpt.ArrayLike: 48 | # # forward scan over tree: accumulate link center of mass acceleration 49 | # def cacc_fn(cacc, cdof_dot, qvel): 50 | # if cacc is None: 51 | # if m.opt.disableflags & DisableBit.GRAVITY: 52 | # cacc = jnp.zeros((6,)) 53 | # else: 54 | # cacc = jnp.concatenate((jnp.zeros((3,)), -m.opt.gravity)) 55 | 56 | # cacc += jnp.sum(jax.vmap(jnp.multiply)(cdof_dot, qvel), axis=0) 57 | 58 | # return cacc 59 | 60 | # return scan.body_tree(m, cacc_fn, "vv", "b", d.cdof_dot, d.qvel) 61 | 62 | 63 | def object_acceleration(mjx_model: mjx.Model, mjx_data: mjx.Data, bodyid) -> jpt.ArrayLike: 64 | pos = mjx_data.xpos[bodyid] 65 | rot = mjx_data.xmat[bodyid] # TODO: maybe reshape is required here 66 | 67 | # transform spatial 68 | accel = mjx_data.cacc[bodyid] 69 | oldpos = mjx_data.subtree_com[mjx_model.body_rootid[bodyid]] 70 | 71 | velocity = object_velocity(mjx_model, mjx_data, bodyid) 72 | 73 | # transform com-based acceleration to local frame 74 | acceleration = transform_motion(accel, pos - oldpos, rot) 75 | 76 | # acc_tran += vel_rot x vel_tran 77 | correction = jnp.cross(velocity[:3], velocity[3:]) 78 | acceleration[3:] += correction 79 | 80 | return acceleration 81 | 82 | 83 | def body_energyRegressor( 84 | v: jpt.ArrayLike, 85 | w: jpt.ArrayLike, 86 | r: jpt.ArrayLike, 87 | R: jpt.ArrayLike, 88 | gravity: jpt.ArrayLike | None = None, 89 | ) -> tuple[jpt.ArrayLike, jpt.ArrayLike]: 90 | """ 91 | Computes the kinetic and potential energy regressors for a single body. 92 | 93 | The energy regressors are computed to represent the kinetic and potential energy of the system 94 | as a linear combination of the inertial parameters of the bodies. 95 | 96 | The kinetic energy of a rigid body is given by: 97 | K = 0.5 * v^T * M * v 98 | 99 | The potential energy of a rigid body in a uniform gravitational field is given by: 100 | U = m * g^T * r 101 | 102 | where: 103 | v is the spatial velocity of the body 104 | M is the spatial inertia matrix of the body 105 | m is the mass of the body 106 | g is the gravitational acceleration vector 107 | r is the position of the body's center of mass 108 | 109 | The energy regressors Y_k and Y_u are matrices such that: 110 | K = Y_k * theta 111 | U = Y_u * theta 112 | 113 | Args: 114 | v (npt.ArrayLike): Linear velocity of the body. 115 | w (npt.ArrayLike): Angular velocity of the body. 116 | r (npt.ArrayLike): Position of the body. 117 | R (npt.ArrayLike): Rotation matrix of the body. 118 | gravity (npt.ArrayLike, optional): Gravity vector. Defaults to np.array([0, 0, 9.81]). 119 | 120 | Returns: 121 | tuple[npt.ArrayLike, npt.ArrayLike]: Kinetic and potential energy regressors. 122 | """ 123 | if gravity is None: 124 | gravity = jnp.array([0, 0, 9.81]) 125 | 126 | kinetic = jnp.array( 127 | [ 128 | 0.5 * (v[0] ** 2 + v[1] ** 2 + v[2] ** 2), 129 | -w[1] * v[2] + w[2] * v[1], 130 | w[0] * v[2] - w[2] * v[0], 131 | -w[0] * v[1] + w[1] * v[0], 132 | 0.5 * w[0] ** 2, 133 | w[0] * w[1], 134 | 0.5 * w[1] ** 2, 135 | w[0] * w[2], 136 | w[1] * w[2], 137 | 0.5 * w[2] ** 2, 138 | ] 139 | ) 140 | 141 | potential = jnp.array([*(gravity.T @ r).flatten(), *(gravity.T @ R).flatten(), 0, 0, 0, 0, 0, 0]) 142 | 143 | return kinetic, potential 144 | 145 | 146 | def energy_regressor(mjx_model: mjx.Model, mjx_data: mjx.Data) -> jpt.ArrayLike: 147 | """Get the energy regressor matrices 148 | 149 | The energy regressors Y_k and Y_u are matrices such that: 150 | K = Y_k * theta 151 | U = Y_u * theta 152 | 153 | where: 154 | theta is the vector of inertial parameters of the bodies (10 parameters per body): 155 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 156 | 157 | The total energy regressor Y_e is simply the sum of Y_k and Y_u: 158 | E = K + U = Y_e * theta 159 | 160 | 161 | Args: 162 | mjx_model (mujoco.mjx.MjModel): The mujoco model 163 | mjx_data (mujoco.mjx.MjData): The mujoco data 164 | 165 | Returns: 166 | jpt.ArrayLike: The regressor matrix 167 | """ 168 | 169 | njoints = mjx_model.njnt 170 | # energy_regressor = jnp.zeros(njoints * 10) 171 | kinetic_regressor = jnp.zeros(njoints * 10) 172 | potential_regressor = jnp.zeros(njoints * 10) 173 | 174 | for jnt_id, bodyid in enumerate(mjx_model.jnt_bodyid): 175 | # get the spatial velocity 176 | velocity = object_velocity(mjx_model, mjx_data, bodyid) 177 | v, w = velocity[3:], velocity[:3] 178 | 179 | rot = mjx_data.xmat[bodyid] 180 | pos = mjx_data.xpos[bodyid] 181 | 182 | kinetic, potential = body_energyRegressor(v, w, pos, rot) 183 | kinetic_regressor = kinetic_regressor.at[10 * jnt_id : 10 * jnt_id + 10].set(kinetic) 184 | potential_regressor = potential_regressor.at[10 * jnt_id : 10 * jnt_id + 10].set(potential) 185 | 186 | return kinetic_regressor + potential_regressor 187 | 188 | 189 | def potential_energy_bias(mjxmodel: mjx.Model): 190 | """ 191 | The bodies before the first joint are considered to be fixed in space. 192 | They are included in potential energy calculation, but not in the regressor. 193 | """ 194 | 195 | bias = 0 196 | for i in range(mjxmodel.nbody): 197 | if i not in mjxmodel.jnt_bodyid: 198 | bias += mjxmodel.body_mass[i] * mjxmodel.opt.gravity[2] * mjxmodel.body_ipos[i][2] 199 | 200 | return bias 201 | 202 | 203 | def body_regressor( 204 | v_lin: jpt.ArrayLike, 205 | v_ang: jpt.ArrayLike, 206 | a_lin: jpt.ArrayLike, 207 | a_ang: jpt.ArrayLike, 208 | ) -> jpt.ArrayLike: 209 | """Y_body returns a regressor for a single rigid body 210 | 211 | Newton-Euler equations for a rigid body are given by: 212 | M * a_g + v x M * v = f 213 | 214 | where: 215 | M is the spatial inertia matrix of the body 216 | a_g is the acceleration of the body 217 | v is the spatial velocity of the body 218 | f is the spatial force acting on the body 219 | 220 | The regressor is a matrix Y such that: 221 | Y \theta = f 222 | 223 | where: 224 | \theta is the vector of inertial parameters of the body (10 parameters) 225 | 226 | More expressive derivation is given here: 227 | https://colab.research.google.com/drive/1xFte2FT0nQ0ePs02BoOx4CmLLw5U-OUZ?usp=sharing 228 | 229 | Args: 230 | v_lin (npt.ArrayLike): linear velocity of the body 231 | v_ang (npt.ArrayLike): angular velocity of the body 232 | a_lin (npt.ArrayLike): linear acceleration of the body 233 | a_ang (npt.ArrayLike): angular acceleration of the body 234 | 235 | Returns: 236 | npt.ArrayLike: regressor for the body 237 | """ 238 | v1, v2, v3 = v_lin 239 | v4, v5, v6 = v_ang 240 | 241 | a1, a2, a3 = a_lin 242 | a4, a5, a6 = a_ang 243 | 244 | # fmt: off 245 | return jnp.array([ 246 | [a1 - v2*v6 + v3*v5, -v5**2 - v6**2, -a6 + v4*v5, a5 + v4*v6, 0, 0, 0, 0, 0, 0], 247 | [a2 + v1*v6 - v3*v4, a6 + v4*v5, -v4**2 - v6**2, -a4 + v5*v6, 0, 0, 0, 0, 0, 0], 248 | [a3 - v1*v5 + v2*v4, -a5 + v4*v6, a4 + v5*v6, -v4**2 - v5**2, 0, 0, 0, 0, 0, 0], 249 | [0, 0, a3 - v1*v5 + v2*v4, -a2 - v1*v6 + v3*v4, a4, a5 - v4*v6, -v5*v6, a6 + v4*v5, v5**2 - v6**2, v5*v6], 250 | [0, -a3 + v1*v5 - v2*v4, 0, a1 - v2*v6 + v3*v5, v4*v6, a4 + v5*v6, a5, -v4**2 + v6**2, a6 - v4*v5, -v4*v6], 251 | [0, a2 + v1*v6 - v3*v4, -a1 + v2*v6 - v3*v5, 0, -v4*v5, v4**2 - v5**2, v4*v5, a4 - v5*v6, a5 + v4*v6, a6] 252 | ]) 253 | # fmt: on 254 | 255 | 256 | # def joint_body_regressor(mjxmodel: mjx.Model, mjxdata: mjx.Data, bodyid) -> jpt.ArrayLike: 257 | # """mj_bodyRegressor returns a regressor for a single rigid body 258 | 259 | # This function calculates the regressor for a single rigid body in the MuJoCo model. 260 | # Given the index of body we compute the velocity and acceleration of the body and 261 | # then calculate the regressor using the Y_body function. 262 | 263 | # Args: 264 | # mj_model: MuJoCo model 265 | # mj_data: MuJoCo data 266 | # body_id: ID of the body 267 | 268 | # Returns: 269 | # npt.ArrayLike: regressor for the body 270 | # """ 271 | # velocity = object_velocity(mjxmodel, mjxdata, bodyid) 272 | # # accel = mjx. 273 | 274 | # # velocity = np.zeros(6) 275 | # # accel = np.zeros(6) 276 | # # _cross = np.zeros(3) 277 | 278 | # # mujoco.mj_objectVelocity(mj_model, mj_data, 2, body_id, velocity, 1) 279 | # # mujoco.mj_objectAcceleration(mj_model, mj_data, 2, body_id, accel, 1) 280 | 281 | # # v, w = velocity[3:], velocity[:3] 282 | # # # dv - classical acceleration, already contains g 283 | # # dv, dw = accel[3:], accel[:3] 284 | # # mujoco.mju_cross(_cross, w, v) 285 | 286 | # # # for floating base, this is already included in dv 287 | # # if mj_model.nq == mj_model.nv: 288 | # # dv -= _cross 289 | 290 | # # return body_regressor(v, w, dv, dw) 291 | -------------------------------------------------------------------------------- /mujoco_sysid/parameters.py: -------------------------------------------------------------------------------- 1 | import mujoco 2 | import numpy as np 3 | import numpy.typing as npt 4 | 5 | 6 | def skew(vector): 7 | return np.cross(np.eye(vector.size), vector.reshape(-1)) 8 | 9 | 10 | def get_dynamic_parameters(mjmodel, body_id) -> npt.ArrayLike: 11 | """Get the dynamic parameters \theta of a body 12 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 13 | 14 | Args: 15 | mjmodel (mujoco.MjModel): The mujoco model 16 | body_id (int): The id of the body 17 | 18 | Returns: 19 | npt.ArrayLike: theta of the body 20 | """ 21 | mass = mjmodel.body(body_id).mass[0] 22 | rc = mjmodel.body(body_id).ipos 23 | diag_inertia = mjmodel.body(body_id).inertia 24 | 25 | # get the orientation of the body 26 | r_flat = np.zeros(9) 27 | mujoco.mju_quat2Mat(r_flat, mjmodel.body(body_id).iquat) 28 | 29 | R = r_flat.reshape(3, 3) 30 | 31 | shift = mass * skew(rc) @ skew(rc) 32 | mjinertia = R @ np.diag(diag_inertia) @ R.T - shift 33 | 34 | upper_triangular = np.array( 35 | [ 36 | mjinertia[0, 0], 37 | mjinertia[0, 1], 38 | mjinertia[1, 1], 39 | mjinertia[0, 2], 40 | mjinertia[1, 2], 41 | mjinertia[2, 2], 42 | ] 43 | ) 44 | 45 | return np.concatenate([[mass], mass * rc, upper_triangular]) 46 | 47 | 48 | def set_dynamic_parameters(mjmodel, body_id, theta: npt.ArrayLike) -> None: 49 | """Set the dynamic parameters to a body 50 | 51 | Args: 52 | mjmodel (mujoco.MjModel): The mujoco model 53 | body_id (int): The id of the body 54 | theta (npt.ArrayLike): The dynamic parameters of the body 55 | """ 56 | 57 | mass = theta[0] 58 | rc = theta[1:4] / mass 59 | inertia = theta[4:] 60 | inertia_full = np.array( 61 | [ 62 | [inertia[0], inertia[1], inertia[3]], 63 | [inertia[1], inertia[2], inertia[4]], 64 | [inertia[3], inertia[4], inertia[5]], 65 | ] 66 | ) 67 | 68 | # shift the inertia 69 | inertia_full += mass * skew(rc) @ skew(rc) 70 | 71 | # eigen decomposition 72 | eigval, eigvec = np.linalg.eigh(inertia_full) 73 | R = eigvec 74 | diag_inertia = eigval 75 | 76 | # check if singular, then abort 77 | if np.any(np.isclose(diag_inertia, 0)): 78 | raise ValueError("Cannot deduce inertia matrix because RIR^T is singular.") 79 | 80 | # set the mass 81 | mjmodel.body(body_id).mass = np.array([mass]) 82 | mjmodel.body(body_id).ipos = rc 83 | 84 | # set the orientation 85 | mujoco.mju_mat2Quat(mjmodel.body(body_id).iquat, R.flatten()) 86 | 87 | # set the inertia 88 | mjmodel.body(body_id).inertia = diag_inertia 89 | -------------------------------------------------------------------------------- /mujoco_sysid/regressors.py: -------------------------------------------------------------------------------- 1 | """This module contains functions of the regressor representation of inverse dynamics for single rigid body: 2 | M_b * a_g + v x M_b * v = Y_b(v, a_g)*theta = f 3 | 4 | and multiple bodies in generilized coordinates: 5 | M(q) * ddq + h(q, dq) = Y(q, dq, ddq)*theta = tau 6 | 7 | with vector of inertial parameters (stacked for multiple bodies): 8 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 9 | 10 | A linear parametrization of inverse dynamics is pivotal in SysID for robotic systems. 11 | To the best of our knowledge, dedicated functions for this representation are not available in MuJoCo, 12 | prompting us to develop this prototype. 13 | 14 | References: 15 | - Traversaro, Silvio, et al. "Identification of fully physical consistent inertial parameters using optimization on manifolds." 16 | 2016 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS). IEEE, 2016. 17 | - Garofalo G, Ott C, Albu-Schäffer A. On the closed form computation of the dynamic matrices and their differentiations. 18 | In2013 IEEE/RSJ International Conference on Intelligent Robots and Systems 2013 Nov 3 (pp. 2364-2359). IEEE. 19 | """ # noqa 20 | 21 | import mujoco 22 | import numpy as np 23 | from numpy import typing as npt 24 | 25 | 26 | def body_regressor( 27 | v_lin: npt.ArrayLike, v_ang: npt.ArrayLike, a_lin: npt.ArrayLike, a_ang: npt.ArrayLike 28 | ) -> npt.ArrayLike: 29 | """Y_body returns a regressor for a single rigid body 30 | 31 | Newton-Euler equations for a rigid body are given by: 32 | M * a_g + v x M * v = f 33 | 34 | where: 35 | M is the spatial inertia matrix of the body 36 | a_g is the acceleration of the body 37 | v is the spatial velocity of the body 38 | f is the spatial force acting on the body 39 | 40 | The regressor is a matrix Y such that: 41 | Y \theta = f 42 | 43 | where: 44 | \theta is the vector of inertial parameters of the body (10 parameters) 45 | 46 | More expressive derivation is given here: 47 | https://colab.research.google.com/drive/1xFte2FT0nQ0ePs02BoOx4CmLLw5U-OUZ?usp=sharing 48 | 49 | Args: 50 | v_lin (npt.ArrayLike): linear velocity of the body 51 | v_ang (npt.ArrayLike): angular velocity of the body 52 | a_lin (npt.ArrayLike): linear acceleration of the body 53 | a_ang (npt.ArrayLike): angular acceleration of the body 54 | 55 | Returns: 56 | npt.ArrayLike: regressor for the body 57 | """ 58 | v1, v2, v3 = v_lin 59 | v4, v5, v6 = v_ang 60 | 61 | a1, a2, a3 = a_lin 62 | a4, a5, a6 = a_ang 63 | 64 | # fmt: off 65 | return np.array([ 66 | [a1 - v2*v6 + v3*v5, -v5**2 - v6**2, -a6 + v4*v5, a5 + v4*v6, 0, 0, 0, 0, 0, 0], 67 | [a2 + v1*v6 - v3*v4, a6 + v4*v5, -v4**2 - v6**2, -a4 + v5*v6, 0, 0, 0, 0, 0, 0], 68 | [a3 - v1*v5 + v2*v4, -a5 + v4*v6, a4 + v5*v6, -v4**2 - v5**2, 0, 0, 0, 0, 0, 0], 69 | [0, 0, a3 - v1*v5 + v2*v4, -a2 - v1*v6 + v3*v4, a4, a5 - v4*v6, -v5*v6, a6 + v4*v5, v5**2 - v6**2, v5*v6], 70 | [0, -a3 + v1*v5 - v2*v4, 0, a1 - v2*v6 + v3*v5, v4*v6, a4 + v5*v6, a5, -v4**2 + v6**2, a6 - v4*v5, -v4*v6], 71 | [0, a2 + v1*v6 - v3*v4, -a1 + v2*v6 - v3*v5, 0, -v4*v5, v4**2 - v5**2, v4*v5, a4 - v5*v6, a5 + v4*v6, a6] 72 | ]) 73 | # fmt: on 74 | 75 | 76 | def joint_body_regressor(mj_model, mj_data, body_id) -> npt.ArrayLike: 77 | """mj_bodyRegressor returns a regressor for a single rigid body 78 | 79 | This function calculates the regressor for a single rigid body in the MuJoCo model. 80 | Given the index of body we compute the velocity and acceleration of the body and 81 | then calculate the regressor using the Y_body function. 82 | 83 | Args: 84 | mj_model: MuJoCo model 85 | mj_data: MuJoCo data 86 | body_id: ID of the body 87 | 88 | Returns: 89 | npt.ArrayLike: regressor for the body 90 | """ 91 | 92 | velocity = np.zeros(6) 93 | accel = np.zeros(6) 94 | _cross = np.zeros(3) 95 | 96 | mujoco.mj_objectVelocity(mj_model, mj_data, 2, body_id, velocity, 1) 97 | mujoco.mj_objectAcceleration(mj_model, mj_data, 2, body_id, accel, 1) 98 | 99 | v, w = velocity[3:], velocity[:3] 100 | # dv - classical acceleration, already contains g 101 | dv, dw = accel[3:], accel[:3] 102 | mujoco.mju_cross(_cross, w, v) 103 | 104 | # for floating base, this is already included in dv 105 | if mj_model.nq == mj_model.nv: 106 | dv -= _cross 107 | 108 | return body_regressor(v, w, dv, dw) 109 | 110 | 111 | def get_jacobian(mjmodel, mjdata, bodyid): 112 | R = mjdata.xmat[bodyid].reshape(3, 3) 113 | 114 | jac_lin, jac_rot = np.zeros((3, mjmodel.nv)), np.zeros((3, mjmodel.nv)) 115 | mujoco.mj_jacBody(mjmodel, mjdata, jac_lin, jac_rot, bodyid) 116 | 117 | return np.vstack([R.T @ jac_lin, R.T @ jac_rot]) 118 | 119 | 120 | def joint_torque_regressor(mj_model, mj_data) -> npt.ArrayLike: 121 | """mj_jointRegressor returns a regressor for the whole model 122 | 123 | This function calculates the regressor for the whole model in the MuJoCo model. 124 | 125 | This regressor is computed to use in joint-space calculations. It is a matrix that 126 | maps the inertial parameters of the bodies to the generalized forces. 127 | 128 | Newton-Euler equations for a rigid body are given by: 129 | M * a_g + v x M * v = f 130 | 131 | Expressing the spatial quantities in terms of the generalized quantities 132 | we can rewrite the equation for the system of bodies as: 133 | M * q_dot_dot + h = tau 134 | 135 | Where 136 | M is the mass matrix 137 | h is the bias term 138 | tau is the generalized forces 139 | 140 | Then, the regressor is a matrix Y such that: 141 | Y * theta = tau 142 | 143 | where: 144 | theta is the vector of inertial parameters of the bodies (10 parameters per body): 145 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 146 | 147 | 148 | Args: 149 | mj_model: MuJoCo model 150 | mj_data: MuJoCo data 151 | 152 | Returns: 153 | npt.ArrayLike: regressor for the whole model 154 | """ 155 | 156 | njoints = mj_model.njnt 157 | body_regressors = np.zeros((6 * njoints, njoints * 10)) 158 | col_jac = np.zeros((6 * njoints, mj_model.nv)) 159 | 160 | for i, body_id in enumerate(mj_model.jnt_bodyid): 161 | # calculate cody regressors 162 | body_regressors[6 * i : 6 * (i + 1), 10 * i : 10 * (i + 1)] = joint_body_regressor(mj_model, mj_data, body_id) 163 | 164 | col_jac[6 * i : 6 * i + 6, :] = get_jacobian(mj_model, mj_data, body_id) 165 | 166 | return col_jac.T @ body_regressors 167 | 168 | 169 | def body_energyRegressor( 170 | v: npt.ArrayLike, 171 | w: npt.ArrayLike, 172 | r: npt.ArrayLike, 173 | R: npt.ArrayLike, 174 | gravity: npt.ArrayLike | None = None, 175 | ) -> tuple[npt.ArrayLike, npt.ArrayLike]: 176 | """ 177 | Computes the kinetic and potential energy regressors for a single body. 178 | 179 | This function calculates the energy regressors for the whole model in the MuJoCo model. 180 | 181 | The energy regressors are computed to represent the kinetic and potential energy of the system 182 | as a linear combination of the inertial parameters of the bodies. 183 | 184 | The kinetic energy of a rigid body is given by: 185 | K = 0.5 * v^T * M * v 186 | 187 | The potential energy of a rigid body in a uniform gravitational field is given by: 188 | U = m * g^T * r 189 | 190 | where: 191 | v is the spatial velocity of the body 192 | M is the spatial inertia matrix of the body 193 | m is the mass of the body 194 | g is the gravitational acceleration vector 195 | r is the position of the body's center of mass 196 | 197 | The energy regressors Y_k and Y_u are matrices such that: 198 | K = Y_k * theta 199 | U = Y_u * theta 200 | 201 | Args: 202 | v (npt.ArrayLike): Linear velocity of the body. 203 | w (npt.ArrayLike): Angular velocity of the body. 204 | r (npt.ArrayLike): Position of the body. 205 | R (npt.ArrayLike): Rotation matrix of the body. 206 | gravity (npt.ArrayLike, optional): Gravity vector. Defaults to np.array([0, 0, 9.81]). 207 | 208 | Returns: 209 | tuple[npt.ArrayLike, npt.ArrayLike]: Kinetic and potential energy regressors. 210 | """ 211 | if gravity is None: 212 | gravity = np.array([0, 0, 9.81]) 213 | 214 | kinetic = np.array( 215 | [ 216 | 0.5 * (v[0] ** 2 + v[1] ** 2 + v[2] ** 2), 217 | -w[1] * v[2] + w[2] * v[1], 218 | w[0] * v[2] - w[2] * v[0], 219 | -w[0] * v[1] + w[1] * v[0], 220 | 0.5 * w[0] ** 2, 221 | w[0] * w[1], 222 | 0.5 * w[1] ** 2, 223 | w[0] * w[2], 224 | w[1] * w[2], 225 | 0.5 * w[2] ** 2, 226 | ] 227 | ) 228 | 229 | potential = np.array([*(gravity.T @ r).flatten(), *(gravity.T @ R).flatten(), 0, 0, 0, 0, 0, 0]) 230 | 231 | return kinetic, potential 232 | 233 | 234 | def energy_regressor(mj_model, mj_data) -> npt.ArrayLike: 235 | """ 236 | mj_energyRegressor returns kinetic, potential, and total energy regressors for the whole model. 237 | 238 | The energy regressors Y_k and Y_u are matrices such that: 239 | K = Y_k * theta 240 | U = Y_u * theta 241 | 242 | where: 243 | theta is the vector of inertial parameters of the bodies (10 parameters per body): 244 | theta = [m, h_x, h_y, h_z, I_xx, I_xy, I_yy, I_xz, I_yz, I_zz] 245 | 246 | The total energy regressor Y_e is simply the sum of Y_k and Y_u: 247 | E = K + U = Y_e * theta 248 | 249 | Args: 250 | mj_model: MuJoCo model 251 | mj_data: MuJoCo data 252 | body_offset (int, optional): Starting index of the body, useful when some dummy bodies are introduced. Defaults to 0. 253 | 254 | Returns: 255 | npt.ArrayLike: total energy regressor for the whole model 256 | """ # noqa 257 | 258 | njoints = mj_model.njnt 259 | energy_regressor = np.zeros(njoints * 10) 260 | velocity = np.zeros(6) 261 | 262 | for i, bodyid in enumerate(mj_model.jnt_bodyid): 263 | # same as jacobian @ qvel 264 | mujoco.mj_objectVelocity(mj_model, mj_data, 2, bodyid, velocity, 1) 265 | v, w = velocity[3:], velocity[:3] 266 | 267 | rotation = mj_data.xmat[bodyid].reshape(3, 3) 268 | position = mj_data.xpos[bodyid] 269 | 270 | kinetic, potential = body_energyRegressor(v, w, position, rotation) 271 | energy_regressor[10 * i : 10 * (i + 1)] = kinetic + potential 272 | 273 | return energy_regressor 274 | 275 | 276 | def potential_energy_bias(mjmodel): 277 | """ 278 | The bodies before the first joint are considered to be fixed in space. 279 | They are included in potential energy calculation, but not in the regressor. 280 | """ 281 | 282 | bias = 0 283 | for i in range(mjmodel.nbody): 284 | if i not in mjmodel.jnt_bodyid: 285 | bias += mjmodel.body(i).mass[0] * mjmodel.opt.gravity[2] * mjmodel.body(i).ipos[2] 286 | 287 | return bias 288 | -------------------------------------------------------------------------------- /mujoco_sysid/utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from quaternion import as_rotation_matrix, quaternion 3 | 4 | 5 | def muj2pin(qpos: np.ndarray, qvel: np.ndarray, qacc: np.ndarray | None = None) -> tuple: 6 | """ 7 | Converts Mujoco state to Pinocchio state by adjusting the quaternion representation and rotating the velocity. 8 | 9 | This function assumes that the quaternion representation of orientation in the Mujoco state uses a scalar-first 10 | format (w, x, y, z), while the Pinocchio state uses a scalar-last format (x, y, z, w). It also rotates the linear 11 | velocity from the world frame to the local frame. 12 | 13 | Args: 14 | qpos (numpy.ndarray): Mujoco qpos array, which includes position and orientation. 15 | qvel (numpy.ndarray): Mujoco qvel array, which includes linear and angular velocity. 16 | qacc (numpy.ndarray, optional): Mujoco qacc array, which includes linear and angular acceleration. 17 | 18 | Returns: 19 | tuple: A tuple containing two numpy.ndarrays: 20 | - pin_pos (numpy.ndarray): Pinocchio qpos array, with adjusted quaternion and position. 21 | - pin_vel (numpy.ndarray): Pinocchio qvel array, with velocity rotated to the local frame. 22 | 23 | """ 24 | # Copy the position and velocity to avoid modifying the original arrays 25 | pin_pos = qpos.copy() 26 | pin_vel = qvel.copy() 27 | pin_acc = qacc.copy() if qacc is not None else None 28 | 29 | if len(pin_pos) == len(qvel): 30 | if qacc is None: 31 | return pin_pos, pin_vel 32 | 33 | return pin_pos, pin_vel, pin_acc 34 | 35 | # Create a quaternion object from the Mujoco orientation (scalar-first) 36 | q = quaternion(*pin_pos[3:7]) 37 | # Obtain the corresponding rotation matrix 38 | R = as_rotation_matrix(q) 39 | # Rotate the world frame linear velocity to the local frame 40 | pin_vel[0:3] = R.T @ pin_vel[0:3] 41 | 42 | # Reorder quaternion from scalar-first (Mujoco) to scalar-last (Pinocchio) 43 | pin_pos[[3, 4, 5, 6]] = pin_pos[[4, 5, 6, 3]] 44 | 45 | if qacc is not None: 46 | # Rotate the world frame linear acceleration to the local frame 47 | pin_acc[0:3] = R.T @ pin_acc[0:3] 48 | return pin_pos, pin_vel, pin_acc 49 | 50 | return pin_pos, pin_vel 51 | 52 | 53 | def pin2muj(pin_pos: np.ndarray, pin_vel: np.ndarray, pin_acc: np.ndarray | None = None) -> tuple: 54 | """ 55 | Converts Pinocchio state to Mujoco state by adjusting the quaternion representation and rotating the velocity. 56 | 57 | This function assumes that the quaternion representation of orientation in the Pinocchio state uses a scalar-last 58 | format (x, y, z, w), while the Mujoco state uses a scalar-first format (w, x, y, z). It also rotates the local 59 | frame linear velocity to the world frame. 60 | 61 | Args: 62 | pin_pos (numpy.ndarray): Pinocchio qpos array, which includes position and orientation. 63 | pin_vel (numpy.ndarray): Pinocchio qvel array, which includes linear and angular velocity. 64 | pin_acc (numpy.ndarray, optional): Pinocchio qacc array, which includes linear and angular acceleration. 65 | 66 | Returns: 67 | tuple: A tuple containing two numpy.ndarrays: 68 | - qpos (numpy.ndarray): Mujoco qpos array, with adjusted quaternion and position. 69 | - qvel (numpy.ndarray): Mujoco qvel array, with velocity rotated to the world frame. 70 | """ 71 | # Copy the position and velocity to avoid modifying the original arrays 72 | qpos = pin_pos.copy() 73 | qvel = pin_vel.copy() 74 | qacc = pin_acc.copy() if pin_acc is not None else None 75 | 76 | if len(pin_pos) == len(qvel): 77 | if qacc is None: 78 | return qpos, qvel 79 | 80 | return qpos, qvel, qacc 81 | 82 | # Reorder quaternion from scalar-last (Pinocchio) to scalar-first (Mujoco) 83 | qpos[[3, 4, 5, 6]] = qpos[[6, 3, 4, 5]] 84 | 85 | # Create a quaternion object from the Pinocchio orientation (scalar-last) 86 | q = quaternion(*qpos[3:7]) 87 | # Obtain the corresponding rotation matrix 88 | R = as_rotation_matrix(q) 89 | # Rotate the local frame linear velocity to the world frame 90 | qvel[0:3] = R @ qvel[0:3] 91 | 92 | if qacc is not None: 93 | qacc[0:3] = R @ qacc[0:3] 94 | return qpos, qvel, qacc 95 | 96 | return qpos, qvel 97 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "mujoco_sysid" 3 | description = "MuJoCo System Identification tools" 4 | version = "0.2.1" 5 | authors = [ 6 | { name = "Lev Kozlov", email = "kozlov.l.a10@gmail.com" }, 7 | { name = "Simeon Nedelchev", email = "simkaned@gmail.com" }, 8 | ] 9 | dependencies = ["mujoco", "numpy-quaternion"] 10 | readme = "README.md" 11 | requires-python = ">=3.8" 12 | 13 | [project.optional-dependencies] 14 | examples = [ 15 | "ipykernel", 16 | "dm_control", 17 | "scipy", 18 | "matplotlib", 19 | "robot_descriptions", 20 | "mediapy", 21 | ] 22 | mjx = ["mujoco-mjx", "jax[cuda12]"] 23 | mjx_cpu = ["mujoco-mjx", "jax"] 24 | 25 | [project.urls] 26 | homepage = "https://github.com/lvjonok/mujoco-sysid" 27 | repository = "https://github.com/lvjonok/mujoco-sysid.git" 28 | 29 | [build-system] 30 | requires = ["setuptools>=43.0.0", "wheel"] 31 | build-backend = "setuptools.build_meta" 32 | 33 | [tool.setuptools.packages.find] 34 | where = ["."] 35 | include = ["mujoco_sysid*"] 36 | 37 | 38 | [tool.pytest.ini_options] 39 | pythonpath = ["."] 40 | 41 | [tool.ruff] 42 | select = [ 43 | "E", # pycodestyle errors 44 | "W", # pycodestyle warnings 45 | "F", # pyflakes 46 | "I", # isort 47 | "B", # flake8-bugbear 48 | "C4", # flake8-comprehensions 49 | "UP", # pyupgrade 50 | ] 51 | line-length = 120 52 | 53 | [tool.ruff.lint.per-file-ignores] 54 | "__init__.py" = ["F401"] 55 | -------------------------------------------------------------------------------- /tests/mjx/test_mjx_dynamics.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from mujoco_sysid import regressors 4 | from mujoco_sysid.mjx import regressors as regx 5 | 6 | np.random.seed(0) 7 | 8 | 9 | def test_body_regressor(): 10 | SAMPLES = 100 11 | 12 | for _ in range(SAMPLES): 13 | spatial_v, spatial_dv = np.random.rand(6), np.random.rand(6) 14 | 15 | mjY = regressors.body_regressor(spatial_v[:3], spatial_v[3:], spatial_dv[:3], spatial_dv[3:]) 16 | 17 | mjxY = regx.body_regressor(spatial_v[:3], spatial_v[3:], spatial_dv[:3], spatial_dv[3:]) 18 | 19 | assert np.allclose(mjY, mjxY, atol=1e-6) 20 | 21 | 22 | # def test_joint_body_regressor(): 23 | # import pinocchio as pin 24 | 25 | # # Test for fixed base manipulator 26 | # from robot_descriptions.z1_description import URDF_PATH 27 | # from robot_descriptions.z1_mj_description import MJCF_PATH 28 | 29 | # pinmodel = pin.buildModelFromUrdf(URDF_PATH) 30 | # pindata = pinmodel.createData() 31 | 32 | # mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 33 | # mjdata = mujoco.MjData(mjmodel) 34 | 35 | # SAMPLES = 10000 36 | 37 | # for _ in range(SAMPLES): 38 | # q, v, dv = np.random.rand(pinmodel.nq), np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv) 39 | # pin.rnea(pinmodel, pindata, q, v, dv) 40 | 41 | # mjdata.qpos[:] = q 42 | # mjdata.qvel[:] = v 43 | # mjdata.qacc[:] = dv 44 | # mujoco.mj_inverse(mjmodel, mjdata) 45 | 46 | # for jnt_id in range(pinmodel.njoints - 1): 47 | # pinY = pin.jointBodyRegressor(pinmodel, pindata, jnt_id + 1) 48 | # mjY = regressors.joint_body_regressor(mjmodel, mjdata, mjmodel.jnt_bodyid[jnt_id]) 49 | 50 | # assert np.allclose(mjY, pinY, atol=1e-6), f"Joint {jnt_id} failed. Pinocchio: {pinY}, Mujoco: {mjY}" 51 | 52 | 53 | # # Test for the floating base joint 54 | # from robot_descriptions.skydio_x2_description import URDF_PATH 55 | # from robot_descriptions.skydio_x2_mj_description import MJCF_PATH 56 | 57 | # pinmodel = pin.buildModelFromUrdf(URDF_PATH) 58 | # pindata = pinmodel.createData() 59 | 60 | # mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 61 | # mjdata = mujoco.MjData(mjmodel) 62 | 63 | # SAMPLES = 10000 64 | 65 | # for _ in range(SAMPLES): 66 | # q = np.random.rand(pinmodel.nq) 67 | # # normalize the quaternion 68 | # q[3:7] /= np.linalg.norm(q[3:7]) 69 | # v, dv = np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv) 70 | 71 | # pinq, pinv, pindv = muj2pin(q, v, dv) 72 | 73 | # mjdata.qpos[:] = q 74 | # mjdata.qvel[:] = v 75 | # mjdata.qacc[:] = dv 76 | # mujoco.mj_inverse(mjmodel, mjdata) 77 | 78 | # pin.rnea(pinmodel, pindata, pinq, pinv, pindv) 79 | 80 | # for jnt_id in range(pinmodel.njoints - 1): 81 | # pinY = pin.jointBodyRegressor(pinmodel, pindata, jnt_id + 1) 82 | 83 | # mjY = regressors.joint_body_regressor(mjmodel, mjdata, mjmodel.jnt_bodyid[jnt_id]) 84 | 85 | # assert np.allclose(mjY, pinY, atol=1e-6), f"Joint {jnt_id} failed. Norm diff: {np.linalg.norm(mjY - pinY)}" 86 | 87 | 88 | # def test_joint_torque_regressor(): 89 | # import pinocchio as pin 90 | 91 | # # Test for fixed base manipulator 92 | # from robot_descriptions.z1_description import URDF_PATH 93 | # from robot_descriptions.z1_mj_description import MJCF_PATH 94 | 95 | # pinmodel = pin.buildModelFromUrdf(URDF_PATH) 96 | # pindata = pinmodel.createData() 97 | 98 | # mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 99 | # mjdata = mujoco.MjData(mjmodel) 100 | 101 | # SAMPLES = 10000 102 | 103 | # for _ in range(SAMPLES): 104 | # q, v, dv = np.random.rand(pinmodel.nq), np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv) 105 | # pin.rnea(pinmodel, pindata, q, v, dv) 106 | 107 | # mjdata.qpos[:] = q 108 | # mjdata.qvel[:] = v 109 | # mjdata.qacc[:] = dv 110 | # mujoco.mj_inverse(mjmodel, mjdata) 111 | 112 | # pinY = pin.computeJointTorqueRegressor(pinmodel, pindata, q, v, dv) 113 | # mjY = regressors.joint_torque_regressor(mjmodel, mjdata) 114 | 115 | # assert np.allclose(mjY, pinY, atol=1e-6), f"Norm diff: {np.linalg.norm(mjY - pinY)}" 116 | 117 | 118 | if __name__ == "__main__": 119 | test_body_regressor() 120 | # test_joint_body_regressor() 121 | # test_joint_torque_regressor() 122 | -------------------------------------------------------------------------------- /tests/mjx/test_mjx_energy.py: -------------------------------------------------------------------------------- 1 | import jax 2 | import jax.numpy as jnp 3 | import mujoco 4 | import mujoco.mjx as mjx 5 | 6 | from mujoco_sysid import regressors as mj_reg 7 | from mujoco_sysid.mjx import parameters, regressors 8 | 9 | key = jax.random.PRNGKey(0) 10 | 11 | 12 | def test_energy_match_z1(): 13 | from robot_descriptions.z1_mj_description import MJCF_PATH 14 | 15 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 16 | # alter the model so it becomes mjx compatible 17 | mjmodel.dof_frictionloss = 0 18 | mjmodel.opt.integrator = 0 19 | # disable friction, contact and limits 20 | mjmodel.opt.disableflags |= ( 21 | mujoco.mjtDisableBit.mjDSBL_CONTACT 22 | | mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS 23 | | mujoco.mjtDisableBit.mjDSBL_LIMIT 24 | ) 25 | mjdata = mujoco.MjData(mjmodel) 26 | 27 | mjxmodel = mjx.put_model(mjmodel) 28 | assert len(mjxmodel.jnt_bodyid) == mjmodel.njnt, f"{len(mjxmodel.jnt_bodyid)} != {mjmodel.njnt}" 29 | # mjxdata = mjx.put_data(mjmodel, mjdata) 30 | 31 | # enable energy 32 | mjmodel.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_ENERGY 33 | 34 | # get vector of dynamic parameters 35 | theta = jnp.concatenate([parameters.get_dynamic_parameters(mjxmodel, i) for i in mjxmodel.jnt_bodyid]) 36 | assert theta.shape == (mjxmodel.njnt * 10,) 37 | 38 | # generate random samples of configuration and velocity 39 | N_SAMPLES = 10000 40 | sampled_q, sampled_v = ( 41 | jax.random.normal(key, (mjmodel.nq, N_SAMPLES)), 42 | jax.random.normal(key, (mjmodel.nv, N_SAMPLES)), 43 | ) 44 | 45 | # compute the expected energy for each sample 46 | expected_energy = [] 47 | for i in range(N_SAMPLES): 48 | mjdata.qpos[:] = sampled_q[:, i] 49 | mjdata.qvel[:] = sampled_v[:, i] 50 | 51 | mujoco.mj_step(mjmodel, mjdata) 52 | 53 | mj_en = mjdata.energy.copy() 54 | mj_en[0] += mj_reg.potential_energy_bias(mjmodel) 55 | 56 | expected_energy.append(jnp.sum(mj_en)) 57 | 58 | expected_energy = jnp.array(expected_energy) 59 | # compute the batched regressors using mjx and find the computed energy 60 | 61 | @jax.vmap 62 | def compute_energy(q, v): 63 | mjx_data = mjx.make_data(mjxmodel) 64 | mjx_data = mjx_data.replace(qpos=q, qvel=v) 65 | mjx_data = mjx.step(mjxmodel, mjx_data) 66 | 67 | # compute the regressor 68 | reg_en = regressors.energy_regressor(mjxmodel, mjx_data) 69 | computed_energy = jnp.dot(reg_en, theta) 70 | 71 | return computed_energy 72 | 73 | # compute the energy for each sample 74 | computed_energy = jax.jit(compute_energy)(sampled_q.T, sampled_v.T) 75 | 76 | # compare the computed energy with the expected energy 77 | assert jnp.allclose(expected_energy, computed_energy, atol=1e-5) 78 | 79 | 80 | def test_energy_match_skydio(): 81 | from robot_descriptions.skydio_x2_mj_description import MJCF_PATH 82 | 83 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 84 | # alter the model so it becomes mjx compatible 85 | mjmodel.dof_frictionloss = 0 86 | mjmodel.opt.integrator = 0 87 | # disable friction, contact and limits 88 | mjmodel.opt.disableflags |= ( 89 | mujoco.mjtDisableBit.mjDSBL_CONTACT 90 | | mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS 91 | | mujoco.mjtDisableBit.mjDSBL_LIMIT 92 | ) 93 | mjdata = mujoco.MjData(mjmodel) 94 | 95 | mjxmodel = mjx.put_model(mjmodel) 96 | assert len(mjxmodel.jnt_bodyid) == mjmodel.njnt, f"{len(mjxmodel.jnt_bodyid)} != {mjmodel.njnt}" 97 | # mjxdata = mjx.put_data(mjmodel, mjdata) 98 | 99 | # enable energy 100 | mjmodel.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_ENERGY 101 | 102 | # get vector of dynamic parameters 103 | theta = jnp.concatenate([parameters.get_dynamic_parameters(mjxmodel, i) for i in mjxmodel.jnt_bodyid]) 104 | assert theta.shape == (mjxmodel.njnt * 10,) 105 | 106 | # generate random samples of configuration and velocity 107 | N_SAMPLES = 10000 108 | sampled_q, sampled_v = ( 109 | jax.random.normal(key, (mjmodel.nq, N_SAMPLES)), 110 | jax.random.normal(key, (mjmodel.nv, N_SAMPLES)), 111 | ) 112 | 113 | # compute the expected energy for each sample 114 | expected_energy = [] 115 | for i in range(N_SAMPLES): 116 | mjdata.qpos[:] = sampled_q[:, i] 117 | mjdata.qvel[:] = sampled_v[:, i] 118 | 119 | mujoco.mj_step(mjmodel, mjdata) 120 | 121 | mj_en = mjdata.energy.copy() 122 | mj_en[0] += mj_reg.potential_energy_bias(mjmodel) 123 | 124 | expected_energy.append(jnp.sum(mj_en)) 125 | 126 | expected_energy = jnp.array(expected_energy) 127 | # compute the batched regressors using mjx and find the computed energy 128 | 129 | @jax.vmap 130 | def compute_energy(q, v): 131 | mjx_data = mjx.make_data(mjxmodel) 132 | mjx_data = mjx_data.replace(qpos=q, qvel=v) 133 | mjx_data = mjx.step(mjxmodel, mjx_data) 134 | 135 | # compute the regressor 136 | reg_en = regressors.energy_regressor(mjxmodel, mjx_data) 137 | computed_energy = jnp.dot(reg_en, theta) 138 | 139 | return computed_energy 140 | 141 | # compute the energy for each sample 142 | computed_energy = jax.jit(compute_energy)(sampled_q.T, sampled_v.T) 143 | 144 | # compare the computed energy with the expected energy 145 | assert jnp.allclose(expected_energy, computed_energy, atol=1e-5) 146 | 147 | 148 | if __name__ == "__main__": 149 | test_energy_match_z1() 150 | test_energy_match_skydio() 151 | -------------------------------------------------------------------------------- /tests/mjx/test_mjx_parameters.py: -------------------------------------------------------------------------------- 1 | import jax.numpy as jnp 2 | import mujoco 3 | from mujoco import mjx 4 | 5 | from mujoco_sysid.mjx import parameters 6 | 7 | 8 | def test_mjx_get_dynamic_parameters(): 9 | # get z1 manipulator model 10 | import pinocchio as pin 11 | from robot_descriptions.z1_description import URDF_PATH 12 | from robot_descriptions.z1_mj_description import MJCF_PATH 13 | 14 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 15 | # alter the model so it becomes mjx compatible 16 | mjmodel.dof_frictionloss = 0 17 | mjmodel.opt.integrator = 1 18 | 19 | mjxmodel = mjx.put_model(mjmodel) 20 | 21 | pinmodel = pin.buildModelFromUrdf(URDF_PATH) 22 | 23 | for i in range(6): 24 | pinparams = pinmodel.inertias[i].toDynamicParameters() 25 | mjparams = parameters.get_dynamic_parameters(mjxmodel, i + 1) 26 | 27 | # check that mass is close 28 | assert jnp.isclose(pinparams[0], mjparams[0]) 29 | 30 | # check that the center of mass is close 31 | assert jnp.allclose(pinparams[1:4], mjparams[1:4]) 32 | 33 | # check that the inertia matrix is close 34 | assert jnp.allclose(pinparams[4:], mjparams[4:], atol=1e-6), "Inertia matrix is not the same" 35 | 36 | 37 | def test_mjx_set_dynamic_parameters(): 38 | import numpy as np 39 | from robot_descriptions.z1_mj_description import MJCF_PATH 40 | 41 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 42 | # alter the model so it becomes mjx compatible 43 | mjmodel.dof_frictionloss = 0 44 | mjmodel.opt.integrator = 1 45 | 46 | mjxmodel = mjx.put_model(mjmodel) 47 | 48 | param = parameters.get_dynamic_parameters(mjxmodel, 1) 49 | parameters.set_dynamic_parameters(mjxmodel, 1, param) 50 | 51 | next_param = parameters.get_dynamic_parameters(mjxmodel, 1) 52 | 53 | # mass is the same 54 | assert np.isclose(param[0], next_param[0]) 55 | 56 | # center of mass is the same 57 | assert np.allclose(param[1:4], next_param[1:4]) 58 | 59 | # inertia matrix is the same 60 | assert np.allclose(param[4:], next_param[4:], atol=1e-6), "Inertia matrix is not the same" 61 | -------------------------------------------------------------------------------- /tests/test_dynamics.py: -------------------------------------------------------------------------------- 1 | import mujoco 2 | import numpy as np 3 | 4 | from mujoco_sysid import regressors 5 | from mujoco_sysid import parameters 6 | from mujoco_sysid.utils import muj2pin 7 | 8 | np.random.seed(0) 9 | 10 | 11 | def test_body_regressor(): 12 | import pinocchio as pin 13 | 14 | SAMPLES = 100 15 | 16 | for _ in range(SAMPLES): 17 | spatial_v, spatial_dv = np.random.rand(6), np.random.rand(6) 18 | 19 | pinY = pin.bodyRegressor(pin.Motion(spatial_v), pin.Motion(spatial_dv)) 20 | 21 | mjY = regressors.body_regressor(spatial_v[:3], spatial_v[3:], spatial_dv[:3], spatial_dv[3:]) 22 | 23 | assert np.allclose(mjY, pinY, atol=1e-6) 24 | 25 | 26 | def test_joint_body_regressor(): 27 | import pinocchio as pin 28 | 29 | # Test for fixed base manipulator 30 | from robot_descriptions.z1_description import URDF_PATH 31 | from robot_descriptions.z1_mj_description import MJCF_PATH 32 | 33 | pinmodel = pin.buildModelFromUrdf(URDF_PATH) 34 | pindata = pinmodel.createData() 35 | 36 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 37 | mjdata = mujoco.MjData(mjmodel) 38 | 39 | SAMPLES = 10000 40 | 41 | for _ in range(SAMPLES): 42 | q, v, dv = np.random.rand(pinmodel.nq), np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv) 43 | pin.rnea(pinmodel, pindata, q, v, dv) 44 | 45 | mjdata.qpos[:] = q 46 | mjdata.qvel[:] = v 47 | mjdata.qacc[:] = dv 48 | mujoco.mj_inverse(mjmodel, mjdata) 49 | mujoco.mj_rnePostConstraint(mjmodel, mjdata) 50 | 51 | for jnt_id in range(pinmodel.njoints - 1): 52 | pinY = pin.jointBodyRegressor(pinmodel, pindata, jnt_id + 1) 53 | mjY = regressors.joint_body_regressor(mjmodel, mjdata, mjmodel.jnt_bodyid[jnt_id]) 54 | 55 | assert np.allclose(mjY, pinY, atol=1e-6), f"Joint {jnt_id} failed. Pinocchio: {pinY}, Mujoco: {mjY}" 56 | 57 | # Test for the floating base joint 58 | from robot_descriptions.skydio_x2_description import URDF_PATH 59 | from robot_descriptions.skydio_x2_mj_description import MJCF_PATH 60 | 61 | pinmodel = pin.buildModelFromUrdf(URDF_PATH) 62 | pindata = pinmodel.createData() 63 | 64 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 65 | mjdata = mujoco.MjData(mjmodel) 66 | 67 | SAMPLES = 10000 68 | 69 | for _ in range(SAMPLES): 70 | q = np.random.rand(pinmodel.nq) 71 | # normalize the quaternion 72 | q[3:7] /= np.linalg.norm(q[3:7]) 73 | v, dv = np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv) 74 | 75 | pinq, pinv, pindv = muj2pin(q, v, dv) 76 | 77 | mjdata.qpos[:] = q 78 | mjdata.qvel[:] = v 79 | mjdata.qacc[:] = dv 80 | mujoco.mj_inverse(mjmodel, mjdata) 81 | mujoco.mj_rnePostConstraint(mjmodel, mjdata) 82 | 83 | pin.rnea(pinmodel, pindata, pinq, pinv, pindv) 84 | 85 | for jnt_id in range(pinmodel.njoints - 1): 86 | pinY = pin.jointBodyRegressor(pinmodel, pindata, jnt_id + 1) 87 | 88 | mjY = regressors.joint_body_regressor(mjmodel, mjdata, mjmodel.jnt_bodyid[jnt_id]) 89 | 90 | assert np.allclose(mjY, pinY, atol=1e-6), f"Joint {jnt_id} failed. Norm diff: {np.linalg.norm(mjY - pinY)}" 91 | 92 | 93 | def test_joint_torque_regressor(): 94 | import pinocchio as pin 95 | 96 | # Test for fixed base manipulator 97 | from robot_descriptions.z1_description import URDF_PATH 98 | from robot_descriptions.z1_mj_description import MJCF_PATH 99 | 100 | pinmodel = pin.buildModelFromUrdf(URDF_PATH) 101 | pindata = pinmodel.createData() 102 | 103 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 104 | mjdata = mujoco.MjData(mjmodel) 105 | 106 | SAMPLES = 10000 107 | 108 | theta = np.concatenate([parameters.get_dynamic_parameters(mjmodel, i) for i in mjmodel.jnt_bodyid]) 109 | 110 | for _ in range(SAMPLES): 111 | q, v, dv = np.random.rand(pinmodel.nq), np.random.rand(pinmodel.nv), np.random.rand(pinmodel.nv) 112 | pin.rnea(pinmodel, pindata, q, v, dv) 113 | 114 | mjdata.qpos[:] = q 115 | mjdata.qvel[:] = v 116 | mjdata.qacc[:] = dv 117 | mujoco.mj_inverse(mjmodel, mjdata) 118 | mujoco.mj_rnePostConstraint(mjmodel, mjdata) 119 | 120 | pinY = pin.computeJointTorqueRegressor(pinmodel, pindata, q, v, dv) 121 | mjY = regressors.joint_torque_regressor(mjmodel, mjdata) 122 | 123 | tau = pin.rnea(pinmodel, pindata, q, v, dv) 124 | 125 | assert np.allclose(mjY @ theta, tau, atol=1e-6), f"Norm diff: {np.linalg.norm(mjY @ theta - tau)}" 126 | assert np.allclose(mjY, pinY, atol=1e-6), f"Norm diff: {np.linalg.norm(mjY - pinY)}" 127 | 128 | 129 | if __name__ == "__main__": 130 | test_body_regressor() 131 | test_joint_body_regressor() 132 | test_joint_torque_regressor() 133 | -------------------------------------------------------------------------------- /tests/test_energy.py: -------------------------------------------------------------------------------- 1 | import mujoco 2 | import numpy as np 3 | 4 | from mujoco_sysid import parameters, regressors 5 | 6 | 7 | def test_energy_match_z1(): 8 | from robot_descriptions.z1_mj_description import MJCF_PATH 9 | 10 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 11 | # enable energy 12 | mjmodel.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_ENERGY 13 | # disable friction, contact and limits 14 | mjmodel.opt.disableflags |= ( 15 | mujoco.mjtDisableBit.mjDSBL_CONTACT 16 | | mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS 17 | | mujoco.mjtDisableBit.mjDSBL_LIMIT 18 | ) 19 | mjdata = mujoco.MjData(mjmodel) 20 | 21 | # generate random samples of configuration and velocity 22 | N_SAMPLES = 1000 23 | for _ in range(N_SAMPLES): 24 | q, v = np.random.randn(mjmodel.nq), np.zeros(mjmodel.nv) 25 | 26 | mjdata.qpos[:] = q.copy() 27 | mjdata.qvel[:] = v.copy() 28 | 29 | mujoco.mj_step(mjmodel, mjdata) 30 | 31 | mj_en = mjdata.energy.copy() 32 | mj_en[0] += regressors.potential_energy_bias(mjmodel) 33 | 34 | # energy we expect and computed through simulator 35 | expected_energy = np.sum(mj_en) 36 | 37 | # get vector of dynamic parameters 38 | theta = np.concatenate([parameters.get_dynamic_parameters(mjmodel, i) for i in mjmodel.jnt_bodyid]) 39 | # compute regressor of the total energy 40 | reg_en = regressors.energy_regressor(mjmodel, mjdata) 41 | 42 | computed_energy = reg_en @ theta 43 | 44 | assert np.isclose( 45 | expected_energy, computed_energy, atol=1e-6 46 | ), f"Expected energy: {expected_energy}, Computed energy: {computed_energy}" 47 | 48 | 49 | def test_energy_match_skydio(): 50 | from robot_descriptions.skydio_x2_mj_description import MJCF_PATH 51 | 52 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 53 | # enable energy 54 | mjmodel.opt.enableflags |= mujoco.mjtEnableBit.mjENBL_ENERGY 55 | # disable friction, contact and limits 56 | mjmodel.opt.disableflags |= ( 57 | mujoco.mjtDisableBit.mjDSBL_CONTACT 58 | | mujoco.mjtDisableBit.mjDSBL_FRICTIONLOSS 59 | | mujoco.mjtDisableBit.mjDSBL_LIMIT 60 | ) 61 | mjdata = mujoco.MjData(mjmodel) 62 | 63 | # generate random samples of configuration and velocity 64 | N_SAMPLES = 1000 65 | for _ in range(N_SAMPLES): 66 | q, v = np.random.randn(mjmodel.nq), np.zeros(mjmodel.nv) 67 | 68 | mjdata.qpos[:] = q.copy() 69 | mjdata.qvel[:] = v.copy() 70 | 71 | mujoco.mj_step(mjmodel, mjdata) 72 | 73 | mj_en = mjdata.energy.copy() 74 | mj_en[0] += regressors.potential_energy_bias(mjmodel) 75 | 76 | # energy we expect and computed through simulator 77 | expected_energy = np.sum(mj_en) 78 | 79 | # get vector of dynamic parameters 80 | theta = np.concatenate([parameters.get_dynamic_parameters(mjmodel, i) for i in mjmodel.jnt_bodyid]) 81 | # compute regressor of the total energy 82 | reg_en = regressors.energy_regressor(mjmodel, mjdata) 83 | 84 | computed_energy = reg_en @ theta 85 | 86 | assert np.isclose( 87 | expected_energy, computed_energy, atol=1e-6 88 | ), f"Expected energy: {expected_energy}, Computed energy: {computed_energy}" 89 | -------------------------------------------------------------------------------- /tests/test_parameters.py: -------------------------------------------------------------------------------- 1 | import mujoco 2 | import numpy as np 3 | 4 | from mujoco_sysid import parameters 5 | 6 | 7 | def test_get_dynamic_parameters(): 8 | # get z1 manipulator model 9 | import pinocchio as pin 10 | from robot_descriptions.z1_description import URDF_PATH 11 | from robot_descriptions.z1_mj_description import MJCF_PATH 12 | 13 | mjmodel = mujoco.MjModel.from_xml_path(MJCF_PATH) 14 | pinmodel = pin.buildModelFromUrdf(URDF_PATH) 15 | 16 | for i in range(6): 17 | pinparams = pinmodel.inertias[i].toDynamicParameters() 18 | mjparams = parameters.get_dynamic_parameters(mjmodel, i + 1) 19 | 20 | # check that mass is close 21 | assert np.isclose(pinparams[0], mjparams[0]) 22 | 23 | # check that the center of mass is close 24 | assert np.allclose(pinparams[1:4], mjparams[1:4]) 25 | 26 | # check that the inertia matrix is close 27 | assert np.allclose(pinparams[4:], mjparams[4:], atol=1e-6), "Inertia matrix is not the same" 28 | 29 | 30 | def test_set_dynamic_parameters(): 31 | import numpy as np 32 | from robot_descriptions.z1_mj_description import MJCF_PATH 33 | 34 | # Load the model 35 | model = mujoco.MjModel.from_xml_path(MJCF_PATH) 36 | 37 | param = parameters.get_dynamic_parameters(model, 1) 38 | parameters.set_dynamic_parameters(model, 1, param) 39 | 40 | next_param = parameters.get_dynamic_parameters(model, 1) 41 | 42 | # mass is the same 43 | assert np.isclose(param[0], next_param[0]) 44 | 45 | # center of mass is the same 46 | assert np.allclose(param[1:4], next_param[1:4]) 47 | 48 | # inertia matrix is the same 49 | assert np.allclose(param[4:], next_param[4:], atol=1e-6), "Inertia matrix is not the same" 50 | --------------------------------------------------------------------------------