From 146a5c08f7f46c84da9d8e9088f286043679d19f Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Wed, 4 Feb 2026 20:10:06 -0800 Subject: [PATCH 1/2] System identification toolbox for MuJoCo. This resulted from a lengthy collaboration with @kevinzakka, @jonathanembleyriches, @nimrod-gileadi, @gizemozd, @quagla, and @yuval. --- python/build_requirements.txt | 663 +++++++++++++++++ python/mujoco/sysid/README.md | 637 ++++++++++++++++ python/mujoco/sysid/__init__.py | 72 ++ python/mujoco/sysid/_src/__init__.py | 0 python/mujoco/sysid/_src/io.py | 57 ++ python/mujoco/sysid/_src/model_modifier.py | 507 +++++++++++++ python/mujoco/sysid/_src/optimize.py | 236 ++++++ python/mujoco/sysid/_src/parameter.py | 606 +++++++++++++++ python/mujoco/sysid/_src/plotting.py | 692 +++++++++++++++++ python/mujoco/sysid/_src/residual.py | 408 ++++++++++ python/mujoco/sysid/_src/signal_modifier.py | 244 ++++++ python/mujoco/sysid/_src/signal_transform.py | 257 +++++++ python/mujoco/sysid/_src/timeseries.py | 700 ++++++++++++++++++ python/mujoco/sysid/_src/trajectory.py | 501 +++++++++++++ python/mujoco/sysid/py.typed | 0 python/mujoco/sysid/report/builder.py | 104 +++ python/mujoco/sysid/report/defaults.py | 376 ++++++++++ .../mujoco/sysid/report/sections/__init__.py | 0 python/mujoco/sysid/report/sections/base.py | 47 ++ .../sysid/report/sections/covariance.py | 152 ++++ python/mujoco/sysid/report/sections/group.py | 43 ++ .../mujoco/sysid/report/sections/insights.py | 95 +++ .../report/sections/optimization_trace.py | 428 +++++++++++ .../report/sections/parameter_distribution.py | 398 ++++++++++ .../sysid/report/sections/parameters.py | 218 ++++++ python/mujoco/sysid/report/sections/row.py | 46 ++ .../mujoco/sysid/report/sections/signals.py | 206 ++++++ python/mujoco/sysid/report/sections/video.py | 172 +++++ .../sysid/report/templates/covariance.html | 204 +++++ .../mujoco/sysid/report/templates/group.html | 13 + .../sysid/report/templates/insights.html | 90 +++ .../mujoco/sysid/report/templates/layout.html | 699 +++++++++++++++++ .../report/templates/optimization_trace.html | 20 + .../templates/parameter_confidence.html | 3 + .../report/templates/parameters_table.html | 147 ++++ .../sysid/report/templates/plot_generic.html | 8 + python/mujoco/sysid/report/templates/row.html | 13 + .../sysid/report/templates/signals.html | 3 + .../mujoco/sysid/report/templates/video.html | 11 + python/mujoco/sysid/report/utils.py | 41 + python/mujoco/sysid/tests/__init__.py | 0 python/mujoco/sysid/tests/conftest.py | 219 ++++++ python/mujoco/sysid/tests/test_integration.py | 237 ++++++ .../mujoco/sysid/tests/test_model_modifier.py | 122 +++ python/mujoco/sysid/tests/test_parameter.py | 149 ++++ python/mujoco/sysid/tests/test_signal.py | 377 ++++++++++ python/mujoco/sysid/tests/test_timeseries.py | 332 +++++++++ python/mujoco/sysid/tests/test_trajectory.py | 146 ++++ python/pyproject.toml | 12 + 49 files changed, 10711 insertions(+) create mode 100644 python/mujoco/sysid/README.md create mode 100644 python/mujoco/sysid/__init__.py create mode 100644 python/mujoco/sysid/_src/__init__.py create mode 100644 python/mujoco/sysid/_src/io.py create mode 100644 python/mujoco/sysid/_src/model_modifier.py create mode 100644 python/mujoco/sysid/_src/optimize.py create mode 100644 python/mujoco/sysid/_src/parameter.py create mode 100644 python/mujoco/sysid/_src/plotting.py create mode 100644 python/mujoco/sysid/_src/residual.py create mode 100644 python/mujoco/sysid/_src/signal_modifier.py create mode 100644 python/mujoco/sysid/_src/signal_transform.py create mode 100644 python/mujoco/sysid/_src/timeseries.py create mode 100644 python/mujoco/sysid/_src/trajectory.py create mode 100644 python/mujoco/sysid/py.typed create mode 100644 python/mujoco/sysid/report/builder.py create mode 100644 python/mujoco/sysid/report/defaults.py create mode 100644 python/mujoco/sysid/report/sections/__init__.py create mode 100644 python/mujoco/sysid/report/sections/base.py create mode 100644 python/mujoco/sysid/report/sections/covariance.py create mode 100644 python/mujoco/sysid/report/sections/group.py create mode 100644 python/mujoco/sysid/report/sections/insights.py create mode 100644 python/mujoco/sysid/report/sections/optimization_trace.py create mode 100644 python/mujoco/sysid/report/sections/parameter_distribution.py create mode 100644 python/mujoco/sysid/report/sections/parameters.py create mode 100644 python/mujoco/sysid/report/sections/row.py create mode 100644 python/mujoco/sysid/report/sections/signals.py create mode 100644 python/mujoco/sysid/report/sections/video.py create mode 100644 python/mujoco/sysid/report/templates/covariance.html create mode 100644 python/mujoco/sysid/report/templates/group.html create mode 100644 python/mujoco/sysid/report/templates/insights.html create mode 100644 python/mujoco/sysid/report/templates/layout.html create mode 100644 python/mujoco/sysid/report/templates/optimization_trace.html create mode 100644 python/mujoco/sysid/report/templates/parameter_confidence.html create mode 100644 python/mujoco/sysid/report/templates/parameters_table.html create mode 100644 python/mujoco/sysid/report/templates/plot_generic.html create mode 100644 python/mujoco/sysid/report/templates/row.html create mode 100644 python/mujoco/sysid/report/templates/signals.html create mode 100644 python/mujoco/sysid/report/templates/video.html create mode 100644 python/mujoco/sysid/report/utils.py create mode 100644 python/mujoco/sysid/tests/__init__.py create mode 100644 python/mujoco/sysid/tests/conftest.py create mode 100644 python/mujoco/sysid/tests/test_integration.py create mode 100644 python/mujoco/sysid/tests/test_model_modifier.py create mode 100644 python/mujoco/sysid/tests/test_parameter.py create mode 100644 python/mujoco/sysid/tests/test_signal.py create mode 100644 python/mujoco/sysid/tests/test_timeseries.py create mode 100644 python/mujoco/sysid/tests/test_trajectory.py diff --git a/python/build_requirements.txt b/python/build_requirements.txt index d886424e..45dbdc3d 100644 --- a/python/build_requirements.txt +++ b/python/build_requirements.txt @@ -101,3 +101,666 @@ iniconfig==2.0.0 \ --hash=sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374 pluggy==1.5.0 \ --hash=sha256:44e1ad92c8ca002de6377e165f3e0f1be63266ab4d554740532335b9d75ea669 + +# sysid optional-dependency direct deps +imageio==2.37.2 \ + --hash=sha256:0212ef2727ac9caa5ca4b2c75ae89454312f440a756fcfc8ef1993e718f50f8a \ + --hash=sha256:ad9adfb20335d718c03de457358ed69f141021a333c40a53e57273d8a5bd0b9b +imageio-ffmpeg==0.6.0 \ + --hash=sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a \ + --hash=sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2 \ + --hash=sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc \ + --hash=sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61 \ + --hash=sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742 \ + --hash=sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282 \ + --hash=sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755 +jinja2==3.1.6 \ + --hash=sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d \ + --hash=sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67 +matplotlib==3.10.8 \ + --hash=sha256:00270d217d6b20d14b584c521f810d60c5c78406dc289859776550df837dcda7 \ + --hash=sha256:0a33deb84c15ede243aead39f77e990469fff93ad1521163305095b77b72ce4a \ + --hash=sha256:113bb52413ea508ce954a02c10ffd0d565f9c3bc7f2eddc27dfe1731e71c7b5f \ + --hash=sha256:12d90df9183093fcd479f4172ac26b322b1248b15729cb57f42f71f24c7e37a3 \ + --hash=sha256:15d30132718972c2c074cd14638c7f4592bd98719e2308bccea40e0538bc0cb5 \ + --hash=sha256:18821ace09c763ec93aef5eeff087ee493a24051936d7b9ebcad9662f66501f9 \ + --hash=sha256:1ae029229a57cd1e8fe542485f27e7ca7b23aa9e8944ddb4985d0bc444f1eca2 \ + --hash=sha256:2299372c19d56bcd35cf05a2738308758d32b9eaed2371898d8f5bd33f084aa3 \ + --hash=sha256:238b7ce5717600615c895050239ec955d91f321c209dd110db988500558e70d6 \ + --hash=sha256:24d50994d8c5816ddc35411e50a86ab05f575e2530c02752e02538122613371f \ + --hash=sha256:25d380fe8b1dc32cf8f0b1b448470a77afb195438bafdf1d858bfb876f3edf7b \ + --hash=sha256:2c1998e92cd5999e295a731bcb2911c75f597d937341f3030cc24ef2733d78a8 \ + --hash=sha256:2cf5bd12cecf46908f286d7838b2abc6c91cda506c0445b8223a7c19a00df008 \ + --hash=sha256:32f8dce744be5569bebe789e46727946041199030db8aeb2954d26013a0eb26b \ + --hash=sha256:37b3c1cc42aa184b3f738cfa18c1c1d72fd496d85467a6cf7b807936d39aa656 \ + --hash=sha256:3a48a78d2786784cc2413e57397981fb45c79e968d99656706018d6e62e57958 \ + --hash=sha256:3ab4aabc72de4ff77b3ec33a6d78a68227bf1123465887f9905ba79184a1cc04 \ + --hash=sha256:3c624e43ed56313651bc18a47f838b60d7b8032ed348911c54906b130b20071b \ + --hash=sha256:3f2e409836d7f5ac2f1c013110a4d50b9f7edc26328c108915f9075d7d7a91b6 \ + --hash=sha256:3f5c3e4da343bba819f0234186b9004faba952cc420fbc522dc4e103c1985908 \ + --hash=sha256:41703cc95688f2516b480f7f339d8851a6035f18e100ee6a32bc0b8536a12a9c \ + --hash=sha256:495672de149445ec1b772ff2c9ede9b769e3cb4f0d0aa7fa730d7f59e2d4e1c1 \ + --hash=sha256:4cf267add95b1c88300d96ca837833d4112756045364f5c734a2276038dae27d \ + --hash=sha256:56271f3dac49a88d7fca5060f004d9d22b865f743a12a23b1e937a0be4818ee1 \ + --hash=sha256:595ba4d8fe983b88f0eec8c26a241e16d6376fe1979086232f481f8f3f67494c \ + --hash=sha256:5f62550b9a30afde8c1c3ae450e5eb547d579dd69b25c2fc7a1c67f934c1717a \ + --hash=sha256:646d95230efb9ca614a7a594d4fcacde0ac61d25e37dd51710b36477594963ce \ + --hash=sha256:64fcc24778ca0404ce0cb7b6b77ae1f4c7231cdd60e6778f999ee05cbd581b9a \ + --hash=sha256:6be43b667360fef5c754dda5d25a32e6307a03c204f3c0fc5468b78fa87b4160 \ + --hash=sha256:6da7c2ce169267d0d066adcf63758f0604aa6c3eebf67458930f9d9b79ad1db1 \ + --hash=sha256:83d282364ea9f3e52363da262ce32a09dfe241e4080dcedda3c0db059d3c1f11 \ + --hash=sha256:9153c3292705be9f9c64498a8872118540c3f4123d1a1c840172edf262c8be4a \ + --hash=sha256:99eefd13c0dc3b3c1b4d561c1169e65fe47aab7b8158754d7c084088e2329466 \ + --hash=sha256:a0a7f52498f72f13d4a25ea70f35f4cb60642b466cbb0a9be951b5bc3f45a486 \ + --hash=sha256:a2b336e2d91a3d7006864e0990c83b216fcdca64b5a6484912902cef87313d78 \ + --hash=sha256:a48f2b74020919552ea25d222d5cc6af9ca3f4eb43a93e14d068457f545c2a17 \ + --hash=sha256:ad3d9833a64cf48cc4300f2b406c3d0f4f4724a91c0bd5640678a6ba7c102077 \ + --hash=sha256:b44d07310e404ba95f8c25aa5536f154c0a8ec473303535949e52eb71d0a1565 \ + --hash=sha256:b53285e65d4fa4c86399979e956235deb900be5baa7fc1218ea67fbfaeaadd6f \ + --hash=sha256:b5a2b97dbdc7d4f353ebf343744f1d1f1cca8aa8bfddb4262fcf4306c3761d50 \ + --hash=sha256:b9a5ca4ac220a0cdd1ba6bcba3608547117d30468fefce49bb26f55c1a3d5c58 \ + --hash=sha256:bab485bcf8b1c7d2060b4fcb6fc368a9e6f4cd754c9c2fea281f4be21df394a2 \ + --hash=sha256:c108a1d6fa78a50646029cb6d49808ff0fc1330fda87fa6f6250c6b5369b6645 \ + --hash=sha256:d56a1efd5bfd61486c8bc968fa18734464556f0fb8e51690f4ac25d85cbbbbc2 \ + --hash=sha256:d9050fee89a89ed57b4fb2c1bfac9a3d0c57a0d55aed95949eedbc42070fea39 \ + --hash=sha256:dd80ecb295460a5d9d260df63c43f4afbdd832d725a531f008dad1664f458adf \ + --hash=sha256:e8ea3e2d4066083e264e75c829078f9e149fa119d27e19acd503de65e0b13149 \ + --hash=sha256:eb3823f11823deade26ce3b9f40dcb4a213da7a670013929f31d5f5ed1055b22 \ + --hash=sha256:ee40c27c795bda6a5292e9cff9890189d32f7e3a0bf04e0e3c9430c4a00c37df \ + --hash=sha256:efb30e3baaea72ce5928e32bab719ab4770099079d66726a62b11b1ef7273be4 \ + --hash=sha256:f254d118d14a7f99d616271d6c3c27922c092dac11112670b157798b89bf4933 \ + --hash=sha256:f89c151aab2e2e23cb3fe0acad1e8b82841fd265379c4cecd0f3fcb34c15e0f6 \ + --hash=sha256:f97aeb209c3d2511443f8797e3e5a569aebb040d4f8bc79aa3ee78a8fb9e3dd8 \ + --hash=sha256:f9b587c9c7274c1613a30afabf65a272114cd6cdbe67b3406f818c79d7ab2e2a \ + --hash=sha256:fb061f596dad3a0f52b60dc6a5dec4a0c300dec41e058a7efe09256188d170b7 +plotly==6.5.2 \ + --hash=sha256:7478555be0198562d1435dee4c308268187553cc15516a2f4dd034453699e393 \ + --hash=sha256:91757653bd9c550eeea2fa2404dba6b85d1e366d54804c340b2c874e5a7eb4a4 +pyyaml==6.0.3 \ + --hash=sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c \ + --hash=sha256:0150219816b6a1fa26fb4699fb7daa9caf09eb1999f3b70fb6e786805e80375a \ + --hash=sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3 \ + --hash=sha256:02ea2dfa234451bbb8772601d7b8e426c2bfa197136796224e50e35a78777956 \ + --hash=sha256:0f29edc409a6392443abf94b9cf89ce99889a1dd5376d94316ae5145dfedd5d6 \ + --hash=sha256:10892704fc220243f5305762e276552a0395f7beb4dbf9b14ec8fd43b57f126c \ + --hash=sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65 \ + --hash=sha256:1d37d57ad971609cf3c53ba6a7e365e40660e3be0e5175fa9f2365a379d6095a \ + --hash=sha256:1ebe39cb5fc479422b83de611d14e2c0d3bb2a18bbcb01f229ab3cfbd8fee7a0 \ + --hash=sha256:214ed4befebe12df36bcc8bc2b64b396ca31be9304b8f59e25c11cf94a4c033b \ + --hash=sha256:2283a07e2c21a2aa78d9c4442724ec1eb15f5e42a723b99cb3d822d48f5f7ad1 \ + --hash=sha256:22ba7cfcad58ef3ecddc7ed1db3409af68d023b7f940da23c6c2a1890976eda6 \ + --hash=sha256:27c0abcb4a5dac13684a37f76e701e054692a9b2d3064b70f5e4eb54810553d7 \ + --hash=sha256:28c8d926f98f432f88adc23edf2e6d4921ac26fb084b028c733d01868d19007e \ + --hash=sha256:2e71d11abed7344e42a8849600193d15b6def118602c4c176f748e4583246007 \ + --hash=sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310 \ + --hash=sha256:37503bfbfc9d2c40b344d06b2199cf0e96e97957ab1c1b546fd4f87e53e5d3e4 \ + --hash=sha256:3c5677e12444c15717b902a5798264fa7909e41153cdf9ef7ad571b704a63dd9 \ + --hash=sha256:3ff07ec89bae51176c0549bc4c63aa6202991da2d9a6129d7aef7f1407d3f295 \ + --hash=sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea \ + --hash=sha256:418cf3f2111bc80e0933b2cd8cd04f286338bb88bdc7bc8e6dd775ebde60b5e0 \ + --hash=sha256:44edc647873928551a01e7a563d7452ccdebee747728c1080d881d68af7b997e \ + --hash=sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac \ + --hash=sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9 \ + --hash=sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7 \ + --hash=sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35 \ + --hash=sha256:5498cd1645aa724a7c71c8f378eb29ebe23da2fc0d7a08071d89469bf1d2defb \ + --hash=sha256:5cf4e27da7e3fbed4d6c3d8e797387aaad68102272f8f9752883bc32d61cb87b \ + --hash=sha256:5e0b74767e5f8c593e8c9b5912019159ed0533c70051e9cce3e8b6aa699fcd69 \ + --hash=sha256:5ed875a24292240029e4483f9d4a4b8a1ae08843b9c54f43fcc11e404532a8a5 \ + --hash=sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b \ + --hash=sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c \ + --hash=sha256:6344df0d5755a2c9a276d4473ae6b90647e216ab4757f8426893b5dd2ac3f369 \ + --hash=sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd \ + --hash=sha256:652cb6edd41e718550aad172851962662ff2681490a8a711af6a4d288dd96824 \ + --hash=sha256:66291b10affd76d76f54fad28e22e51719ef9ba22b29e1d7d03d6777a9174198 \ + --hash=sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065 \ + --hash=sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c \ + --hash=sha256:79005a0d97d5ddabfeeea4cf676af11e647e41d81c9a7722a193022accdb6b7c \ + --hash=sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764 \ + --hash=sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196 \ + --hash=sha256:8098f252adfa6c80ab48096053f512f2321f0b998f98150cea9bd23d83e1467b \ + --hash=sha256:850774a7879607d3a6f50d36d04f00ee69e7fc816450e5f7e58d7f17f1ae5c00 \ + --hash=sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac \ + --hash=sha256:8da9669d359f02c0b91ccc01cac4a67f16afec0dac22c2ad09f46bee0697eba8 \ + --hash=sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e \ + --hash=sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28 \ + --hash=sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3 \ + --hash=sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5 \ + --hash=sha256:9c57bb8c96f6d1808c030b1687b9b5fb476abaa47f0db9c0101f5e9f394e97f4 \ + --hash=sha256:9c7708761fccb9397fe64bbc0395abcae8c4bf7b0eac081e12b809bf47700d0b \ + --hash=sha256:9f3bfb4965eb874431221a3ff3fdcddc7e74e3b07799e0e84ca4a0f867d449bf \ + --hash=sha256:a33284e20b78bd4a18c8c2282d549d10bc8408a2a7ff57653c0cf0b9be0afce5 \ + --hash=sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702 \ + --hash=sha256:b30236e45cf30d2b8e7b3e85881719e98507abed1011bf463a8fa23e9c3e98a8 \ + --hash=sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788 \ + --hash=sha256:b865addae83924361678b652338317d1bd7e79b1f4596f96b96c77a5a34b34da \ + --hash=sha256:b8bb0864c5a28024fac8a632c443c87c5aa6f215c0b126c449ae1a150412f31d \ + --hash=sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc \ + --hash=sha256:bdb2c67c6c1390b63c6ff89f210c8fd09d9a1217a465701eac7316313c915e4c \ + --hash=sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba \ + --hash=sha256:c2514fceb77bc5e7a2f7adfaa1feb2fb311607c9cb518dbc378688ec73d8292f \ + --hash=sha256:c3355370a2c156cffb25e876646f149d5d68f5e0a3ce86a5084dd0b64a994917 \ + --hash=sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5 \ + --hash=sha256:d0eae10f8159e8fdad514efdc92d74fd8d682c933a6dd088030f3834bc8e6b26 \ + --hash=sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f \ + --hash=sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b \ + --hash=sha256:eda16858a3cab07b80edaf74336ece1f986ba330fdb8ee0d6c0d68fe82bc96be \ + --hash=sha256:ee2922902c45ae8ccada2c5b501ab86c36525b883eff4255313a253a3160861c \ + --hash=sha256:efd7b85f94a6f21e4932043973a7ba2613b059c4a000551892ac9f1d11f5baf3 \ + --hash=sha256:f7057c9a337546edc7973c0d3ba84ddcdf0daa14533c2065749c9075001090e6 \ + --hash=sha256:fa160448684b4e94d80416c0fa4aac48967a969efe22931448d853ada8baf926 \ + --hash=sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0 +scipy==1.17.0 \ + --hash=sha256:00fb5f8ec8398ad90215008d8b6009c9db9fa924fd4c7d6be307c6f945f9cd73 \ + --hash=sha256:031121914e295d9791319a1875444d55079885bbae5bdc9c5e0f2ee5f09d34ff \ + --hash=sha256:0937a0b0d8d593a198cededd4c439a0ea216a3f36653901ea1f3e4be949056f8 \ + --hash=sha256:0cf46c8013fec9d3694dc572f0b54100c28405d55d3e2cb15e2895b25057996e \ + --hash=sha256:0d5018a57c24cb1dd828bcf51d7b10e65986d549f52ef5adb6b4d1ded3e32a57 \ + --hash=sha256:130d12926ae34399d157de777472bf82e9061c60cc081372b3118edacafe1d00 \ + --hash=sha256:13c4096ac6bc31d706018f06a49abe0485f96499deb82066b94d19b02f664209 \ + --hash=sha256:13e861634a2c480bd237deb69333ac79ea1941b94568d4b0efa5db5e263d4fd1 \ + --hash=sha256:1f9586a58039d7229ce77b52f8472c972448cded5736eaf102d5658bbac4c269 \ + --hash=sha256:1ff269abf702f6c7e67a4b7aad981d42871a11b9dd83c58d2d2ea624efbd1088 \ + --hash=sha256:255c0da161bd7b32a6c898e7891509e8a9289f0b1c6c7d96142ee0d2b114c2ea \ + --hash=sha256:2591060c8e648d8b96439e111ac41fd8342fdeff1876be2e19dea3fe8930454e \ + --hash=sha256:272a9f16d6bb4667e8b50d25d71eddcc2158a214df1b566319298de0939d2ab7 \ + --hash=sha256:2abd71643797bd8a106dff97894ff7869eeeb0af0f7a5ce02e4227c6a2e9d6fd \ + --hash=sha256:2b531f57e09c946f56ad0b4a3b2abee778789097871fc541e267d2eca081cff1 \ + --hash=sha256:30509da9dbec1c2ed8f168b8d8aa853bc6723fede1dbc23c7d43a56f5ab72a67 \ + --hash=sha256:33af70d040e8af9d5e7a38b5ed3b772adddd281e3062ff23fec49e49681c38cf \ + --hash=sha256:357ca001c6e37601066092e7c89cca2f1ce74e2a520ca78d063a6d2201101df2 \ + --hash=sha256:3625c631a7acd7cfd929e4e31d2582cf00f42fcf06011f59281271746d77e061 \ + --hash=sha256:363ad4ae2853d88ebcde3ae6ec46ccca903ea9835ee8ba543f12f575e7b07e4e \ + --hash=sha256:40052543f7bbe921df4408f46003d6f01c6af109b9e2c8a66dd1cf6cf57f7d5d \ + --hash=sha256:423ca1f6584fc03936972b5f7c06961670dbba9f234e71676a7c7ccf938a0d61 \ + --hash=sha256:474da16199f6af66601a01546144922ce402cb17362e07d82f5a6cf8f963e449 \ + --hash=sha256:4e00562e519c09da34c31685f6acc3aa384d4d50604db0f245c14e1b4488bfa2 \ + --hash=sha256:5194c445d0a1c7a6c1a4a4681b6b7c71baad98ff66d96b949097e7513c9d6742 \ + --hash=sha256:5fb10d17e649e1446410895639f3385fd2bf4c3c7dfc9bea937bddcbc3d7b9ba \ + --hash=sha256:65ec32f3d32dfc48c72df4291345dae4f048749bc8d5203ee0a3f347f96c5ce6 \ + --hash=sha256:6680f2dfd4f6182e7d6db161344537da644d1cf85cf293f015c60a17ecf08752 \ + --hash=sha256:6e886000eb4919eae3a44f035e63f0fd8b651234117e8f6f29bad1cd26e7bc45 \ + --hash=sha256:7204fddcbec2fe6598f1c5fdf027e9f259106d05202a959a9f1aecf036adc9f6 \ + --hash=sha256:819fc26862b4b3c73a60d486dbb919202f3d6d98c87cf20c223511429f2d1a97 \ + --hash=sha256:8547e7c57f932e7354a2319fab613981cde910631979f74c9b542bb167a8b9db \ + --hash=sha256:85b0ac3ad17fa3be50abd7e69d583d98792d7edc08367e01445a1e2076005379 \ + --hash=sha256:87b411e42b425b84777718cc41516b8a7e0795abfa8e8e1d573bf0ef014f0812 \ + --hash=sha256:88c22af9e5d5a4f9e027e26772cc7b5922fab8bcc839edb3ae33de404feebd9e \ + --hash=sha256:9244608d27eafe02b20558523ba57f15c689357c85bdcfe920b1828750aa26eb \ + --hash=sha256:979c3a0ff8e5ba254d45d59ebd38cde48fce4f10b5125c680c7a4bfe177aab07 \ + --hash=sha256:9eeb9b5f5997f75507814ed9d298ab23f62cf79f5a3ef90031b1ee2506abdb5b \ + --hash=sha256:9fad7d3578c877d606b1150135c2639e9de9cecd3705caa37b66862977cc3e72 \ + --hash=sha256:a38c3337e00be6fd8a95b4ed66b5d988bac4ec888fd922c2ea9fe5fb1603dd67 \ + --hash=sha256:aabf057c632798832f071a8dde013c2e26284043934f53b00489f1773b33527e \ + --hash=sha256:c17514d11b78be8f7e6331b983a65a7f5ca1fd037b95e27b280921fe5606286a \ + --hash=sha256:c5e8647f60679790c2f5c76be17e2e9247dc6b98ad0d3b065861e082c56e078d \ + --hash=sha256:cacbaddd91fcffde703934897c5cd2c7cb0371fac195d383f4e1f1c5d3f3bd04 \ + --hash=sha256:d7425fcafbc09a03731e1bc05581f5fad988e48c6a861f441b7ab729a49a55ea \ + --hash=sha256:dac97a27520d66c12a34fd90a4fe65f43766c18c0d6e1c0a80f114d2260080e4 \ + --hash=sha256:dbf133ced83889583156566d2bdf7a07ff89228fe0c0cb727f777de92092ec6b \ + --hash=sha256:e8c0b331c2c1f531eb51f1b4fc9ba709521a712cce58f1aa627bc007421a5306 \ + --hash=sha256:eb2651271135154aa24f6481cbae5cc8af1f0dd46e6533fb7b56aa9727b6a232 \ + --hash=sha256:ebb7446a39b3ae0fe8f416a9a3fdc6fba3f11c634f680f16a239c5187bc487c0 \ + --hash=sha256:ec0827aa4d36cb79ff1b81de898e948a51ac0b9b1c43e4a372c0508c38c0f9a3 \ + --hash=sha256:edce1a1cf66298cccdc48a1bdf8fb10a3bf58e8b58d6c3883dd1530e103f87c0 \ + --hash=sha256:eec3842ec9ac9de5917899b277428886042a93db0b227ebbe3a333b64ec7643d \ + --hash=sha256:ef28d815f4d2686503e5f4f00edc387ae58dfd7a2f42e348bb53359538f01558 \ + --hash=sha256:f2a4942b0f5f7c23c7cd641a0ca1955e2ae83dedcff537e3a0259096635e186b \ + --hash=sha256:f3cd947f20fe17013d401b64e857c6b2da83cae567adbb75b9dcba865abc66d8 \ + --hash=sha256:f603d8a5518c7426414d1d8f82e253e454471de682ce5e39c29adb0df1efb86b \ + --hash=sha256:f7df7941d71314e60a481e02d5ebcb3f0185b8d799c70d03d8258f6c80f3d467 \ + --hash=sha256:f9eb55bb97d00f8b7ab95cb64f873eb0bf54d9446264d9f3609130381233483f \ + --hash=sha256:fc02c37a5639ee67d8fb646ffded6d793c06c5622d36b35cfa8fe5ececb8f042 \ + --hash=sha256:fe508b5690e9eaaa9467fc047f833af58f1152ae51a0d0aed67aa5801f4dd7d6 +tabulate==0.9.0 \ + --hash=sha256:0095b12bf5966de529c0feb1fa08671671b3368eec77d7ef7ab114be2c068b3c \ + --hash=sha256:024ca478df22e9340661486f85298cff5f6dcdba14f3813e8830015b9ed1948f + +# Transitive dependencies of sysid optional deps +contourpy==1.3.3 \ + --hash=sha256:023b44101dfe49d7d53932be418477dba359649246075c996866106da069af69 \ + --hash=sha256:07ce5ed73ecdc4a03ffe3e1b3e3c1166db35ae7584be76f65dbbe28a7791b0cc \ + --hash=sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880 \ + --hash=sha256:0bf67e0e3f482cb69779dd3061b534eb35ac9b17f163d851e2a547d56dba0a3a \ + --hash=sha256:0c1fc238306b35f246d61a1d416a627348b5cf0648648a031e14bb8705fcdfe8 \ + --hash=sha256:13b68d6a62db8eafaebb8039218921399baf6e47bf85006fd8529f2a08ef33fc \ + --hash=sha256:15ff10bfada4bf92ec8b31c62bf7c1834c244019b4a33095a68000d7075df470 \ + --hash=sha256:177fb367556747a686509d6fef71d221a4b198a3905fe824430e5ea0fda54eb5 \ + --hash=sha256:1cadd8b8969f060ba45ed7c1b714fe69185812ab43bd6b86a9123fe8f99c3263 \ + --hash=sha256:1fd43c3be4c8e5fd6e4f2baeae35ae18176cf2e5cced681cca908addf1cdd53b \ + --hash=sha256:22e9b1bd7a9b1d652cd77388465dc358dafcd2e217d35552424aa4f996f524f5 \ + --hash=sha256:23416f38bfd74d5d28ab8429cc4d63fa67d5068bd711a85edb1c3fb0c3e2f381 \ + --hash=sha256:283edd842a01e3dcd435b1c5116798d661378d83d36d337b8dde1d16a5fc9ba3 \ + --hash=sha256:2a2a8b627d5cc6b7c41a4beff6c5ad5eb848c88255fda4a8745f7e901b32d8e4 \ + --hash=sha256:2b7e9480ffe2b0cd2e787e4df64270e3a0440d9db8dc823312e2c940c167df7e \ + --hash=sha256:322ab1c99b008dad206d406bb61d014cf0174df491ae9d9d0fac6a6fda4f977f \ + --hash=sha256:33c82d0138c0a062380332c861387650c82e4cf1747aaa6938b9b6516762e772 \ + --hash=sha256:348ac1f5d4f1d66d3322420f01d42e43122f43616e0f194fc1c9f5d830c5b286 \ + --hash=sha256:3519428f6be58431c56581f1694ba8e50626f2dd550af225f82fb5f5814d2a42 \ + --hash=sha256:3c30273eb2a55024ff31ba7d052dde990d7d8e5450f4bbb6e913558b3d6c2301 \ + --hash=sha256:3d1a3799d62d45c18bafd41c5fa05120b96a28079f2393af559b843d1a966a77 \ + --hash=sha256:451e71b5a7d597379ef572de31eeb909a87246974d960049a9848c3bc6c41bf7 \ + --hash=sha256:459c1f020cd59fcfe6650180678a9993932d80d44ccde1fa1868977438f0b411 \ + --hash=sha256:4d00e655fcef08aba35ec9610536bfe90267d7ab5ba944f7032549c55a146da1 \ + --hash=sha256:4debd64f124ca62069f313a9cb86656ff087786016d76927ae2cf37846b006c9 \ + --hash=sha256:4feffb6537d64b84877da813a5c30f1422ea5739566abf0bd18065ac040e120a \ + --hash=sha256:50ed930df7289ff2a8d7afeb9603f8289e5704755c7e5c3bbd929c90c817164b \ + --hash=sha256:51e79c1f7470158e838808d4a996fa9bac72c498e93d8ebe5119bc1e6becb0db \ + --hash=sha256:556dba8fb6f5d8742f2923fe9457dbdd51e1049c4a43fd3986a0b14a1d815fc6 \ + --hash=sha256:598c3aaece21c503615fd59c92a3598b428b2f01bfb4b8ca9c4edeecc2438620 \ + --hash=sha256:5ed3657edf08512fc3fe81b510e35c2012fbd3081d2e26160f27ca28affec989 \ + --hash=sha256:626d60935cf668e70a5ce6ff184fd713e9683fb458898e4249b63be9e28286ea \ + --hash=sha256:644a6853d15b2512d67881586bd03f462c7ab755db95f16f14d7e238f2852c67 \ + --hash=sha256:655456777ff65c2c548b7c454af9c6f33f16c8884f11083244b5819cc214f1b5 \ + --hash=sha256:66c8a43a4f7b8df8b71ee1840e4211a3c8d93b214b213f590e18a1beca458f7d \ + --hash=sha256:6afc576f7b33cf00996e5c1102dc2a8f7cc89e39c0b55df93a0b78c1bd992b36 \ + --hash=sha256:6c3d53c796f8647d6deb1abe867daeb66dcc8a97e8455efa729516b997b8ed99 \ + --hash=sha256:709a48ef9a690e1343202916450bc48b9e51c049b089c7f79a267b46cffcdaa1 \ + --hash=sha256:70f9aad7de812d6541d29d2bbf8feb22ff7e1c299523db288004e3157ff4674e \ + --hash=sha256:8153b8bfc11e1e4d75bcb0bff1db232f9e10b274e0929de9d608027e0d34ff8b \ + --hash=sha256:87acf5963fc2b34825e5b6b048f40e3635dd547f590b04d2ab317c2619ef7ae8 \ + --hash=sha256:88df9880d507169449d434c293467418b9f6cbe82edd19284aa0409e7fdb933d \ + --hash=sha256:929ddf8c4c7f348e4c0a5a3a714b5c8542ffaa8c22954862a46ca1813b667ee7 \ + --hash=sha256:92d9abc807cf7d0e047b95ca5d957cf4792fcd04e920ca70d48add15c1a90ea7 \ + --hash=sha256:95b181891b4c71de4bb404c6621e7e2390745f887f2a026b2d99e92c17892339 \ + --hash=sha256:9e999574eddae35f1312c2b4b717b7885d4edd6cb46700e04f7f02db454e67c1 \ + --hash=sha256:a15459b0f4615b00bbd1e91f1b9e19b7e63aea7483d03d804186f278c0af2659 \ + --hash=sha256:a22738912262aa3e254e4f3cb079a95a67132fc5a063890e224393596902f5a4 \ + --hash=sha256:ab2fd90904c503739a75b7c8c5c01160130ba67944a7b77bbf36ef8054576e7f \ + --hash=sha256:ab3074b48c4e2cf1a960e6bbeb7f04566bf36b1861d5c9d4d8ac04b82e38ba20 \ + --hash=sha256:afe5a512f31ee6bd7d0dda52ec9864c984ca3d66664444f2d72e0dc4eb832e36 \ + --hash=sha256:b08a32ea2f8e42cf1d4be3169a98dd4be32bafe4f22b6c4cb4ba810fa9e5d2cb \ + --hash=sha256:b20c7c9a3bf701366556e1b1984ed2d0cedf999903c51311417cf5f591d8c78d \ + --hash=sha256:b2e8faa0ed68cb29af51edd8e24798bb661eac3bd9f65420c1887b6ca89987c8 \ + --hash=sha256:b7301b89040075c30e5768810bc96a8e8d78085b47d8be6e4c3f5a0b4ed478a0 \ + --hash=sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b \ + --hash=sha256:ca0fdcd73925568ca027e0b17ab07aad764be4706d0a925b89227e447d9737b7 \ + --hash=sha256:ca658cd1a680a5c9ea96dc61cdbae1e85c8f25849843aa799dfd3cb370ad4fbe \ + --hash=sha256:cbedb772ed74ff5be440fa8eee9bd49f64f6e3fc09436d9c7d8f1c287b121d77 \ + --hash=sha256:cd5dfcaeb10f7b7f9dc8941717c6c2ade08f587be2226222c12b25f0483ed497 \ + --hash=sha256:cf9022ef053f2694e31d630feaacb21ea24224be1c3ad0520b13d844274614fd \ + --hash=sha256:d002b6f00d73d69333dac9d0b8d5e84d9724ff9ef044fd63c5986e62b7c9e1b1 \ + --hash=sha256:d06bb1f751ba5d417047db62bca3c8fde202b8c11fb50742ab3ab962c81e8216 \ + --hash=sha256:d304906ecc71672e9c89e87c4675dc5c2645e1f4269a5063b99b0bb29f232d13 \ + --hash=sha256:e4e6b05a45525357e382909a4c1600444e2a45b4795163d3b22669285591c1ae \ + --hash=sha256:e74a9a0f5e3fff48fb5a7f2fd2b9b70a3fe014a67522f79b7cca4c0c7e43c9ae \ + --hash=sha256:ea37e7b45949df430fe649e5de8351c423430046a2af20b1c1961cae3afcda77 \ + --hash=sha256:f64836de09927cba6f79dcd00fdd7d5329f3fccc633468507079c829ca4db4e3 \ + --hash=sha256:fd6ec6be509c787f1caf6b247f0b1ca598bef13f4ddeaa126b7658215529ba0f \ + --hash=sha256:fd907ae12cd483cd83e414b12941c632a969171bf90fc937d0c9f268a31cafff \ + --hash=sha256:fd914713266421b7536de2bfa8181aa8c699432b6763a0ea64195ebe28bff6a9 \ + --hash=sha256:fde6c716d51c04b1c25d0b90364d0be954624a0ee9d60e23e850e8d48353d07a +cycler==0.12.1 \ + --hash=sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30 \ + --hash=sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c +fonttools==4.61.1 \ + --hash=sha256:0de30bfe7745c0d1ffa2b0b7048fb7123ad0d71107e10ee090fa0b16b9452e87 \ + --hash=sha256:10d88e55330e092940584774ee5e8a6971b01fc2f4d3466a1d6c158230880796 \ + --hash=sha256:11f35ad7805edba3aac1a3710d104592df59f4b957e30108ae0ba6c10b11dd75 \ + --hash=sha256:15acc09befd16a0fb8a8f62bc147e1a82817542d72184acca9ce6e0aeda9fa6d \ + --hash=sha256:17d2bf5d541add43822bcf0c43d7d847b160c9bb01d15d5007d84e2217aaa371 \ + --hash=sha256:2180f14c141d2f0f3da43f3a81bc8aa4684860f6b0e6f9e165a4831f24e6a23b \ + --hash=sha256:21e7c8d76f62ab13c9472ccf74515ca5b9a761d1bde3265152a6dc58700d895b \ + --hash=sha256:41a7170d042e8c0024703ed13b71893519a1a6d6e18e933e3ec7507a2c26a4b2 \ + --hash=sha256:41ed4b5ec103bd306bb68f81dc166e77409e5209443e5773cb4ed837bcc9b0d3 \ + --hash=sha256:497c31ce314219888c0e2fce5ad9178ca83fe5230b01a5006726cdf3ac9f24d9 \ + --hash=sha256:4c1b526c8d3f615a7b1867f38a9410849c8f4aef078535742198e942fba0e9bd \ + --hash=sha256:4d7092bb38c53bbc78e9255a59158b150bcdc115a1e3b3ce0b5f267dc35dd63c \ + --hash=sha256:4f5686e1fe5fce75d82d93c47a438a25bf0d1319d2843a926f741140b2b16e0c \ + --hash=sha256:58b0ee0ab5b1fc9921eccfe11d1435added19d6494dde14e323f25ad2bc30c56 \ + --hash=sha256:5ce02f38a754f207f2f06557523cd39a06438ba3aafc0639c477ac409fc64e37 \ + --hash=sha256:5fade934607a523614726119164ff621e8c30e8fa1ffffbbd358662056ba69f0 \ + --hash=sha256:5fe9fd43882620017add5eabb781ebfbc6998ee49b35bd7f8f79af1f9f99a958 \ + --hash=sha256:64102ca87e84261419c3747a0d20f396eb024bdbeb04c2bfb37e2891f5fadcb5 \ + --hash=sha256:664c5a68ec406f6b1547946683008576ef8b38275608e1cee6c061828171c118 \ + --hash=sha256:6675329885c44657f826ef01d9e4fb33b9158e9d93c537d84ad8399539bc6f69 \ + --hash=sha256:75c1a6dfac6abd407634420c93864a1e274ebc1c7531346d9254c0d8f6ca00f9 \ + --hash=sha256:75da8f28eff26defba42c52986de97b22106cb8f26515b7c22443ebc9c2d3261 \ + --hash=sha256:77efb033d8d7ff233385f30c62c7c79271c8885d5c9657d967ede124671bbdfb \ + --hash=sha256:78a7d3ab09dc47ac1a363a493e6112d8cabed7ba7caad5f54dbe2f08676d1b47 \ + --hash=sha256:7c7db70d57e5e1089a274cbb2b1fd635c9a24de809a231b154965d415d6c6d24 \ + --hash=sha256:8c56c488ab471628ff3bfa80964372fc13504ece601e0d97a78ee74126b2045c \ + --hash=sha256:91669ccac46bbc1d09e9273546181919064e8df73488ea087dcac3e2968df9ba \ + --hash=sha256:9b666a475a65f4e839d3d10473fad6d47e0a9db14a2f4a224029c5bfde58ad2c \ + --hash=sha256:9cfef3ab326780c04d6646f68d4b4742aae222e8b8ea1d627c74e38afcbc9d91 \ + --hash=sha256:a13fc8aeb24bad755eea8f7f9d409438eb94e82cf86b08fe77a03fbc8f6a96b1 \ + --hash=sha256:a75c301f96db737e1c5ed5fd7d77d9c34466de16095a266509e13da09751bd19 \ + --hash=sha256:a76d4cb80f41ba94a6691264be76435e5f72f2cb3cab0b092a6212855f71c2f6 \ + --hash=sha256:aed04cabe26f30c1647ef0e8fbb207516fd40fe9472e9439695f5c6998e60ac5 \ + --hash=sha256:b148b56f5de675ee16d45e769e69f87623a4944f7443850bf9a9376e628a89d2 \ + --hash=sha256:b501c862d4901792adaec7c25b1ecc749e2662543f68bb194c42ba18d6eec98d \ + --hash=sha256:b846a1fcf8beadeb9ea4f44ec5bdde393e2f1569e17d700bfc49cd69bde75881 \ + --hash=sha256:b931ae8f62db78861b0ff1ac017851764602288575d65b8e8ff1963fed419063 \ + --hash=sha256:c33ab3ca9d3ccd581d58e989d67554e42d8d4ded94ab3ade3508455fe70e65f7 \ + --hash=sha256:c6604b735bb12fef8e0efd5578c9fb5d3d8532d5001ea13a19cddf295673ee09 \ + --hash=sha256:d8db08051fc9e7d8bc622f2112511b8107d8f27cd89e2f64ec45e9825e8288da \ + --hash=sha256:d9203500f7c63545b4ce3799319fe4d9feb1a1b89b28d3cb5abd11b9dd64147e \ + --hash=sha256:dc492779501fa723b04d0ab1f5be046797fee17d27700476edc7ee9ae535a61e \ + --hash=sha256:e6bcdf33aec38d16508ce61fd81838f24c83c90a1d1b8c68982857038673d6b8 \ + --hash=sha256:e76ce097e3c57c4bcb67c5aa24a0ecdbd9f74ea9219997a707a4061fbe2707aa \ + --hash=sha256:eff1ac3cc66c2ac7cda1e64b4e2f3ffef474b7335f92fc3833fc632d595fcee6 \ + --hash=sha256:f3cb4a569029b9f291f88aafc927dd53683757e640081ca8c412781ea144565e \ + --hash=sha256:f79b168428351d11e10c5aeb61a74e1851ec221081299f4cf56036a95431c43a \ + --hash=sha256:fa646ecec9528bef693415c79a86e733c70a4965dd938e9a226b0fc64c9d2e6c \ + --hash=sha256:fe2efccb324948a11dd09d22136fe2ac8a97d6c1347cf0b58a911dcd529f66b7 \ + --hash=sha256:fff4f534200a04b4a36e7ae3cb74493afe807b517a09e99cb4faa89a34ed6ecd +kiwisolver==1.4.9 \ + --hash=sha256:0749fd8f4218ad2e851e11cc4dc05c7cbc0cbc4267bdfdb31782e65aace4ee9c \ + --hash=sha256:0763515d4df10edf6d06a3c19734e2566368980d21ebec439f33f9eb936c07b7 \ + --hash=sha256:0856e241c2d3df4efef7c04a1e46b1936b6120c9bcf36dd216e3acd84bc4fb21 \ + --hash=sha256:0a590506f303f512dff6b7f75fd2fd18e16943efee932008fe7140e5fa91d80e \ + --hash=sha256:0ab74e19f6a2b027ea4f845a78827969af45ce790e6cb3e1ebab71bdf9f215ff \ + --hash=sha256:0ae37737256ba2de764ddc12aed4956460277f00c4996d51a197e72f62f5eec7 \ + --hash=sha256:0e4e2bf29574a6a7b7f6cb5fa69293b9f96c928949ac4a53ba3f525dffb87f9c \ + --hash=sha256:15163165efc2f627eb9687ea5f3a28137217d217ac4024893d753f46bce9de26 \ + --hash=sha256:17680d737d5335b552994a2008fab4c851bcd7de33094a82067ef3a576ff02fa \ + --hash=sha256:1a12cf6398e8a0a001a059747a1cbf24705e18fe413bc22de7b3d15c67cffe3f \ + --hash=sha256:1b11d6a633e4ed84fc0ddafd4ebfd8ea49b3f25082c04ad12b8315c11d504dc1 \ + --hash=sha256:1fa333e8b2ce4d9660f2cda9c0e1b6bafcfb2457a9d259faa82289e73ec24891 \ + --hash=sha256:2327a4a30d3ee07d2fbe2e7933e8a37c591663b96ce42a00bc67461a87d7df77 \ + --hash=sha256:2405a7d98604b87f3fc28b1716783534b1b4b8510d8142adca34ee0bc3c87543 \ + --hash=sha256:2489e4e5d7ef9a1c300a5e0196e43d9c739f066ef23270607d45aba368b91f2d \ + --hash=sha256:24c175051354f4a28c5d6a31c93906dc653e2bf234e8a4bbfb964892078898ce \ + --hash=sha256:2635d352d67458b66fd0667c14cb1d4145e9560d503219034a18a87e971ce4f3 \ + --hash=sha256:2c1a4f57df73965f3f14df20b80ee29e6a7930a57d2d9e8491a25f676e197c60 \ + --hash=sha256:2c93f00dcba2eea70af2be5f11a830a742fe6b579a1d4e00f47760ef13be247a \ + --hash=sha256:39a219e1c81ae3b103643d2aedb90f1ef22650deb266ff12a19e7773f3e5f089 \ + --hash=sha256:3b3115b2581ea35bb6d1f24a4c90af37e5d9b49dcff267eeed14c3893c5b86ab \ + --hash=sha256:40092754720b174e6ccf9e845d0d8c7d8e12c3d71e7fc35f55f3813e96376f78 \ + --hash=sha256:412f287c55a6f54b0650bd9b6dce5aceddb95864a1a90c87af16979d37c89771 \ + --hash=sha256:464415881e4801295659462c49461a24fb107c140de781d55518c4b80cb6790f \ + --hash=sha256:497d05f29a1300d14e02e6441cf0f5ee81c1ff5a304b0d9fb77423974684e08b \ + --hash=sha256:4a2899935e724dd1074cb568ce7ac0dce28b2cd6ab539c8e001a8578eb106d14 \ + --hash=sha256:4a48a2ce79d65d363597ef7b567ce3d14d68783d2b2263d98db3d9477805ba32 \ + --hash=sha256:4d1d9e582ad4d63062d34077a9a1e9f3c34088a2ec5135b1f7190c07cf366527 \ + --hash=sha256:52a15b0f35dad39862d376df10c5230155243a2c1a436e39eb55623ccbd68185 \ + --hash=sha256:540c7c72324d864406a009d72f5d6856f49693db95d1fbb46cf86febef873634 \ + --hash=sha256:5656aa670507437af0207645273ccdfee4f14bacd7f7c67a4306d0dcaeaf6eed \ + --hash=sha256:5a0f2724dfd4e3b3ac5a82436a8e6fd16baa7d507117e4279b660fe8ca38a3a1 \ + --hash=sha256:60c439763a969a6af93b4881db0eed8fadf93ee98e18cbc35bc8da868d0c4f0c \ + --hash=sha256:61874cdb0a36016354853593cffc38e56fc9ca5aa97d2c05d3dcf6922cd55a11 \ + --hash=sha256:67bb8b474b4181770f926f7b7d2f8c0248cbcb78b660fdd41a47054b28d2a752 \ + --hash=sha256:720e05574713db64c356e86732c0f3c5252818d05f9df320f0ad8380641acea5 \ + --hash=sha256:72d0eb9fba308b8311685c2268cf7d0a0639a6cd027d8128659f72bdd8a024b4 \ + --hash=sha256:767c23ad1c58c9e827b649a9ab7809fd5fd9db266a9cf02b0e926ddc2c680d58 \ + --hash=sha256:77937e5e2a38a7b48eef0585114fe7930346993a88060d0bf886086d2aa49ef5 \ + --hash=sha256:7a08b491ec91b1d5053ac177afe5290adacf1f0f6307d771ccac5de30592d198 \ + --hash=sha256:7b4da0d01ac866a57dd61ac258c5607b4cd677f63abaec7b148354d2b2cdd536 \ + --hash=sha256:7cf974dd4e35fa315563ac99d6287a1024e4dc2077b8a7d7cd3d2fb65d283134 \ + --hash=sha256:84fd60810829c27ae375114cd379da1fa65e6918e1da405f356a775d49a62bcf \ + --hash=sha256:858e4c22fb075920b96a291928cb7dea5644e94c0ee4fcd5af7e865655e4ccf2 \ + --hash=sha256:85b5352f94e490c028926ea567fc569c52ec79ce131dadb968d3853e809518c2 \ + --hash=sha256:85bd218b5ecfbee8c8a82e121802dcb519a86044c9c3b2e4aef02fa05c6da370 \ + --hash=sha256:8a1f570ce4d62d718dce3f179ee78dac3b545ac16c0c04bb363b7607a949c0d1 \ + --hash=sha256:8fdca1def57a2e88ef339de1737a1449d6dbf5fab184c54a1fca01d541317154 \ + --hash=sha256:90f47e70293fc3688b71271100a1a5453aa9944a81d27ff779c108372cf5567b \ + --hash=sha256:92a2f997387a1b79a75e7803aa7ded2cfbe2823852ccf1ba3bcf613b62ae3197 \ + --hash=sha256:9928fe1eb816d11ae170885a74d074f57af3a0d65777ca47e9aeb854a1fba386 \ + --hash=sha256:9af39d6551f97d31a4deebeac6f45b156f9755ddc59c07b402c148f5dbb6482a \ + --hash=sha256:9cf554f21be770f5111a1690d42313e140355e687e05cf82cb23d0a721a64a48 \ + --hash=sha256:a30fd6fdef1430fd9e1ba7b3398b5ee4e2887783917a687d86ba69985fb08748 \ + --hash=sha256:a31d512c812daea6d8b3be3b2bfcbeb091dbb09177706569bcfc6240dcf8b41c \ + --hash=sha256:a5d0432ccf1c7ab14f9949eec60c5d1f924f17c037e9f8b33352fa05799359b8 \ + --hash=sha256:a60ea74330b91bd22a29638940d115df9dc00af5035a9a2a6ad9399ffb4ceca5 \ + --hash=sha256:ac5a486ac389dddcc5bef4f365b6ae3ffff2c433324fb38dd35e3fab7c957999 \ + --hash=sha256:aedff62918805fb62d43a4aa2ecd4482c380dc76cd31bd7c8878588a61bd0369 \ + --hash=sha256:b34e51affded8faee0dfdb705416153819d8ea9250bbbf7ea1b249bdeb5f1122 \ + --hash=sha256:b4b4d74bda2b8ebf4da5bd42af11d02d04428b2c32846e4c2c93219df8a7987b \ + --hash=sha256:b67e6efbf68e077dd71d1a6b37e43e1a99d0bff1a3d51867d45ee8908b931098 \ + --hash=sha256:b78efa4c6e804ecdf727e580dbb9cba85624d2e1c6b5cb059c66290063bd99a9 \ + --hash=sha256:bb4ae2b57fc1d8cbd1cf7b1d9913803681ffa903e7488012be5b76dedf49297f \ + --hash=sha256:bdd1a81a1860476eb41ac4bc1e07b3f07259e6d55bbf739b79c8aaedcf512799 \ + --hash=sha256:bdee92c56a71d2b24c33a7d4c2856bd6419d017e08caa7802d2963870e315028 \ + --hash=sha256:be6a04e6c79819c9a8c2373317d19a96048e5a3f90bec587787e86a1153883c2 \ + --hash=sha256:bfc08add558155345129c7803b3671cf195e6a56e7a12f3dde7c57d9b417f525 \ + --hash=sha256:c3b22c26c6fd6811b0ae8363b95ca8ce4ea3c202d3d0975b2914310ceb1bcc4d \ + --hash=sha256:c9e7cdf45d594ee04d5be1b24dd9d49f3d1590959b2271fb30b5ca2b262c00fb \ + --hash=sha256:cb27e7b78d716c591e88e0a09a2139c6577865d7f2e152488c2cc6257f460872 \ + --hash=sha256:cc9617b46837c6468197b5945e196ee9ca43057bb7d9d1ae688101e4e1dddf64 \ + --hash=sha256:ccd09f20ccdbbd341b21a67ab50a119b64a403b09288c27481575105283c1586 \ + --hash=sha256:ce6a3a4e106cf35c2d9c4fa17c05ce0b180db622736845d4315519397a77beaf \ + --hash=sha256:d0005b053977e7b43388ddec89fa567f43d4f6d5c2c0affe57de5ebf290dc552 \ + --hash=sha256:d4188e73af84ca82468f09cadc5ac4db578109e52acb4518d8154698d3a87ca2 \ + --hash=sha256:d4efec7bcf21671db6a3294ff301d2fc861c31faa3c8740d1a94689234d1b415 \ + --hash=sha256:d75aa530ccfaa593da12834b86a0724f58bff12706659baa9227c2ccaa06264c \ + --hash=sha256:d84cd4061ae292d8ac367b2c3fa3aad11cb8625a95d135fe93f286f914f3f5a6 \ + --hash=sha256:d8aacd3d4b33b772542b2e01beb50187536967b514b00003bdda7589722d2a64 \ + --hash=sha256:d8fc5c867c22b828001b6a38d2eaeb88160bf5783c6cb4a5e440efc981ce286d \ + --hash=sha256:d976bbb382b202f71c67f77b0ac11244021cfa3f7dfd9e562eefcea2df711548 \ + --hash=sha256:dba5ee5d3981160c28d5490f0d1b7ed730c22470ff7f6cc26cfcfaacb9896a07 \ + --hash=sha256:dc1ae486f9abcef254b5618dfb4113dd49f94c68e3e027d03cf0143f3f772b61 \ + --hash=sha256:dd0a578400839256df88c16abddf9ba14813ec5f21362e1fe65022e00c883d4d \ + --hash=sha256:deed0c7258ceb4c44ad5ec7d9918f9f14fd05b2be86378d86cf50e63d1e7b771 \ + --hash=sha256:e09c2279a4d01f099f52d5c4b3d9e208e91edcbd1a175c9662a8b16e000fece9 \ + --hash=sha256:e2ea9f7ab7fbf18fffb1b5434ce7c69a07582f7acc7717720f1d69f3e806f90c \ + --hash=sha256:e6b93f13371d341afee3be9f7c5964e3fe61d5fa30f6a30eb49856935dfe4fc3 \ + --hash=sha256:eb14a5da6dc7642b0f3a18f13654847cd8b7a2550e2645a5bda677862b03ba16 \ + --hash=sha256:ed0fecd28cc62c54b262e3736f8bb2512d8dcfdc2bcf08be5f47f96bf405b145 \ + --hash=sha256:ede8c6d533bc6601a47ad4046080d36b8fc99f81e6f1c17b0ac3c2dc91ac7611 \ + --hash=sha256:efb3a45b35622bb6c16dbfab491a8f5a391fe0e9d45ef32f4df85658232ca0e2 \ + --hash=sha256:f117e1a089d9411663a3207ba874f31be9ac8eaa5b533787024dc07aeb74f464 \ + --hash=sha256:f2ba92255faa7309d06fe44c3a4a97efe1c8d640c2a79a5ef728b685762a6fd2 \ + --hash=sha256:f6008a4919fdbc0b0097089f67a1eb55d950ed7e90ce2cc3e640abadd2757a04 \ + --hash=sha256:f68208a520c3d86ea51acf688a3e3002615a7f0238002cccc17affecc86a8a54 \ + --hash=sha256:f68e4f3eeca8fb22cc3d731f9715a13b652795ef657a13df1ad0c7dc0e9731df \ + --hash=sha256:fb3b8132019ea572f4611d770991000d7f58127560c4889729248eb5852a102f \ + --hash=sha256:fb940820c63a9590d31d88b815e7a3aa5915cad3ce735ab45f0c730b39547de1 \ + --hash=sha256:fc1795ac5cd0510207482c3d1d3ed781143383b8cfd36f5c645f3897ce066220 +markupsafe==3.0.3 \ + --hash=sha256:0303439a41979d9e74d18ff5e2dd8c43ed6c6001fd40e5bf2e43f7bd9bbc523f \ + --hash=sha256:068f375c472b3e7acbe2d5318dea141359e6900156b5b2ba06a30b169086b91a \ + --hash=sha256:0bf2a864d67e76e5c9a34dc26ec616a66b9888e25e7b9460e1c76d3293bd9dbf \ + --hash=sha256:0db14f5dafddbb6d9208827849fad01f1a2609380add406671a26386cdf15a19 \ + --hash=sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf \ + --hash=sha256:0f4b68347f8c5eab4a13419215bdfd7f8c9b19f2b25520968adfad23eb0ce60c \ + --hash=sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175 \ + --hash=sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219 \ + --hash=sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb \ + --hash=sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6 \ + --hash=sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab \ + --hash=sha256:15d939a21d546304880945ca1ecb8a039db6b4dc49b2c5a400387cdae6a62e26 \ + --hash=sha256:177b5253b2834fe3678cb4a5f0059808258584c559193998be2601324fdeafb1 \ + --hash=sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce \ + --hash=sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218 \ + --hash=sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634 \ + --hash=sha256:1ba88449deb3de88bd40044603fafffb7bc2b055d626a330323a9ed736661695 \ + --hash=sha256:1cc7ea17a6824959616c525620e387f6dd30fec8cb44f649e31712db02123dad \ + --hash=sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73 \ + --hash=sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c \ + --hash=sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe \ + --hash=sha256:2a15a08b17dd94c53a1da0438822d70ebcd13f8c3a95abe3a9ef9f11a94830aa \ + --hash=sha256:2f981d352f04553a7171b8e44369f2af4055f888dfb147d55e42d29e29e74559 \ + --hash=sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa \ + --hash=sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37 \ + --hash=sha256:3537e01efc9d4dccdf77221fb1cb3b8e1a38d5428920e0657ce299b20324d758 \ + --hash=sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f \ + --hash=sha256:38664109c14ffc9e7437e86b4dceb442b0096dfe3541d7864d9cbe1da4cf36c8 \ + --hash=sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d \ + --hash=sha256:3b562dd9e9ea93f13d53989d23a7e775fdfd1066c33494ff43f5418bc8c58a5c \ + --hash=sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97 \ + --hash=sha256:4bd4cd07944443f5a265608cc6aab442e4f74dff8088b0dfc8238647b8f6ae9a \ + --hash=sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19 \ + --hash=sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9 \ + --hash=sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9 \ + --hash=sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc \ + --hash=sha256:591ae9f2a647529ca990bc681daebdd52c8791ff06c2bfa05b65163e28102ef2 \ + --hash=sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4 \ + --hash=sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354 \ + --hash=sha256:6b5420a1d9450023228968e7e6a9ce57f65d148ab56d2313fcd589eee96a7a50 \ + --hash=sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698 \ + --hash=sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9 \ + --hash=sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b \ + --hash=sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc \ + --hash=sha256:7be7b61bb172e1ed687f1754f8e7484f1c8019780f6f6b0786e76bb01c2ae115 \ + --hash=sha256:7c3fb7d25180895632e5d3148dbdc29ea38ccb7fd210aa27acbd1201a1902c6e \ + --hash=sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485 \ + --hash=sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f \ + --hash=sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12 \ + --hash=sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025 \ + --hash=sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009 \ + --hash=sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d \ + --hash=sha256:949b8d66bc381ee8b007cd945914c721d9aba8e27f71959d750a46f7c282b20b \ + --hash=sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a \ + --hash=sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5 \ + --hash=sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f \ + --hash=sha256:a320721ab5a1aba0a233739394eb907f8c8da5c98c9181d1161e77a0c8e36f2d \ + --hash=sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1 \ + --hash=sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287 \ + --hash=sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6 \ + --hash=sha256:bc51efed119bc9cfdf792cdeaa4d67e8f6fcccab66ed4bfdd6bde3e59bfcbb2f \ + --hash=sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581 \ + --hash=sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed \ + --hash=sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b \ + --hash=sha256:c0c0b3ade1c0b13b936d7970b1d37a57acde9199dc2aecc4c336773e1d86049c \ + --hash=sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026 \ + --hash=sha256:c4ffb7ebf07cfe8931028e3e4c85f0357459a3f9f9490886198848f4fa002ec8 \ + --hash=sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676 \ + --hash=sha256:d2ee202e79d8ed691ceebae8e0486bd9a2cd4794cec4824e1c99b6f5009502f6 \ + --hash=sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e \ + --hash=sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d \ + --hash=sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d \ + --hash=sha256:de8a88e63464af587c950061a5e6a67d3632e36df62b986892331d4620a35c01 \ + --hash=sha256:df2449253ef108a379b8b5d6b43f4b1a8e81a061d6537becd5582fba5f9196d7 \ + --hash=sha256:e1c1493fb6e50ab01d20a22826e57520f1284df32f2d8601fdd90b6304601419 \ + --hash=sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795 \ + --hash=sha256:e2103a929dfa2fcaf9bb4e7c091983a49c9ac3b19c9061b6d5427dd7d14d81a1 \ + --hash=sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5 \ + --hash=sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d \ + --hash=sha256:e8fc20152abba6b83724d7ff268c249fa196d8259ff481f3b1476383f8f24e42 \ + --hash=sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe \ + --hash=sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda \ + --hash=sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e \ + --hash=sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737 \ + --hash=sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523 \ + --hash=sha256:f42d0984e947b8adf7dd6dde396e720934d12c506ce84eea8476409563607591 \ + --hash=sha256:f71a396b3bf33ecaa1626c255855702aca4d3d9fea5e051b41ac59a9c1c41edc \ + --hash=sha256:f9e130248f4462aaa8e2552d547f36ddadbeaa573879158d721bbd33dfe4743a \ + --hash=sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50 +narwhals==2.16.0 \ + --hash=sha256:155bb45132b370941ba0396d123cf9ed192bf25f39c4cea726f2da422ca4e145 \ + --hash=sha256:846f1fd7093ac69d63526e50732033e86c30ea0026a44d9b23991010c7d1485d +pillow==12.1.0 \ + --hash=sha256:00162e9ca6d22b7c3ee8e61faa3c3253cd19b6a37f126cad04f2f88b306f557d \ + --hash=sha256:079af2fb0c599c2ec144ba2c02766d1b55498e373b3ac64687e43849fbbef5bc \ + --hash=sha256:0b022eaaf709541b391ee069f0022ee5b36c709df71986e3f7be312e46f42c84 \ + --hash=sha256:0c27407a2d1b96774cbc4a7594129cc027339fd800cd081e44497722ea1179de \ + --hash=sha256:0ddedfaa8b5f0b4ffbc2fa87b556dc59f6bb4ecb14a53b33f9189713ae8053c0 \ + --hash=sha256:0deedf2ea233722476b3a81e8cdfbad786f7adbed5d848469fa59fe52396e4ef \ + --hash=sha256:0ed07dca4a8464bada6139ab38f5382f83e5f111698caf3191cb8dbf27d908b4 \ + --hash=sha256:0fde7ec5538ab5095cc02df38ee99b0443ff0e1c847a045554cf5f9af1f4aa82 \ + --hash=sha256:15c794d74303828eaa957ff8070846d0efe8c630901a1c753fdc63850e19ecd9 \ + --hash=sha256:1a949604f73eb07a8adab38c4fe50791f9919344398bdc8ac6b307f755fc7030 \ + --hash=sha256:1f345e7bc9d7f368887c712aa5054558bad44d2a301ddf9248599f4161abc7c0 \ + --hash=sha256:1fcc52d86ce7a34fd17cb04e87cfdb164648a3662a6f20565910a99653d66c18 \ + --hash=sha256:21e686a21078b0f9cb8c8a961d99e6a4ddb88e0fc5ea6e130172ddddc2e5221a \ + --hash=sha256:2415373395a831f53933c23ce051021e79c8cd7979822d8cc478547a3f4da8ef \ + --hash=sha256:277518bf4fe74aa91489e1b20577473b19ee70fb97c374aa50830b279f25841b \ + --hash=sha256:27b9baecb428899db6c0de572d6d305cfaf38ca1596b5c0542a5182e3e74e8c6 \ + --hash=sha256:29a4cef9cb672363926f0470afc516dbf7305a14d8c54f7abbb5c199cd8f8179 \ + --hash=sha256:3413c2ae377550f5487991d444428f1a8ae92784aac79caa8b1e3b89b175f77e \ + --hash=sha256:351889afef0f485b84078ea40fe33727a0492b9af3904661b0abbafee0355b72 \ + --hash=sha256:3ffaa2f0659e2f740473bcf03c702c39a8d4b2b7ffc629052028764324842c64 \ + --hash=sha256:40a8e3b9e8773876d6e30daed22f016509e3987bab61b3b7fe309d7019a87451 \ + --hash=sha256:414b9a78e14ffeb98128863314e62c3f24b8a86081066625700b7985b3f529bd \ + --hash=sha256:43aca0a55ce1eefc0aefa6253661cb54571857b1a7b2964bd8a1e3ef4b729924 \ + --hash=sha256:43b4899cfd091a9693a1278c4982f3e50f7fb7cff5153b05174b4afc9593b616 \ + --hash=sha256:461f9dfdafa394c59cd6d818bdfdbab4028b83b02caadaff0ffd433faf4c9a7a \ + --hash=sha256:4f9f6a650743f0ddee5593ac9e954ba1bdbc5e150bc066586d4f26127853ab94 \ + --hash=sha256:53d8b764726d3af1a138dd353116f774e3862ec7e3794e0c8781e30db0f35dfc \ + --hash=sha256:565c986f4b45c020f5421a4cea13ef294dde9509a8577f29b2fc5edc7587fff8 \ + --hash=sha256:5c5ae0a06e9ea030ab786b0251b32c7e4ce10e58d983c0d5c56029455180b5b9 \ + --hash=sha256:5cb7bc1966d031aec37ddb9dcf15c2da5b2e9f7cc3ca7c54473a20a927e1eb91 \ + --hash=sha256:5da841d81b1a05ef940a8567da92decaa15bc4d7dedb540a8c219ad83d91808a \ + --hash=sha256:5fee4c04aad8932da9f8f710af2c1a15a83582cfb884152a9caa79d4efcdbf9c \ + --hash=sha256:609e89d9f90b581c8d16358c9087df76024cf058fa693dd3e1e1620823f39670 \ + --hash=sha256:6258f3260986990ba2fa8a874f8b6e808cf5abb51a94015ca3dc3c68aa4f30ea \ + --hash=sha256:64efdf00c09e31efd754448a383ea241f55a994fd079866b92d2bbff598aad91 \ + --hash=sha256:65b80c1ee7e14a87d6a068dd3b0aea268ffcabfe0498d38661b00c5b4b22e74c \ + --hash=sha256:6741e6f3074a35e47c77b23a4e4f2d90db3ed905cb1c5e6e0d49bff2045632bc \ + --hash=sha256:681088909d7e8fa9e31b9799aaa59ba5234c58e5e4f1951b4c4d1082a2e980e0 \ + --hash=sha256:6b7a9d1db5dad90e2991645874f708e87d9a3c370c243c2d7684d28f7e133e6b \ + --hash=sha256:7315f9137087c4e0ee73a761b163fc9aa3b19f5f606a7fc08d83fd3e4379af65 \ + --hash=sha256:742aea052cf5ab5034a53c3846165bc3ce88d7c38e954120db0ab867ca242661 \ + --hash=sha256:75af0b4c229ac519b155028fa1be632d812a519abba9b46b20e50c6caa184f19 \ + --hash=sha256:7b5dd7cbae20285cdb597b10eb5a2c13aa9de6cde9bb64a3c1317427b1db1ae1 \ + --hash=sha256:7d6daa89a00b58c37cb1747ec9fb7ac3bc5ffd5949f5888657dfddde6d1312e0 \ + --hash=sha256:800429ac32c9b72909c671aaf17ecd13110f823ddb7db4dfef412a5587c2c24e \ + --hash=sha256:806f3987ffe10e867bab0ddad45df1148a2b98221798457fa097ad85d6e8bc75 \ + --hash=sha256:808b99604f7873c800c4840f55ff389936ef1948e4e87645eaf3fccbc8477ac4 \ + --hash=sha256:80941e6d573197a0c28f394753de529bb436b1ca990ed6e765cf42426abc39f8 \ + --hash=sha256:84cabc7095dd535ca934d57e9ce2a72ffd216e435a84acb06b2277b1de2689bd \ + --hash=sha256:8637e29d13f478bc4f153d8daa9ffb16455f0a6cb287da1b432fdad2bfbd66c7 \ + --hash=sha256:896866d2d436563fa2a43a9d72f417874f16b5545955c54a64941e87c1376c61 \ + --hash=sha256:8e178e3e99d3c0ea8fc64b88447f7cac8ccf058af422a6cedc690d0eadd98c51 \ + --hash=sha256:907bfa8a9cb790748a9aa4513e37c88c59660da3bcfffbd24a7d9e6abf224551 \ + --hash=sha256:9212d6b86917a2300669511ed094a9406888362e085f2431a7da985a6b124f45 \ + --hash=sha256:92a7fe4225365c5e3a8e598982269c6d6698d3e783b3b1ae979e7819f9cd55c1 \ + --hash=sha256:935b9d1aed48fcfb3f838caac506f38e29621b44ccc4f8a64d575cb1b2a88644 \ + --hash=sha256:97e9993d5ed946aba26baf9c1e8cf18adbab584b99f452ee72f7ee8acb882796 \ + --hash=sha256:983976c2ab753166dc66d36af6e8ec15bb511e4a25856e2227e5f7e00a160587 \ + --hash=sha256:9f5fefaca968e700ad1a4a9de98bf0869a94e397fe3524c4c9450c1445252304 \ + --hash=sha256:a332ac4ccb84b6dde65dbace8431f3af08874bf9770719d32a635c4ef411b18b \ + --hash=sha256:a40905599d8079e09f25027423aed94f2823adaf2868940de991e53a449e14a8 \ + --hash=sha256:a6dfc2af5b082b635af6e08e0d1f9f1c4e04d17d4e2ca0ef96131e85eda6eb17 \ + --hash=sha256:a786bf667724d84aa29b5db1c61b7bfdde380202aaca12c3461afd6b71743171 \ + --hash=sha256:a83e0850cb8f5ac975291ebfc4170ba481f41a28065277f7f735c202cd8e0af3 \ + --hash=sha256:aa0c9cc0b82b14766a99fbe6084409972266e82f459821cd26997a488a7261a7 \ + --hash=sha256:b17fbdbe01c196e7e159aacb889e091f28e61020a8abeac07b68079b6e626988 \ + --hash=sha256:b63e13dd27da389ed9475b3d28510f0f954bca0041e8e551b2a4eb1eab56a39a \ + --hash=sha256:b6e53e82ec2db0717eabb276aa56cf4e500c9a7cec2c2e189b55c24f65a3e8c0 \ + --hash=sha256:bb0984b30e973f7e2884362b7d23d0a348c7143ee559f38ef3eaab640144204c \ + --hash=sha256:bc11908616c8a283cf7d664f77411a5ed2a02009b0097ff8abbba5e79128ccf2 \ + --hash=sha256:bdec5e43377761c5dbca620efb69a77f6855c5a379e32ac5b158f54c84212b14 \ + --hash=sha256:bef9768cab184e7ae6e559c032e95ba8d07b3023c289f79a2bd36e8bf85605a5 \ + --hash=sha256:c990547452ee2800d8506c4150280757f88532f3de2a58e3022e9b179107862a \ + --hash=sha256:ca94b6aac0d7af2a10ba08c0f888b3d5114439b6b3ef39968378723622fed377 \ + --hash=sha256:cad302dc10fac357d3467a74a9561c90609768a6f73a1923b0fd851b6486f8b0 \ + --hash=sha256:d0a7735df32ccbcc98b98a1ac785cc4b19b580be1bdf0aeb5c03223220ea09d5 \ + --hash=sha256:d70347c8a5b7ccd803ec0c85c8709f036e6348f1e6a5bf048ecd9c64d3550b8b \ + --hash=sha256:d70534cea9e7966169ad29a903b99fc507e932069a881d0965a1a84bb57f6c6d \ + --hash=sha256:db44d5c160a90df2d24a24760bbd37607d53da0b34fb546c4c232af7192298ac \ + --hash=sha256:e115c15e3bc727b1ca3e641a909f77f8ca72a64fff150f666fcc85e57701c26c \ + --hash=sha256:e2479c7f02f9d505682dc47df8c0ea1fc5e264c4d1629a5d63fe3e2334b89554 \ + --hash=sha256:e5dcbe95016e88437ecf33544ba5db21ef1b8dd6e1b434a2cb2a3d605299e643 \ + --hash=sha256:e6bdb408f7c9dd2a5ff2b14a3b0bb6d4deb29fb9961e6eb3ae2031ae9a5cec13 \ + --hash=sha256:e75d3dba8fc1ddfec0cd752108f93b83b4f8d6ab40e524a95d35f016b9683b09 \ + --hash=sha256:efdc140e7b63b8f739d09a99033aa430accce485ff78e6d311973a67b6bf3208 \ + --hash=sha256:f10c98f49227ed8383d28174ee95155a675c4ed7f85e2e573b04414f7e371bda \ + --hash=sha256:f188028b5af6b8fb2e9a76ac0f841a575bd1bd396e46ef0840d9b88a48fdbcea \ + --hash=sha256:f188d580bd870cda1e15183790d1cc2fa78f666e76077d103edf048eed9c356e \ + --hash=sha256:f45bd71d1fa5e5749587613037b172e0b3b23159d1c00ef2fc920da6f470e6f0 \ + --hash=sha256:f61333d817698bdcdd0f9d7793e365ac3d2a21c1f1eb02b32ad6aefb8d8ea831 \ + --hash=sha256:fb125d860738a09d363a88daa0f59c4533529a90e564785e20fe875b200b6dbd + # imageio + # matplotlib +psutil==7.2.2 \ + --hash=sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372 \ + --hash=sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9 \ + --hash=sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841 \ + --hash=sha256:1a571f2330c966c62aeda00dd24620425d4b0cc86881c89861fbc04549e5dc63 \ + --hash=sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979 \ + --hash=sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a \ + --hash=sha256:2edccc433cbfa046b980b0df0171cd25bcaeb3a68fe9022db0979e7aa74a826b \ + --hash=sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9 \ + --hash=sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee \ + --hash=sha256:917e891983ca3c1887b4ef36447b1e0873e70c933afc831c6b6da078ba474312 \ + --hash=sha256:ab486563df44c17f5173621c7b198955bd6b613fb87c71c161f827d3fb149a9b \ + --hash=sha256:ae0aefdd8796a7737eccea863f80f81e468a1e4cf14d926bd9b6f5f2d5f90ca9 \ + --hash=sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e \ + --hash=sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc \ + --hash=sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1 \ + --hash=sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf \ + --hash=sha256:e78c8603dcd9a04c7364f1a3e670cea95d51ee865e4efb3556a3a63adef958ea \ + --hash=sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988 \ + --hash=sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486 \ + --hash=sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00 \ + --hash=sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8 +pyparsing==3.3.2 \ + --hash=sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d \ + --hash=sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc +python-dateutil==2.9.0.post0 \ + --hash=sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3 \ + --hash=sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427 +six==1.17.0 \ + --hash=sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274 \ + --hash=sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81 diff --git a/python/mujoco/sysid/README.md b/python/mujoco/sysid/README.md new file mode 100644 index 00000000..b96b9f0f --- /dev/null +++ b/python/mujoco/sysid/README.md @@ -0,0 +1,637 @@ +# Practical System Identification + +A toolbox for system identification built on top of MuJoCo. + +## API Overview + +The library solves a **box-constrained nonlinear least-squares** problem. Given +a parameter vector `θ`, simulated sensor readings `ȳ(θ)`, and recorded sensor +data `y`, the objective is: + +``` +min ½ ‖W (ȳ(θ) − y)‖² + θ + +subject to θ_min ≤ θ ≤ θ_max +``` + +where `W` is a diagonal weighting matrix and the box constraints enforce +physical plausibility (e.g., positive masses). + +The optimizer uses the **Gauss-Newton** method. The residual Jacobian +`J = ∂r/∂θ` is computed by **finite differences**: each column of `J` requires +one perturbed simulation rollout, and these evaluations are independent across +parameters and parallelize naturally across threads. The Gauss-Newton +approximate Hessian is `H ≈ JᵀJ` and the gradient is `g = Jᵀr`, yielding the +update `Δθ = −H⁻¹g`. Box constraints are handled by projected steps. + +The pipeline has five stages: + +``` +Define Parameters ──> Package Data ──> Build Residual ──> Optimize ──> Save / Report + ParameterDict ModelSequences build_residual_fn optimize save_results +``` + +--- + +### What Can You Identify? + +**Anything settable on `MjSpec` can be identified** via modifier callbacks. The +convenience functions handle common cases with correct bounds; for everything +else, write a `modifier` lambda that sets the quantity on the spec. + +**Physics parameters** — these change the model before simulation: + +| Target | Approach | +|---|---| +| Body mass | `body_inertia_param(..., InertiaType.Mass)` | +| Body mass + center of mass | `body_inertia_param(..., InertiaType.MassIpos)` | +| Full body inertia (10-D) | `body_inertia_param(..., InertiaType.Pseudo)` | +| Actuator P/D gains | `Parameter(..., modifier=lambda s, p: apply_pgain(s, "act1", p.value[0]))` | +| Contact friction / solref | `Parameter(..., modifier=lambda s, p: s.pair("cp").friction.__setitem__(0, p.value[0]))` | +| Joint damping / stiffness | `Parameter(..., modifier=lambda s, p: setattr(s.joint("j1"), "damping", p.value[0]))` | + +**Measurement parameters** — real sensors aren't perfect. They may lag behind +the simulation clock, have an unknown scale factor, or sit at a nonzero offset. +These can't be set on `MjSpec` because they aren't physics — they're artifacts +of the measurement system. `SignalTransform` (Section 4) adjusts the simulated +or recorded signals *after* rollout to account for these: + +| Target | Approach | +|---|---| +| Sensor delay | `transform.delay("*_pos", params["delay"])` | +| Sensor gain/scale | `transform.gain("*_torque", params["scale"])` | +| Sensor bias/offset | `transform.bias("*_vel", params["bias"])` | + +#### Common recipes + +**Identify link masses of a robot arm:** + +```python +from mujoco.sysid import body_inertia_param, InertiaType, ParameterDict + +params = ParameterDict() +for link in ["link1", "link2", "link3"]: + params.add(body_inertia_param(spec, model, link, inertia_type=InertiaType.Mass)) +``` + +**Identify contact friction:** + +```python +from mujoco.sysid import Parameter + +params.add(Parameter( + "floor_friction", + nominal=1.0, min_value=0.1, max_value=3.0, + modifier=lambda s, p: s.pair("foot_floor").friction.__setitem__(0, p.value[0]), +)) +``` + +--- + +### 1. Define Parameters + +A **`Parameter`** is a named value (scalar or array) with bounds and an optional +**modifier callback** that knows how to apply itself to a MuJoCo spec. A +**`ParameterDict`** collects parameters into the single vector that the +optimizer sees — it handles flattening them into one array, writing optimizer +updates back, and enforcing bounds: + +```python +from mujoco.sysid import Parameter, ParameterDict + +params = ParameterDict() + +params.add(Parameter( + "box_mass", + nominal=5.0, # starting value + min_value=1.0, + max_value=10.0, + modifier=lambda spec, p: setattr(spec.body("box"), "mass", p.value[0]), +)) + +params.add(Parameter( + "friction", + nominal=[1.6, 0.005], + min_value=[0.0, 0.0], + max_value=[3.0, 0.01], + frozen=True, # excluded from optimization + modifier=lambda spec, p: spec.pair("contact").friction.__setitem__(slice(0, 2), p.value), +)) +``` + +**`frozen`**: A frozen parameter is completely invisible to the optimizer — it +is excluded from `as_vector()`, `update_from_vector()`, `get_bounds()`, and +`randomize()`. Its modifier callback is also **not called** during +`apply_param_modifiers`, so the model uses whatever value is already in the XML +spec for that quantity. The intended workflow: define all parameters you might +ever want to identify up front, then toggle `frozen` on and off as you +iteratively narrow which parameters matter. + +Key `ParameterDict` methods: + +| Method | Description | +|---|---| +| `as_vector()` | Flatten all non-frozen parameters into a 1-D array | +| `update_from_vector(x)` | Write a flat array back into the parameters | +| `get_bounds()` | Returns `(lower, upper)` bound arrays | +| `randomize(rng)` | Sample each non-frozen parameter uniformly within bounds | +| `reset()` | Restore every parameter to its nominal value | +| `copy()` | Deep copy (preserves modifier lambdas) | +| `save_to_disk(path)` / `load_from_disk(path)` | YAML serialization (schema + values) | + +#### Body inertia parameterization + +Rigid-body inertia is tricky to identify: mass, center-of-mass, and the rotational +inertia tensor are coupled, and naively optimizing the 6 independent entries of +the inertia tensor can produce physically impossible results (e.g. negative +eigenvalues). The library implements three parameterizations of increasing +fidelity: + +| `InertiaType` | Params | What it identifies | +|---|---|---| +| `Mass` | 1 | Mass only. Optionally scales the existing rotational inertia proportionally (`scale_rot_inertia=True`). | +| `MassIpos` | 4 | Mass + center-of-mass position (3-D). Optionally scales rotational inertia. | +| `Pseudo` | 10 | Full inertia via the pseudo-inertia Cholesky factor from [Rucker & Wensing 2022](https://ieeexplore.ieee.org/document/9690029). The 10 parameters `θ = [α, d₁, d₂, d₃, s₁₂, s₂₃, s₁₃, t₁, t₂, t₃]` are the entries of a lower-triangular matrix whose product `LLᵀ` is the 4×4 pseudo-inertia matrix. Physical consistency (positive mass, positive-definite inertia tensor) is guaranteed by construction for any `θ`. | + +Use `body_inertia_param` to create a `Parameter` with the right nominal values, +bounds, and modifier already wired up: + +```python +from mujoco.sysid import body_inertia_param, InertiaType + +param = body_inertia_param( + spec, model, "link1", + inertia_type=InertiaType.Pseudo, +) +params.add(param) +``` + +--- + +### 2. Package Data + +Measured data is stored as **`TimeSeries`** objects (frozen dataclass: `times`, +`data`, optional `signal_mapping`): + +```python +from mujoco.sysid import TimeSeries, SignalType + +# For sensor/state observations: +sensordata = TimeSeries.from_names(times, sensor_data, model) # all sensors +sensordata = TimeSeries.from_names(times, data, model, names=["joint1_pos", "joint2_pos"]) + +# Explicit type disambiguation (useful when sensor/state names overlap): +sensordata = TimeSeries.from_names(times, data, model, names=[ + ("joint1_pos", SignalType.MjSensor), # sensor named "joint1_pos" + ("joint1_qpos", SignalType.MjStateQPos), # joint state +]) + +# For control signals: +control = TimeSeries.from_control_names(times, control_data, model) # all actuators +control = TimeSeries.from_control_names(times, data, model, names=["motor1_ctrl"]) + +# For custom/raw data: +ts = TimeSeries.from_custom_map(times, data, ["signal1", "signal2"]) +``` + +`signal_mapping` is a dict `{name: (SignalType, indices)}` that labels which +columns of `data` correspond to which sensor/actuator. + +**`TimeSeries` factory methods:** + +| Constructor | Use case | +|---|---| +| `TimeSeries.from_names(times, data, model, names=None)` | Sensor/state data. If `names=None`, maps all model sensors. | +| `TimeSeries.from_control_names(times, data, model, names=None)` | Control signals. If `names=None`, maps all actuators. | +| `TimeSeries.from_custom_map(times, data, signals)` | Custom data with explicit signal definitions. | +| `TimeSeries(times, data)` | Raw arrays, no signal mapping. | +| `TimeSeries(times, data, signal_mapping)` | Named signals with explicit mapping. | + +**`TimeSeries` methods:** `resample(new_times=, target_dt=)`, `interpolate(t)`, +`get(t)`, `save_to_disk(path)`, `load_from_disk(path)`, `dt_statistics()`, +`remove_from_beginning(t)`, `slice_by_name(ts, names)`. + +Bundle a spec with one or more data sequences into a **`ModelSequences`**: + +```python +from mujoco.sysid import ModelSequences, create_initial_state + +initial_state = create_initial_state(model, qpos, qvel, act) + +ms = ModelSequences( + name="robot", + spec=spec, + sequence_name=["traj_1", "traj_2"], # or a single string + initial_state=[initial_state_1, initial_state_2], + control=[control_1, control_2], + sensordata=[sensordata_1, sensordata_2], +) +``` + +You can pass a single sequence (not wrapped in a list) and it will be +auto-promoted. + +**Multiple `ModelSequences`:** Each `ModelSequences` carries its own `spec`, but +the optimizer applies the **same parameter vector `θ`** to all of them. This +enables joint optimization across different physical configurations. For example, +you might have the same robot arm recorded with and without a known payload +attached — two different specs (one has the payload body), two sets of recorded +data, but the inertial parameters of the arm links are shared. The residuals +from all `ModelSequences` are stacked and minimized jointly, giving a better- +conditioned problem than fitting each dataset independently. + +--- + +### 3. Build the Residual Function + +The **residual** is the vector of differences between simulated sensor readings +and recorded sensor data: `r(θ) = W(ȳ(θ) − y)`. Each element measures how +much the simulation with parameters `θ` disagrees with reality for one sensor +at one timestep. The optimizer's job is to find the `θ` that makes this vector +as small as possible (in the least-squares sense). + +**`build_residual_fn`** captures data and configuration, returning a closure +that the optimizer will call repeatedly: + +```python +from mujoco.sysid import build_residual_fn + +residual_fn = build_residual_fn( + models_sequences=[ms], + # Optional overrides: + modify_residual=..., # custom residual logic + custom_rollout=..., # custom simulation + sensor_weights=..., # per-sensor weighting + enabled_observations=..., # subset of sensors to use +) +``` + +#### How `residual_fn` works internally + +The returned `residual_fn(x, params)` accepts `x` as either a **1-D vector** +(plain function evaluation) or a **2-D matrix** of shape `(n_params, n_fd)` +(batched finite-difference Jacobian evaluation, where each column is a +perturbed parameter vector). This is the key to parallelism. + +For each column `i` of `x`: + +1. `params.update_from_vector(x[:, i])` — writes the optimizer's current + candidate values back into the `Parameter` objects so that each + parameter's `.value` attribute reflects column `i` of `x` +2. `model_i = apply_param_modifiers(params, spec)` — iterates over every + non-frozen parameter and calls its `modifier(spec, param)` callback, + then compiles the mutated spec into an `MjModel` +3. Replicate `model_i` once per trajectory chunk (if you have `C` data + sequences, you get `C` copies) + +This produces a flat list of `n_fd * C` models. All of them are rolled out in +a **single call** to `mujoco.rollout.rollout`: + +```python +datas = [mujoco.MjData(models[0]) for _ in range(n_threads)] # one per thread + +state, sensordata = mujoco.rollout.rollout( + models, # n_fd * C models + datas, # K thread-local scratch MjData objects + initial_states, # n_fd * C initial states + control, # n_fd * C control sequences +) +``` + +MuJoCo's rollout engine distributes the `n_fd * C` independent rollouts across +`K` threads using the `MjData` objects as thread-local scratch space (each +thread gets its own `MjData` to avoid data races). **This is why the Jacobian +computation is fast**: all `n_params + 1` perturbed rollouts (times `C` +trajectory chunks) execute in one batched, multithreaded call. + +After rollout, residuals are computed per-trajectory (predicted vs. measured +sensor data), then stacked and returned. + +#### Concrete example + +Suppose you have `p = 10` parameters and `C = 3` trajectory chunks: + +- **Function eval** (`x` is 1-D): `1 * 3 = 3` rollouts, distributed across + threads. +- **Jacobian eval** (`x` is `(10, 11)` — nominal + 10 perturbations): `11 * 3 + = 33` rollouts in one batched call. On a 16-core machine this is ~2x wall + time of a single rollout. + +#### Three tiers of customization + +| Tier | What you provide | When to use | +|---|---|---| +| **Default** | Nothing extra (or `SignalTransform`) | Standard MuJoCo sensors, optional delays/gains | +| **Custom rollout** | `custom_rollout=fn` | Non-standard simulation (e.g. task-space control) | +| **Custom residual** | `modify_residual=fn` | State-based observations, exotic loss functions | + +--- + +### 4. SignalTransform (Declarative Residual Configuration) + +After simulation, the residual pipeline compares predicted sensor readings to +recorded data. But real sensors aren't ideal — position encoders may lag by a +few milliseconds, torque sensors may have an unknown scale factor, and velocity +estimates may sit at a nonzero offset. These aren't physics parameters (you +can't set "delay" on an `MjSpec`), so they need to be corrected *after* the +rollout, before the residual is computed. + +**`SignalTransform`** lets you declare these corrections and which sensors to +use, without writing a custom residual callback. Internally it: + +1. **Time-shifts** the predicted (or measured) signals by per-sensor delay + parameters, resampling onto a common time grid. +2. **Scales** sensor columns by gain parameters (`target="predicted"` scales + the simulation output, `target="measured"` scales the recording — useful + when the sensor's scale factor is unknown on either side). +3. **Offsets** sensor columns by bias parameters. +4. Computes the weighted difference and normalizes by RMS. + +Patterns use `fnmatch` syntax, so `"*_pos"` matches all sensors whose name +ends in `_pos`: + +```python +from mujoco.sysid import SignalTransform + +transform = SignalTransform() +transform.delay("*_pos", params["delay_pos"]) # fnmatch pattern +transform.delay("*_torque", params["delay_torque"]) +transform.gain("*_torque", params["torque_scale"], target="predicted") +transform.bias("*_vel", params["vel_bias"]) +transform.enable_sensors(["joint1_pos", "joint2_pos", "joint1_torque"]) +transform.set_sensor_weights({"joint1_torque": 0.5}) + +residual_fn = build_residual_fn( + models_sequences=[ms], + modify_residual=transform.apply, # drop-in replacement +) +``` + +`SignalTransform.apply` has the same signature as `ModifyResidualFn`, so it +plugs directly into `build_residual_fn`. + +#### What this replaces + +Without `SignalTransform`, you'd write the same logic by hand as a +`modify_residual` callback using the low-level `signal_modifier` functions +(Section 8): + +```python +from mujoco.sysid._src import signal_modifier + +def modify_residual(params, predicted, measured, model, return_pred_all, **kw): + # 1. Apply delays and resample onto a common time grid. + min_d, max_d = -0.02, 0.05 # must track delay bounds yourself + measured = signal_modifier.apply_delayed_ts_window(measured, predicted, min_d, max_d) + sensor_delays = {"joint1_pos": params["delay_pos"].value[0], ...} + predicted = signal_modifier.apply_resample_and_delay( + predicted, measured.times, default_delay=0.0, sensor_delays=sensor_delays, + ) + # 2. Apply gains and biases. + predicted = signal_modifier.apply_gain(predicted, "joint1_torque", params["torque_scale"]) + predicted = signal_modifier.apply_bias(predicted, "joint1_vel", params["vel_bias"]) + # 3. Compute residual. + diff = signal_modifier.weighted_diff(predicted.data, measured.data, model, weights) + diff = signal_modifier.normalize_residual(diff, measured.data) + return diff, predicted, measured +``` + +`SignalTransform` does all of this — including tracking delay bounds, expanding +fnmatch patterns to sensor names, and handling the windowing/resampling +bookkeeping — from a few declarative lines. + +--- + +### 5. Optimize + +**`optimize`** runs box-constrained Gauss-Newton least-squares on the residual +function: + +```python +from mujoco.sysid import optimize + +opt_params, opt_result = optimize( + initial_params=params, + residual_fn=residual_fn, + optimizer="mujoco", # "mujoco", "scipy", or "scipy_parallel_fd" + max_iters=200, +) +``` + +#### Optimizer backends + +| Backend | Jacobian | Description | +|---|---|---| +| `"mujoco"` (recommended) | Parallel FD, batched | `mujoco.minimize.least_squares`. Calls `residual_fn(x)` with `x` as a 2-D matrix `(n_params, n_params+1)` — the nominal point plus one perturbation per parameter — so the entire Jacobian is computed in a single batched, multithreaded rollout call. | +| `"scipy"` | Sequential 2-point FD | `scipy.optimize.least_squares`. Computes the Jacobian column-by-column (sequential). Slower for problems with many parameters. | +| `"scipy_parallel_fd"` | Parallel FD via MuJoCo, scipy outer loop | Scipy's trust-region solver but with `mujoco.minimize.jacobian_fd` for the Jacobian. Gives scipy's convergence control with MuJoCo's batched FD speed. | + +All three backends solve box-constrained nonlinear least-squares using the +Gauss-Newton Hessian approximation `H ≈ JᵀJ`. They differ in how the +finite-difference Jacobian is computed and how the trust-region step is +handled: `"mujoco"` uses projected Gauss-Newton steps, while `"scipy"` and +`"scipy_parallel_fd"` use scipy's trust-region reflective algorithm (`trf`). + +#### Return value + +Returns `(opt_params, OptimizeResult)` where `opt_params` is a deep copy of +the input `ParameterDict` with `.value` set to the solution, and +`OptimizeResult` contains: +- `.x` — the solution vector +- `.jac` — the Jacobian at the solution (used for confidence intervals) +- `.grad` — the gradient at the solution +- `.extras` — (**mujoco backend only**) dict with `"objective"` (cost per + iteration) and `"candidate"` (parameter vector per iteration), when verbose + +--- + +### 6. Save Results and Report + +**`save_results`** writes everything to disk: + +```python +from mujoco.sysid import save_results + +save_results( + experiment_results_folder="results/exp01", + models_sequences=[ms], + initial_params=params, + opt_params=opt_params, + opt_result=opt_result, + residual_fn=residual_fn, +) +``` + +This creates: +- `params_x_0.yaml` — initial parameter values +- `params_x_hat.yaml` — optimized parameter values +- `results.pkl` — full `OptimizeResult` +- `confidence.pkl` — parameter covariance matrix `Σ_θ = σ²_r H⁻¹` and + per-parameter confidence intervals, computed from the eigendecomposition of + `H = JᵀJ` at the solution. Parameters in near-null-space directions of `H` + receive infinite confidence intervals, making identifiability issues + immediately visible. +- `{model_name}.xml` — identified MuJoCo XML for each model + +**`default_report`** generates an HTML report with sensor comparisons, parameter +tables, and videos: + +```python +from mujoco.sysid import default_report + +default_report( + models_sequences=[ms], + initial_params=params, + opt_params=opt_params, + opt_result=opt_result, + residual_fn=residual_fn, + save_dir="results/exp01", +) +``` + +--- + +### 7. Model Modification + +`apply_param_modifiers` is the default `build_model` implementation — it +iterates over all non-frozen parameters, calls each one's `modifier` callback +on the spec, and compiles. Most users never need to call it directly; it runs +automatically inside the residual pipeline. + +The remaining functions are useful when writing a **custom `build_model`** +(e.g. the box case study manually mutates the spec instead of using modifier +callbacks): + +| Function | Description | +|---|---| +| `apply_param_modifiers(params, spec)` | Run all modifier callbacks, return compiled `MjModel` | +| `apply_param_modifiers_spec(params, spec)` | Run all modifier callbacks, return the `MjSpec` | +| `apply_pgain(spec, name, value)` | Set proportional gain on a position actuator | +| `apply_dgain(spec, name, value)` | Set derivative gain on a position actuator | +| `apply_pdgain(spec, name, value)` | Set both P and D gains | +| `apply_body_inertia(spec, name, param)` | Apply Mass / MassIpos / Pseudo inertia | +| `body_inertia_param(spec, model, name, ...)` | Create a `Parameter` for body inertia | +| `remove_visuals(spec)` | Strip textures, materials, and visual-only geoms | + +--- + +### 8. Signal Modification (Power-User API) + +Low-level functions used internally by `SignalTransform` and the default +residual pipeline. Useful when writing a custom `modify_residual`: + +| Function | Description | +|---|---| +| `get_sensor_indices(model, name)` | Column indices for a named sensor | +| `apply_gain(ts, name, param)` | Multiply sensor columns by `param.value` | +| `apply_bias(ts, name, param)` | Add `param.value` to sensor columns | +| `apply_delay(ts, name, param)` | Time-shift sensor columns | +| `apply_delayed_ts_window(ts, ts_ref, min_d, max_d)` | Crop `ts` to the valid time window | +| `apply_resample_and_delay(ts, times, default_delay, ...)` | Resample with per-sensor delays | +| `weighted_diff(pred, meas, model, weights)` | `measured - predicted`, optionally weighted | +| `normalize_residual(residual, measured)` | Divide by column-wise RMS of measured data | + +--- + +### 9. Additional Utilities + +| Function / Class | Module | Description | +|---|---|---| +| `create_initial_state(model, qpos, qvel, act)` | trajectory | Pack qpos/qvel/act into a flat state vector | +| `SystemTrajectory` | trajectory | Frozen dataclass holding a single rollout (model, control, sensordata, state) | +| `sysid_rollout(models, datas, control, initial_states)` | trajectory | Parallel MuJoCo rollout returning `SystemTrajectory` list | +| `render_rollout(model, data, state, framerate)` | plotting | Render state trajectories to pixel frames | +| `calculate_intervals(residuals, J, alpha)` | optimize | Confidence intervals from Jacobian at the solution | +| `sweep_parameter(params, name, values, residual_fn)` | optimize | 1-D parameter sweep returning cost curve | +| `plot_sensor_comparison(model, ...)` | plotting | Matplotlib overlay of predicted vs. measured sensors | +| `SignalType` | timeseries | Enum: `MjSensor`, `CustomObs`, `MjStateQPos`, `MjStateQVel`, `MjStateAct`, `MjCtrl` | + +--- + +## Type Aliases + +```python +ModifyResidualFn = Callable[ + ..., tuple[np.ndarray, TimeSeries, TimeSeries] +] +# (params, sensordata_predicted, sensordata_measured, model, return_pred_all, state=..., sensor_weights=...) +# Returns (residual_array, pred_timeseries, measured_timeseries) + +CustomRolloutFn = Callable[..., Sequence[SystemTrajectory]] +# (models, datas, control_signal, initial_states, param_dicts, ...) +# Returns list of SystemTrajectory + +BuildModelFn = Callable[[ParameterDict, MjSpec], MjModel] +# Default: apply_param_modifiers +``` + +--- + +## Skeleton Case Study + +Pseudocode showing the five-stage pipeline. Replace the data-loading step with +your own hardware logs or simulation data. For a complete runnable example, see +`case_studies/box/`. + +```python +import mujoco +import numpy as np + +from mujoco.sysid import ( + Parameter, + ParameterDict, + TimeSeries, + ModelSequences, + build_residual_fn, + create_initial_state, + optimize, + save_results, +) + +# 1. Load model. +spec = mujoco.MjSpec.from_file("robot.xml") +model = spec.compile() + +# 2. Define parameters with modifier callbacks. +params = ParameterDict() +params.add(Parameter( + "link1_mass", + nominal=2.0, + min_value=0.5, + max_value=5.0, + modifier=lambda spec, p: setattr(spec.body("link1"), "mass", p.value[0]), +)) + +# 3. Package recorded data. +# times: (N,) timestamps +# ctrl_array: (N, model.nu) control inputs +# sensor_array: (N, model.nsensordata) recorded sensor readings +# qpos_0, qvel_0: initial joint positions and velocities +control = TimeSeries.from_control_names(times, ctrl_array, model) +sensordata = TimeSeries.from_names(times, sensor_array, model) +initial_state = create_initial_state(model, qpos_0, qvel_0) + +ms = ModelSequences( + name="robot", + spec=spec, + sequence_name="traj_1", + initial_state=initial_state, + control=control, + sensordata=sensordata, +) + +# 4. Build the residual function and optimize. +# models_sequences is a list because you can jointly optimize across +# multiple ModelSequences with different specs (see Section 2). +residual_fn = build_residual_fn(models_sequences=[ms]) +opt_params, opt_result = optimize( + initial_params=params, + residual_fn=residual_fn, + optimizer="mujoco", +) + +# 5. Inspect results. +print(opt_params) +save_results("results/", [ms], params, opt_params, opt_result, residual_fn) +``` diff --git a/python/mujoco/sysid/__init__.py b/python/mujoco/sysid/__init__.py new file mode 100644 index 00000000..2632d120 --- /dev/null +++ b/python/mujoco/sysid/__init__.py @@ -0,0 +1,72 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Practical system identification for MuJoCo.""" + +from mujoco.sysid._src import model_modifier as model_modifier +from mujoco.sysid._src import parameter as parameter +from mujoco.sysid._src import plotting as plotting +from mujoco.sysid._src import signal_modifier as signal_modifier +from mujoco.sysid._src.io import save_results as save_results +from mujoco.sysid._src.model_modifier import InertiaType as InertiaType +from mujoco.sysid._src.model_modifier import apply_body_inertia as apply_body_inertia +from mujoco.sysid._src.model_modifier import apply_dgain as apply_dgain +from mujoco.sysid._src.model_modifier import ( + apply_param_modifiers as apply_param_modifiers, +) +from mujoco.sysid._src.model_modifier import ( + apply_param_modifiers_spec as apply_param_modifiers_spec, +) +from mujoco.sysid._src.model_modifier import apply_pdgain as apply_pdgain +from mujoco.sysid._src.model_modifier import apply_pgain as apply_pgain +from mujoco.sysid._src.model_modifier import body_inertia_param as body_inertia_param +from mujoco.sysid._src.model_modifier import remove_visuals as remove_visuals +from mujoco.sysid._src.optimize import calculate_intervals as calculate_intervals +from mujoco.sysid._src.optimize import optimize as optimize +from mujoco.sysid._src.parameter import Parameter as Parameter +from mujoco.sysid._src.parameter import ParameterDict as ParameterDict +from mujoco.sysid._src.plotting import plot_sensor_comparison as plot_sensor_comparison +from mujoco.sysid._src.plotting import render_rollout as render_rollout +from mujoco.sysid._src.residual import BuildModelFn as BuildModelFn +from mujoco.sysid._src.residual import CustomRolloutFn as CustomRolloutFn +from mujoco.sysid._src.residual import ModifyResidualFn as ModifyResidualFn +from mujoco.sysid._src.residual import build_residual_fn as build_residual_fn +from mujoco.sysid._src.residual import ( + construct_ts_from_defaults as construct_ts_from_defaults, +) +from mujoco.sysid._src.residual import model_residual as model_residual +from mujoco.sysid._src.residual import residual as residual +from mujoco.sysid._src.signal_modifier import apply_bias as apply_bias +from mujoco.sysid._src.signal_modifier import apply_delay as apply_delay +from mujoco.sysid._src.signal_modifier import ( + apply_delayed_ts_window as apply_delayed_ts_window, +) +from mujoco.sysid._src.signal_modifier import apply_gain as apply_gain +from mujoco.sysid._src.signal_modifier import ( + apply_resample_and_delay as apply_resample_and_delay, +) +from mujoco.sysid._src.signal_modifier import get_sensor_indices as get_sensor_indices +from mujoco.sysid._src.signal_modifier import normalize_residual as normalize_residual +from mujoco.sysid._src.signal_modifier import weighted_diff as weighted_diff +from mujoco.sysid._src.signal_transform import SignalTransform as SignalTransform +from mujoco.sysid._src.timeseries import SignalType as SignalType +from mujoco.sysid._src.timeseries import TimeSeries as TimeSeries +from mujoco.sysid._src.trajectory import ModelSequences as ModelSequences +from mujoco.sysid._src.trajectory import SystemTrajectory as SystemTrajectory +from mujoco.sysid._src.trajectory import create_initial_state as create_initial_state +from mujoco.sysid._src.trajectory import sysid_rollout as sysid_rollout +from mujoco.sysid.report.defaults import default_report as default_report +from mujoco.sysid.report.defaults import ( + default_report_matplotlib as default_report_matplotlib, +) diff --git a/python/mujoco/sysid/_src/__init__.py b/python/mujoco/sysid/_src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python/mujoco/sysid/_src/io.py b/python/mujoco/sysid/_src/io.py new file mode 100644 index 00000000..4efeda5f --- /dev/null +++ b/python/mujoco/sysid/_src/io.py @@ -0,0 +1,57 @@ +"""I/O utilities for saving system identification results.""" + +import os +import pathlib +import pickle +from collections.abc import Sequence + +import scipy.optimize as scipy_optimize +from absl import logging + +from mujoco.sysid._src import parameter +from mujoco.sysid._src.optimize import calculate_intervals +from mujoco.sysid._src.trajectory import ModelSequences + + +def save_results( + experiment_results_folder: str | os.PathLike, + models_sequences: Sequence[ModelSequences], + initial_params: parameter.ParameterDict, + opt_params: parameter.ParameterDict, + opt_result: scipy_optimize.OptimizeResult, + residual_fn, +): + experiment_results_folder = pathlib.Path(experiment_results_folder) + if not experiment_results_folder.exists(): + experiment_results_folder.mkdir(parents=True, exist_ok=True) + logging.info("Experiment results will be saved to %s", experiment_results_folder) + + initial_params.save_to_disk(experiment_results_folder / "params_x_0.yaml") + opt_params.save_to_disk(experiment_results_folder / "params_x_hat.yaml") + + with open(os.path.join(experiment_results_folder, "results.pkl"), "wb") as handle: + pickle.dump(opt_result, handle, protocol=pickle.HIGHEST_PROTOCOL) + + # TODO: these intervals should be part of the params object. + residuals_star, _, _ = residual_fn(opt_result.x, opt_params, return_pred_all=True) + covariance, intervals = calculate_intervals(residuals_star, opt_result.jac) + with open(os.path.join(experiment_results_folder, "confidence.pkl"), "wb") as handle: + pickle.dump( + {"cov": covariance, "intervals": intervals}, + handle, + protocol=pickle.HIGHEST_PROTOCOL, + ) + + # Dump identified models to disk. + for model_sequences in models_sequences: + model_sequences.spec.to_file( + (experiment_results_folder / f"{model_sequences.name}.xml").as_posix() + ) + + # Log nominal compared to initial. + x0 = initial_params.as_vector() + x_nominal = initial_params.as_nominal_vector() + logging.info( + "Initial Parameters\n%s", + initial_params.compare_parameters(x0, opt_result.x, measured_params=x_nominal), + ) diff --git a/python/mujoco/sysid/_src/model_modifier.py b/python/mujoco/sysid/_src/model_modifier.py new file mode 100644 index 00000000..fd49bb60 --- /dev/null +++ b/python/mujoco/sysid/_src/model_modifier.py @@ -0,0 +1,507 @@ +"""Model modifiers.""" + +from enum import Enum +from typing import Any + +import mujoco +import numpy as np + +from mujoco.sysid._src.parameter import ModifierFn, Parameter, ParameterDict + + +def remove_visuals(in_spec: mujoco.MjSpec) -> mujoco.MjSpec: + """Remove visual elements from a Spec.""" + spec = in_spec.copy() + all_geoms = spec.worldbody.find_all("geom") + for geom in all_geoms: + if geom.contype == 0 and geom.conaffinity == 0: + if geom.type == mujoco.mjtGeom.mjGEOM_MESH and geom.meshname != "": + meshname = geom.meshname + mesh = spec.mesh(meshname) + if mesh: # multiple geoms can ref same mesh. + spec.delete(mesh) + spec.delete(geom) + + for mat in spec.materials: + spec.delete(mat) + for tex in spec.textures: + spec.delete(tex) + + spec.compile() # TODO: is this compile necessary? + return spec + + +def _get_obj_or_raise(spec: mujoco.MjSpec, obj_type: str, obj_name: str) -> Any: + getter = getattr(spec, obj_type, None) + if not callable(getter): + raise AttributeError(f"MjSpec has no method '{obj_type}'") + obj = getter(obj_name) + if obj is None: + raise ValueError(f"{obj_type.capitalize()} '{obj_name}' not found in spec.") + return obj + + +def apply_param_modifiers_spec( + params: ParameterDict, spec: mujoco.MjSpec +) -> mujoco.MjSpec: + for key in params.keys(): + param = params[key] + if not param.frozen: + param.apply_modifier(spec) + return spec + + +def apply_param_modifiers(params: ParameterDict, spec: mujoco.MjSpec) -> mujoco.MjModel: + return apply_param_modifiers_spec(params, spec).compile() + + +def _infer_inertial(spec: mujoco.MjSpec, body_name: str) -> mujoco.MjsBody: + """Override spec inertia using inferred inertia from compiled model.""" + body = _get_obj_or_raise(spec, "body", body_name) + assert isinstance(body, mujoco.MjsBody) + spec.compiler.inertiafromgeom = 2 + model = spec.compile() + body.explicitinertial = True + body.fullinertia = np.full((6, 1), np.nan) + body.mass = model.body(body_name).mass[0] + body.inertia = model.body(body_name).inertia + body.ipos = model.body(body_name).ipos + body.iquat = model.body(body_name).iquat + return body + + +def is_position_actuator(actuator) -> bool: + """Check if an actuator is a position actuator. + + This function works on both model.actuator and spec.actuator objects. + """ + return ( + actuator.gaintype == mujoco.mjtGain.mjGAIN_FIXED + and actuator.biastype == mujoco.mjtBias.mjBIAS_AFFINE + and actuator.dyntype in (mujoco.mjtDyn.mjDYN_NONE, mujoco.mjtDyn.mjDYN_FILTEREXACT) + and actuator.gainprm[0] == -actuator.biasprm[1] + ) + + +def get_actuator_pd_gains( + model: mujoco.MjModel, actuator_name: str +) -> tuple[float, float]: + actuator_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, actuator_name) + if actuator_id == -1: + raise ValueError(f"Actuator {actuator_name} not found in model.") + actuator = model.actuator(actuator_id) + if not is_position_actuator(actuator): + raise ValueError(f"Actuator {actuator_name} is not a position actuator.") + return -actuator.biasprm[1], -actuator.biasprm[2] + + +def apply_pgain( + spec: mujoco.MjSpec, + actuator_name: str, + value: float | np.ndarray, +) -> mujoco.MjSpec: + # TODO: assert scalar + actuator = _get_obj_or_raise(spec, "actuator", actuator_name) + assert isinstance(actuator, mujoco.MjsActuator) + if not is_position_actuator(actuator): + raise ValueError(f"Actuator {actuator_name} is not a position actuator.") + actuator.gainprm[0] = value + actuator.biasprm[1] = -value + return spec + + +def apply_dgain( + spec: mujoco.MjSpec, + actuator_name: str, + value: float | np.ndarray, +) -> mujoco.MjSpec: + # TODO: assert scalar + actuator = _get_obj_or_raise(spec, "actuator", actuator_name) + assert isinstance(actuator, mujoco.MjsActuator) + if not is_position_actuator(actuator): + raise ValueError(f"Actuator {actuator_name} is not a position actuator.") + actuator.biasprm[2] = -value + return spec + + +def apply_pdgain( + spec: mujoco.MjSpec, + actuator_name: str, + value: np.ndarray, +) -> mujoco.MjSpec: + if value.size != 2: + raise ValueError(f"pdgain must be a 2-element array, got {value.size}.") + apply_pgain(spec, actuator_name, value[0]) + apply_dgain(spec, actuator_name, value[1]) + return spec + + +def apply_body_mass_ipos( + spec: mujoco.MjSpec, + body_name: str, + mass: np.ndarray | None = None, + ipos: np.ndarray | None = None, + rot_inertia_scale: bool = False, +) -> mujoco.MjSpec: + # TODO: assert mass and ipos shapes + body = _infer_inertial(spec, body_name) + mass_original = body.mass + if mass is not None: + body.mass = mass + if rot_inertia_scale: + scale = mass / mass_original + body.inertia *= scale + if ipos is not None: + body.ipos = ipos + return spec + + +def scale_body_inertia( + spec: mujoco.MjSpec, + body_name: str, + value: np.ndarray, +) -> mujoco.MjSpec: + # TODO: assert scalar + body = _infer_inertial(spec, body_name) + body.inertia *= value + return spec + + +def pi_from_theta(theta: np.ndarray) -> np.ndarray: + alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3 = theta + exp_alpha = np.exp(alpha) + exp_d1 = np.exp(d1) + exp_d2 = np.exp(d2) + exp_d3 = np.exp(d3) + U = np.zeros((4, 4)) + U[0, 0] = exp_d1 + U[0, 1] = s12 + U[0, 2] = s13 + U[0, 3] = t1 + U[1, 1] = exp_d2 + U[1, 2] = s23 + U[1, 3] = t2 + U[2, 2] = exp_d3 + U[2, 3] = t3 + U[3, 3] = 1 + U *= exp_alpha + + J = U @ U.T + + sigma = J[:3, :3] + I_bar = np.trace(sigma) * np.eye(3) - sigma + h = J[:3, 3] + m = J[3, 3] + + return np.concatenate(([m], h, I_bar.flatten())) + + +def pseudoinertia_from_pi(pi: np.ndarray) -> np.ndarray: + """Converts inertial parameters π to a 4x4 pseudoinertia matrix J. + + Args: + pi: A 10-D array [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz] where: + m: Mass of the body + [hx, hy, hz]: First moment of mass + [Ixx, Iyy, Izz]: Diagonal elements of inertia tensor + [Ixy, Iyz, Ixz]: Off-diagonal elements of inertia tensor + + Returns: + A 4x4 pseudoinertia matrix J of the form: + [[Σ, h], + [hᵀ, m]] + where: + Σ = (tr(I)/2)I₃ - I: + h: The 3x1 first moment of mass vector + m: The scalar mass + """ + m = pi[0] + h = pi[1:4] + I_bar = pi[4:].reshape((3, 3)) + + Sigma = 0.5 * np.trace(I_bar) * np.eye(3) - I_bar + + J = np.zeros((4, 4)) + J[:3, :3] = Sigma + J[:3, 3] = h + J[3, :3] = h + J[3, 3] = m + + return J + + +def cholesky_decompose_upper(J: np.ndarray) -> np.ndarray: + """Perform an upper-triangular Cholesky decomposition of J. + + The returned matrix U is such that J = U @ U.T. + """ + n = J.shape[0] + indices = np.arange(n - 1, -1, -1) + J_reversed = J[indices][:, indices] + L_prime = np.linalg.cholesky(J_reversed) + return L_prime[indices][:, indices] + + +def theta_from_pseudoinertia(J: np.ndarray) -> np.ndarray: + """Extract the 10-D vector of base parameters θ from the pseudoinertia J. + + Args: + J: A 4x4 pseudoinertia. + + Returns: + A 10-D array θ = [alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3] where: + alpha: Scale parameter (log of U[3,3]) + [d1, d2, d3]: Log of diagonal elements + [s12, s23, s13]: Shear parameters from upper triangle + [t1, t2, t3]: Translation parameters from last column + """ + # U: A 4x4 upper-triangular matrix from Cholesky decomposition + U = cholesky_decompose_upper(J) + + # Extract exp(α) from the bottom-right element of U. + exp_alpha = U[3, 3] + alpha = np.log(exp_alpha) + + # Compute the d parameters from the diagonal entries (adjusted by alpha). + d1 = np.log(U[0, 0] / exp_alpha) + d2 = np.log(U[1, 1] / exp_alpha) + d3 = np.log(U[2, 2] / exp_alpha) + + # Extract the shear parameters (off-diagonals in the upper triangle). + s12 = U[0, 1] / exp_alpha + s13 = U[0, 2] / exp_alpha + s23 = U[1, 2] / exp_alpha + + # Extract the translation parameters (last column, except the bottom element). + t1 = U[0, 3] / exp_alpha + t2 = U[1, 3] / exp_alpha + t3 = U[2, 3] / exp_alpha + + return np.array([alpha, d1, d2, d3, s12, s23, s13, t1, t2, t3]) + + +def skew(v: np.ndarray) -> np.ndarray: + """Skew-symmetric matrix from a length-3 vector.""" + return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]]) + + +def inertia_to_fullinertia(q: np.ndarray, inertia: np.ndarray) -> np.ndarray: + xmat = np.empty(9) + mujoco.mju_quat2Mat(xmat, q) + R = xmat.reshape(3, 3) + return R @ np.diag(inertia) @ R.T + + +def pi_from_body(spec: mujoco.MjSpec, body_name: str) -> np.ndarray: + """Extracts the 10-D vector of inertial parameters π from a MuJoCo body. + + Args: + spec: MuJoCo model specification object. + body_name: Name of the body to extract parameters from. + + Returns: + A 10-D numpy array π = [m, hx, hy, hz, Ixx, Iyy, Izz, Ixy, Iyz, Ixz] where: + m: Mass of the body + [hx, hy, hz]: First moment of mass (m * com, where com is center of mass) + [Ixx, Iyy, Izz, Ixy, Iyz, Ixz]: Rotational inertia about the origin of the + body-fixed reference frame. + """ + body = _infer_inertial(spec, body_name) + mass = body.mass + ipos = body.ipos + inertia = body.inertia + iquat = body.iquat + + fullinertia = inertia_to_fullinertia(iquat, inertia) + # Transform inertial from ipos origin to body origin. + I_bar = fullinertia - (mass * skew(ipos) @ skew(ipos)) + + return np.concatenate([[mass], mass * ipos, I_bar.flatten()]) + + +def theta_inertia_from_body(spec: mujoco.MjSpec, body_name: str) -> np.ndarray: + pi = pi_from_body(spec, body_name) + J = pseudoinertia_from_pi(pi) + return theta_from_pseudoinertia(J) + + +def apply_body_theta_inertia( + spec: mujoco.MjSpec, + body_name: str, + theta: np.ndarray, +) -> mujoco.MjSpec: + if theta.size != 10: + raise ValueError(f"theta must be a 10-element array, got {theta.size}.") + pi = pi_from_theta(theta) + + body = _infer_inertial(spec, body_name) + body.mass = pi[0] + body.ipos = pi[1:4] / pi[0] + + # This tells the compiler to ignore the diagonal inertia and instead calculate it + # from the full inertia. + body.inertia[:] = 0.0 + body.iquat[:] = np.nan + + I_bar = pi[4:].reshape((3, 3)) + skew_ipos = skew(body.ipos) + fullinertia = I_bar + (body.mass * skew_ipos @ skew_ipos) + + # MuJoCo's ordering is: M(1,1), M(2,2), M(3,3), M(1,2), M(1,3), M(2,3) which + # corresponds to Ixx, Iyy, Izz, Ixy, Ixz + body.fullinertia[0] = fullinertia[0, 0] # Ixx + body.fullinertia[1] = fullinertia[1, 1] # Iyy + body.fullinertia[2] = fullinertia[2, 2] # Izz + body.fullinertia[3] = fullinertia[0, 1] # Ixy + body.fullinertia[4] = fullinertia[0, 2] # Ixz + body.fullinertia[5] = fullinertia[1, 2] # Iyz + + return spec + + +def apply_body_inertia(spec: mujoco.MjSpec, name: str, param: Parameter): + if not hasattr(param, "inertia_type"): + raise ValueError(f"Parameter {param.name} does not have inertia_type attribute.") + + if param.inertia_type == InertiaType.Mass: + apply_body_mass_ipos( + spec, name, mass=param.value, rot_inertia_scale=param.scale_rot_inertia + ) + + elif param.inertia_type == InertiaType.MassIpos: + apply_body_mass_ipos( + spec, + name, + mass=param.value[0], + ipos=param.value[1:4], + rot_inertia_scale=param.scale_rot_inertia, + ) + + elif param.inertia_type == InertiaType.Pseudo: + apply_body_theta_inertia(spec, name, param.value) + + +class InertiaType(Enum): + Mass = 0 + MassIpos = 1 + Pseudo = 2 + + +def body_inertia_param( + spec: mujoco.MjSpec, + model: mujoco.MjModel, + body_name: str, + inertia_type: InertiaType = InertiaType.MassIpos, + scale_rot_inertia: bool = False, + mass_bound_mult: np.ndarray | None = None, + ipos_bound_off: np.ndarray | None = None, + stretch_bound_mult: np.ndarray | None = None, + shear_bound_off: np.ndarray | None = None, + param_name: str | None = None, + modifier: ModifierFn | None = None, +) -> Parameter: + """Creates Parameter objects for the inertia of a body in a simplified manner. + + Args: + model: The MuJoCo model. + body_name: Name of the body to create the parameter for. + inertia_type: The type of inertia parameterization to use. + scale_rot_inertia: Whether to scale the original inertia when mass changes, + ignored with pseudo inertia. + mass_bound_mult: Multiplicative bounds for the mass parameter. + ipos_bound_off: Additive bounds for the ipos parameter. + stretch_bound_mult: Multiplicative bounds for the stretch parameters in the + pseudo-inertia parameterization. + shear_bound_off: Additive bounds for the shear parameters in the pseudo-inertia + parameterization. + param_name: Optional name for the parameter. Defaults to + ``"{body_name}_inertia"``. + modifier: Optional custom modifier callback. If None, the default + :func:`apply_body_inertia` modifier is registered on the Parameter.""" + + if mass_bound_mult is None: + mass_bound_mult = np.array([0.1, 10.0]) + if ipos_bound_off is None: + ipos_bound_off = np.array([-0.5, 0.5]) + if stretch_bound_mult is None: + stretch_bound_mult = np.array([0.5, 2.0]) + if shear_bound_off is None: + shear_bound_off = np.array([-0.5, 0.5]) + + body = model.body(body_name) + if param_name is None: + param_name = f"{body_name}_inertia" + + if modifier is None: + + def _default_modifier(spec, param): + return apply_body_inertia(spec, body_name, param) + + modifier = _default_modifier + + if inertia_type == InertiaType.Mass: + param = Parameter( + param_name, + body.mass, + body.mass * mass_bound_mult[0], + body.mass * mass_bound_mult[1], + modifier=modifier, + ) + param.inertia_type = inertia_type + param.scale_rot_inertia = scale_rot_inertia + + elif inertia_type == InertiaType.MassIpos: + massipos0 = np.concatenate((body.mass, body.ipos)) + massipos_low = np.concatenate( + (body.mass * mass_bound_mult[0], body.ipos + ipos_bound_off[0]) + ) + massipos_high = np.concatenate( + (body.mass * mass_bound_mult[1], body.ipos + ipos_bound_off[1]) + ) + param = Parameter( + param_name, massipos0, massipos_low, massipos_high, modifier=modifier + ) + param.inertia_type = inertia_type + param.scale_rot_inertia = scale_rot_inertia + + elif inertia_type == InertiaType.Pseudo: + theta_i_0 = theta_inertia_from_body(spec, body_name) + + # mass = exp(2*alpha) + alpha = theta_i_0[0] + mass = np.exp(2 * alpha) + mass_bounds = mass * mass_bound_mult + alpha_bounds = 0.5 * np.log(mass_bounds) + + # d1, d2, d3, stretch = exp(2*d) + # stretches body along principal axes + d = theta_i_0[1 : 1 + 3] + stretch = np.exp(2 * d) + stretch_bounds = stretch[:, np.newaxis] * np.atleast_2d(stretch_bound_mult) + d_bounds = 0.5 * np.log(stretch_bounds) + + # s12, s23, s13 + # shear the body + s_bounds = theta_i_0[4 : 4 + 3, np.newaxis] + np.atleast_2d(shear_bound_off) + + # t1, t2, t3 + # center of mass + t_bounds = theta_i_0[7:10, np.newaxis] + np.atleast_2d(ipos_bound_off) + + theta_bounds = np.vstack( + [ + alpha_bounds, + d_bounds, + s_bounds, + t_bounds, + ] + ) + param = Parameter( + param_name, theta_i_0, theta_bounds[:, 0], theta_bounds[:, 1], modifier=modifier + ) + param.inertia_type = inertia_type + + else: + raise ValueError(f"Unknown inertia_type: {inertia_type}") + + return param diff --git a/python/mujoco/sysid/_src/optimize.py b/python/mujoco/sysid/_src/optimize.py new file mode 100644 index 00000000..5f69a8a7 --- /dev/null +++ b/python/mujoco/sysid/_src/optimize.py @@ -0,0 +1,236 @@ +"""Optimization routines for system identification.""" + +from collections.abc import Callable +from typing import Literal + +import numpy as np +import scipy.optimize as scipy_optimize +from absl import logging +from mujoco import minimize as mujoco_minimize +from scipy.special import stdtrit + +from mujoco.sysid._src import parameter + + +def _scipy_least_squares( + x0: np.ndarray, + residual_fn: Callable, + bounds: tuple[np.ndarray, np.ndarray], + use_mujoco_jac: bool = False, + **kwargs, +) -> scipy_optimize.OptimizeResult: + max_nfev = kwargs.pop("max_iters", 200) + if kwargs.pop("verbose", True): + verbose = 2 + else: + verbose = 0 + x_scale = kwargs.pop("x_scale", "jac") + loss = kwargs.pop("loss", "linear") + + jac_arg: str | Callable + if use_mujoco_jac: + # This is the default step sized for finite difference used in + # scipy's least_squares and mujoco's minimize finite difference + # https://github.com/scipy/scipy/blob/91e18f3bd355477b8b7747ec82d70ac98ffd2422/scipy/optimize/_numdiff.py#L404 + eps = np.finfo(np.float64).eps ** 0.5 + if "diff_step" in kwargs: + eps = kwargs.pop("diff_step") + + def _jac_fn(x): + return mujoco_minimize.jacobian_fd( + residual=residual_fn, + x=x.reshape((-1, 1)), + r=residual_fn(x).reshape((-1, 1)), + eps=eps, + n_res=0, + bounds=[bounds[0].reshape((-1, 1)), bounds[1].reshape((-1, 1))], + )[0] + + jac_arg = _jac_fn + else: + jac_arg = "2-point" + + return scipy_optimize.least_squares( + residual_fn, + x0, + bounds=bounds, + max_nfev=max_nfev, + verbose=verbose, + x_scale=x_scale, + loss=loss, + jac=jac_arg, # pyright: ignore[reportArgumentType] + **kwargs, + ) + + +def _mujoco_least_squares( + x0: np.ndarray, + residual_fn: Callable, + bounds: tuple[np.ndarray, np.ndarray], + **kwargs, +) -> scipy_optimize.OptimizeResult: + if kwargs.pop("verbose", True): + verbose = mujoco_minimize.Verbosity.FULLITER + else: + verbose = mujoco_minimize.Verbosity.SILENT + max_iter = kwargs.pop("max_iters", 200) + x, log = mujoco_minimize.least_squares( + x0=x0, + bounds=bounds, + residual=residual_fn, + verbose=verbose, + max_iter=max_iter, + **kwargs, + ) + + # If verbose, return the full optimization log. + extras = {} + if verbose == mujoco_minimize.Verbosity.FULLITER: + extras["objective"] = [entry.objective for entry in log] + extras["candidate"] = [entry.candidate[:, 0] for entry in log] + + return scipy_optimize.OptimizeResult( + x=x, + jac=log[-1].jacobian, + grad=log[-1].grad, + extras=extras, + ) + + +def _dispatch_optimizer( + x0: np.ndarray, + residual_fn: Callable, + bounds: tuple[np.ndarray, np.ndarray], + optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"], + **kwargs, +) -> scipy_optimize.OptimizeResult: + if optimizer in ["scipy", "scipy_parallel_fd"]: + return _scipy_least_squares( + x0, + residual_fn, + bounds, + use_mujoco_jac=optimizer == "scipy_parallel_fd", + **kwargs, + ) + elif optimizer == "mujoco": + return _mujoco_least_squares(x0, residual_fn, bounds, **kwargs) + else: + raise ValueError( + f"Unsupported optimizer: '{optimizer}'. Expected one of: 'scipy', 'scipy_parallel_fd', or 'mujoco'." + ) + + +def optimize( + initial_params: parameter.ParameterDict, + residual_fn: Callable, + optimizer: Literal["scipy", "mujoco", "scipy_parallel_fd"] = "mujoco", + **optimizer_kwargs, +) -> tuple[parameter.ParameterDict, scipy_optimize.OptimizeResult]: + """Run nonlinear least-squares optimization on the residual. + + Args: + initial_params: Starting parameter values and bounds. + residual_fn: Callable with signature ``(x, params) -> (residuals, ...)`` + as returned by :func:`build_residual_fn`. + optimizer: Backend — ``"mujoco"`` (default), ``"scipy"``, or + ``"scipy_parallel_fd"`` (scipy with MuJoCo finite-difference Jacobian). + **optimizer_kwargs: Forwarded to the backend (e.g. ``max_iters``, + ``verbose``, ``loss``). + + Returns: + ``(opt_params, opt_result)`` — the optimised ParameterDict and a + ``scipy.optimize.OptimizeResult`` with at least ``x``, ``jac``, ``grad``. + """ + x0 = initial_params.as_vector() + bounds = initial_params.get_bounds() + opt_params = initial_params.copy() + + # Check if there are any parameters to optimize. + if len(opt_params) == 0 or opt_params.size == 0: + logging.warning( + "The ParameterDict is empty or contains only frozen Parameters. " + "Please declare all Parameters that need to be optimized." + ) + return opt_params, scipy_optimize.OptimizeResult( + x=x0, + jac=np.zeros((0, x0.shape[0])), + grad=np.zeros_like(x0), + extras={}, + ) + + def optimized_residual_fn(x): + residuals, _, _ = residual_fn(x, opt_params) + return np.concatenate(residuals) + + opt_result = _dispatch_optimizer( + x0, optimized_residual_fn, bounds, optimizer, **optimizer_kwargs + ) + + opt_params.update_from_vector(opt_result.x) + + return opt_params, opt_result + + +def calculate_intervals( + residuals_star, + J, + alpha=0.05, + lambda_zero_thresh=1e-15, + v_zero_thresh=1e-8, +): + if J is None or J.size == 0: + return np.empty((0, 0)), np.empty((0,)) + + # TODO(levi): account for per sensor variance + # Estimate sensor variance by assuming a good model fit, so + # remaining variance in the residual is due to sensor noise. + # Dividing by n - p is an unbiased estimate of the noise. + final_r = np.concatenate(residuals_star) + s2 = np.dot(final_r, final_r) / (final_r.size - J.shape[1]) + H = J.T @ J + + # Calculate the diagonals of the inverse of H + # using the observation that division by zero + # of eig(H) close to zero is canceled by numerically + # zero elements of the eigenvectors + # That is numerically zero eigenvalues only + # cause a confidence bound to be infinite if that eigenvalue + # has a numerically non-zero effect on the considered parameter + lamb, V = np.linalg.eigh(H) + lamb_max = np.max(lamb) + diag_inv_H = [] + for j in range(H.shape[0]): + inv_H_jj = 0.0 + v_j_max = np.max(np.abs(V[:, j])) + for i in range(H.shape[0]): + lambda_i = lamb[i] + if lambda_i / lamb_max < lambda_zero_thresh: + lambda_i = 0.0 + + v_j_i = V[j, i] + if np.abs(v_j_i / v_j_max) < v_zero_thresh: + v_j_i = 0.0 + + if lambda_i == 0.0 and v_j_i != 0.0: + inv_H_jj += np.inf + elif lambda_i == 0.0 and v_j_i == 0.0: + pass + else: + inv_H_jj += v_j_i**2 / lambda_i + diag_inv_H.append(inv_H_jj) + diag_inv_H = np.array(diag_inv_H) + + # In general eigenvalue decomposition should be more accurate + # than calculating the inverse of H using a general method + # TODO(levi): expand the eigenvalue/eigenvector element cancelation above to the full inverse matrix + # inv_H = V @ np.diag(np.divide(1, lamb, out=np.inf*np.zeros_like(lamb), where=lamb != 0.0)) @ V.T + lamb[lamb == 0] = lambda_zero_thresh + inv_H = V @ np.diag(1 / lamb) @ V.T + # print('inv test') + # print(np.diag(inv_H @ H)) + # print(np.diag(np.linalg.inv(H) @ H))) + Sigma_X = s2 * inv_H + intervals = np.sqrt(diag_inv_H * s2) * stdtrit( + final_r.size - J.shape[1], 1 - alpha / 2 + ) + return Sigma_X, intervals diff --git a/python/mujoco/sysid/_src/parameter.py b/python/mujoco/sysid/_src/parameter.py new file mode 100644 index 00000000..ac5ea1b5 --- /dev/null +++ b/python/mujoco/sysid/_src/parameter.py @@ -0,0 +1,606 @@ +"""Parameter utilities.""" + +from __future__ import annotations + +import copy +import pathlib +from typing import TYPE_CHECKING, Callable, TypeAlias + +import colorama +import mujoco +import numpy as np +import numpy.typing as npt +import yaml +from tabulate import tabulate + +if TYPE_CHECKING: + from typing_extensions import Self + + from mujoco.sysid._src.model_modifier import InertiaType + +Fore = colorama.Fore +Style = colorama.Style + +ModifierFn: TypeAlias = Callable[[mujoco.MjSpec, "Parameter"], object] + + +class Parameter: + """A single (possibly multi-dimensional) parameter for system identification. + + A Parameter holds a current ``value``, a ``nominal`` baseline, and box + bounds (``min_value``, ``max_value``). An optional ``modifier`` callback + is invoked during model compilation to apply the parameter to an MjSpec. + + Args: + name: Human-readable identifier (must be unique within a ParameterDict). + nominal: Nominal (initial) value; scalar or array-like. + min_value: Lower bound, same shape as *nominal*. + max_value: Upper bound, same shape as *nominal*. + frozen: If True the parameter is excluded from optimization. + modifier: Optional callback ``(MjSpec, Parameter) -> None`` that writes + the parameter into a spec during model compilation. + """ + + # Type hints for dynamically-added attributes (set by parameter builders). + if TYPE_CHECKING: + inertia_type: InertiaType | None + scale_rot_inertia: bool + + def __init__( + self, + name: str, + nominal: float | npt.ArrayLike, + min_value: float | npt.ArrayLike, + max_value: float | npt.ArrayLike, + frozen: bool = False, + modifier: ModifierFn | None = None, + ): + self.name = name + self.nominal = np.atleast_1d(nominal) + self.min_value = np.atleast_1d(min_value) + self.max_value = np.atleast_1d(max_value) + self.value = self.nominal.copy() + self.frozen = frozen + self.modifier = modifier + + @property + def size(self) -> int: + return self.nominal.size + + @property + def shape(self) -> tuple[int, ...]: + return self.nominal.shape + + def apply_modifier(self, spec: mujoco.MjSpec) -> None: + """Apply this parameter's modifier callback to *spec*, if one is set.""" + if self.modifier: + self.modifier(spec, self) + + def as_vector(self) -> np.ndarray: + """Return the current value as a flat 1-D array.""" + return self.value.flatten() + + def as_nominal_vector(self) -> np.ndarray: + """Return the nominal value as a flat 1-D array.""" + return self.nominal.flatten() + + def update_from_vector(self, vector: np.ndarray) -> None: + vector_array = np.atleast_1d(vector) + if len(vector_array) != self.size: + raise ValueError( + f"Input vector length {vector_array.size} does not match " + f"parameter size {self.size}." + ) + self.value = vector_array.reshape(self.shape) + + def get_bounds(self) -> tuple[np.ndarray, np.ndarray]: + """Return ``(lower, upper)`` bound arrays, each flat 1-D.""" + return ( + self.min_value.flatten(), + self.max_value.flatten(), + ) + + def reset(self) -> None: + """Reset the current value to nominal.""" + self.value = self.nominal.copy() + + def sample(self, rng: np.random.Generator | None = None) -> np.ndarray: + """Sample a random value uniformly within bounds.""" + if rng is None: + rng = np.random.default_rng() + return rng.uniform(self.min_value.flatten(), self.max_value.flatten()) + + def __str__(self) -> str: + """Return a string representation of the parameter.""" + if self.size == 1: + return ( + f"{Fore.CYAN}{self.name}{Style.RESET_ALL}: " + f"{Fore.GREEN}{float(self.value.item()):.3g}{Style.RESET_ALL} " + f"∈ [{Fore.YELLOW}{float(self.min_value.item()):.3g}, " + f"{float(self.max_value.item()):.3g}{Style.RESET_ALL}]" + ) + else: + return ( + f"{Fore.CYAN}{self.name}{Style.RESET_ALL}: " + f"{Fore.GREEN}array(shape={self.shape}){Style.RESET_ALL} " + f"∈ [{Fore.YELLOW}min={np.min(self.min_value):.3g}, " + f"max={np.max(self.max_value):.3g}{Style.RESET_ALL}]" + ) + + def __repr__(self) -> str: + return self.__str__() + + def __getstate__(self): + return { + "name": self.name, + "nominal": self.nominal.tolist() + if isinstance(self.nominal, np.ndarray) + else self.nominal, + "min_value": self.min_value.tolist() + if isinstance(self.min_value, np.ndarray) + else self.min_value, + "max_value": self.max_value.tolist() + if isinstance(self.max_value, np.ndarray) + else self.max_value, + "value": self.value.tolist() + if isinstance(self.value, np.ndarray) + else self.value, + "frozen": self.frozen, + } + + def __setstate__(self, state): + self.name = state["name"] + self.nominal = np.array(state["nominal"]) + self.min_value = np.array(state["min_value"]) + self.max_value = np.array(state["max_value"]) + self.value = np.array(state["value"]) + self.frozen = state["frozen"] + + # Override default deepycopy so lambda references get copied + def __deepcopy__(self, memo): + cls = self.__class__ + result = cls.__new__(cls) + for k, v in self.__dict__.items(): + setattr(result, k, copy.deepcopy(v, memo)) + return result + + +class ParameterDict: + """An ordered collection of :class:`Parameter` objects. + + Behaves like a ``dict[str, Parameter]`` with convenience methods for + vectorised access (``as_vector`` / ``update_from_vector``), serialisation, + and tabular comparison of parameter estimates. + + Frozen parameters are silently skipped by vector/bounds methods so that the + decision-variable dimension seen by optimizers matches only the free params. + """ + + def __init__(self, parameters: dict[str, Parameter] | None = None): + if parameters is None: + self.parameters = {} + else: + self.parameters = parameters + + def __getitem__(self, key: str) -> Parameter: + return self.parameters[key] + + def __setitem__(self, key: str, value: Parameter) -> None: + self.parameters[key] = value + + def __contains__(self, key: str) -> bool: + return key in self.parameters + + def __len__(self) -> int: + return len(self.parameters) + + def copy(self) -> Self: + """Return a deep copy of this ParameterDict.""" + return copy.deepcopy(self) + + def add(self, param: Parameter) -> None: + """Add a Parameter, keyed by its ``name``.""" + self.parameters[param.name] = param + + def update(self, pdict: Self) -> None: + for keys in pdict.keys(): + if keys in self.parameters: + raise ValueError(f"Parameter '{keys}' already exists in the dictionary.") + self.parameters[keys] = pdict[keys] + + def keys(self) -> list[str]: + return list(self.parameters.keys()) + + def values(self) -> list[Parameter]: + return list(self.parameters.values()) + + def items(self) -> list[tuple[str, Parameter]]: + return list(self.parameters.items()) + + @property + def size(self) -> int: + """Get the total size of all non-frozen parameters.""" + return sum(p.size for p in self.parameters.values() if not p.frozen) + + def as_vector(self, include_frozen=False) -> np.ndarray: + """Convert all non-frozen parameters to a flat vector.""" + vectors = [ + p.as_vector() for p in self.parameters.values() if not p.frozen or include_frozen + ] + return np.concatenate(vectors) if vectors else np.array([]) + + def as_nominal_vector(self, include_frozen=False) -> np.ndarray: + """Get the nominal values of parameters as a flat array.""" + vectors = [ + p.as_nominal_vector() + for p in self.parameters.values() + if not p.frozen or include_frozen + ] + return np.concatenate(vectors) if vectors else np.array([]) + + def update_from_vector(self, vector: np.ndarray) -> None: + """Update all non-frozen parameters from a flat vector.""" + start = 0 + for param in self.parameters.values(): + if not param.frozen: + size = param.size + param.update_from_vector(vector[start : start + size]) + start += size + + def save_to_disk(self, path: str | pathlib.Path) -> None: + """Save the parameter dictionary to disk (schema and data). + + Args: + path: Path where the data will be saved. + """ + parameter_dicts = { + name: param.__getstate__() for name, param in self.parameters.items() + } + with open(path, "w") as handle: + yaml.safe_dump(parameter_dicts, handle, default_flow_style=False) + + @classmethod + def load_from_disk(cls, path: str | pathlib.Path) -> "ParameterDict": + """Load parameter dictionary from disk (schema and data). + + Args: + path: Path to the saved data. + + Returns: + A new ParameterDict object. + """ + with open(path, "r") as handle: + parameter_dicts = yaml.safe_load(handle) + + parameters = {} + for name, param_dict in parameter_dicts.items(): + param = Parameter.__new__(Parameter) + param.__setstate__(param_dict) + parameters[name] = param + + return ParameterDict(parameters) + + def get_bounds(self) -> tuple[np.ndarray, np.ndarray]: + """Get the bounds for all non-frozen parameters.""" + lower_bounds = [] + upper_bounds = [] + for param in self.parameters.values(): + if not param.frozen: + lb, ub = param.get_bounds() + lower_bounds.append(lb) + upper_bounds.append(ub) + + return ( + np.concatenate(lower_bounds) if lower_bounds else np.array([]), + np.concatenate(upper_bounds) if upper_bounds else np.array([]), + ) + + def reset(self) -> None: + """Reset all parameters to their nominal values.""" + for param in self.parameters.values(): + param.reset() + + def sample(self, rng: np.random.Generator | None = None) -> np.ndarray: + """Sample parameter values within bounds for non-frozen parameters.""" + if rng is None: + rng = np.random.default_rng() + lower_bounds, upper_bounds = self.get_bounds() + return rng.uniform(lower_bounds, upper_bounds) + + def randomize(self, rng: np.random.Generator | None = None) -> None: + """Randomize parameter values for non-frozen parameters.""" + for param in self.parameters.values(): + if not param.frozen: + param.value = param.sample(rng) + + def compare_parameters( + self, + init_params: np.ndarray, + predicted_params: np.ndarray, + measured_params: np.ndarray | None = None, + sig_digits: int = 4, + ) -> str: + """Compare true and predicted parameter values. + + Args: + init_params: Initial parameter values as a flat array. + predicted_params: Predicted parameter values as a flat array. + measured_params: True parameter values as a flat array. + sig_digits: Number of significant digits to display. + + Returns: + A formatted string with parameter comparison table. + """ + # Get the vector of non-frozen parameters + non_frozen_vector = self.as_vector() + + if non_frozen_vector.size == 0: + return "No non-frozen parameters to compare." + + if len(init_params) != non_frozen_vector.size: + raise ValueError( + f"Initial parameter vector length {len(init_params)} does not match " + f"the number of non-frozen parameters {non_frozen_vector.size}." + ) + + if len(predicted_params) != non_frozen_vector.size: + raise ValueError( + f"Predicted parameter vector length {len(predicted_params)} does not match " + f"the number of non-frozen parameters {non_frozen_vector.size}." + ) + + if measured_params is not None: + if len(measured_params) != non_frozen_vector.size: + raise ValueError( + f"True parameter vector length {len(measured_params)} does not match " + f"the number of non-frozen parameters {non_frozen_vector.size}." + ) + + # Compute error metrics. + rel_deltas = [] + for i in range(predicted_params.shape[0]): + if ( + init_params[i] == 0 + or np.abs(predicted_params[i] - init_params[i]) / np.abs(init_params[i]) > 2e1 + ): + rel_deltas.append(np.nan) + else: + rel_deltas.append( + np.abs(predicted_params[i] - init_params[i]) / np.abs(init_params[i]) + ) + rel_deltas = np.array(rel_deltas) + overall_rms_delta = np.sqrt(np.mean((predicted_params - init_params) ** 2)) + abs_deltas = np.abs(predicted_params - init_params) + + if measured_params is not None: + rel_errors = [] + for i in range(predicted_params.shape[0]): + if ( + measured_params[i] == 0 + or np.abs(predicted_params[i] - measured_params[i]) + / np.abs(measured_params[i]) + > 2e1 + ): + rel_errors.append(np.nan) + else: + rel_errors.append( + np.abs(predicted_params[i] - measured_params[i]) + / np.abs(measured_params[i]) + ) + rel_errors = np.array(rel_errors) + + overall_rmse = np.sqrt(np.mean((predicted_params - measured_params) ** 2)) + abs_errors = np.abs(predicted_params - measured_params) + else: + overall_rmse = np.nan + abs_errors = np.full_like(predicted_params, np.nan) + rel_errors = np.full_like(predicted_params, np.nan) + + lower_bounds, upper_bounds = self.get_bounds() + + def format_number(x): + """Format number with fixed width for proper table alignment.""" + if abs(x) < 0.01: + return f"{x: .{sig_digits}e}" + else: + return f"{x: .{sig_digits}f}" + + def get_color_for_error(error): + """Get color code based on relative error magnitude.""" + if error < 0.02: + return Fore.GREEN + elif error < 0.1: + return Fore.YELLOW + else: + return Fore.RED + + def create_table_row(param_name, idx): + """Create a formatted table row for a parameter at the given index.""" + + true = measured_params[idx] if measured_params is not None else np.nan + init = init_params[idx] + est = predicted_params[idx] + lower_bound = lower_bounds[idx] + upper_bound = upper_bounds[idx] + delta = abs_deltas[idx] + error = abs_errors[idx] if measured_params is not None else np.nan + rel_delta = rel_deltas[idx] + rel_err = rel_errors[idx] if measured_params is not None else np.nan + + # If a parameter is near the boundary make it magneta + if (abs(est - lower_bound) < 1e-8 + 1e-3 * abs(lower_bound)) or ( + abs(est - upper_bound) < 1e-8 + 1e-3 * abs(upper_bound) + ): + color = Fore.MAGENTA + else: + if measured_params is None: + color = get_color_for_error(rel_delta) + else: + color = get_color_for_error(error) + + # Format all values with appropriate colors + if np.isnan(true): + measured_val = "" + else: + measured_val = f"{Fore.BLUE}{format_number(true)}{Style.RESET_ALL}" + init_val = f"{Fore.BLUE}{format_number(init)}{Style.RESET_ALL}" + est_val = f"{color}{format_number(est)}{Style.RESET_ALL}" + + lower_bound_val = f"{Fore.BLUE}{format_number(lower_bound)}{Style.RESET_ALL}" + upper_bound_val = f"{Fore.BLUE}{format_number(upper_bound)}{Style.RESET_ALL}" + + if np.isnan(error): + abs_err_val = "" + else: + abs_err_val = f"{color}{format_number(error)}{Style.RESET_ALL}" + abs_delta_val = f"{color}{format_number(delta)}{Style.RESET_ALL}" + + if np.isnan(rel_err): + rel_err_val = "" + else: + rel_err_val = f"{color}{rel_err * 100:.1f}%{Style.RESET_ALL}" + + if np.isnan(rel_delta): + rel_delta_val = "" + else: + rel_delta_val = f"{color}{rel_delta * 100:.1f}%{Style.RESET_ALL}" + + return [ + f"{Fore.CYAN}{param_name.ljust(20)}{Style.RESET_ALL}", + init_val, + measured_val, + est_val, + lower_bound_val, + upper_bound_val, + abs_err_val, + abs_delta_val, + rel_err_val, + rel_delta_val, + ] + + # Build table data. + table_data = [] + non_frozen_idx = 0 # Index for non-frozen parameters in the arrays + + for param_name, param in self.parameters.items(): + if param.frozen: + continue # Skip frozen parameters + + if param.size == 1: + table_data.append(create_table_row(param_name, non_frozen_idx)) + non_frozen_idx += 1 + else: + for i in range(param.size): + if param.shape == (param.size,): + element_name = f"{param_name}[{i}]" + else: + multi_idx = np.unravel_index(i, param.shape) + idx_str = ",".join(str(x) for x in multi_idx) + element_name = f"{param_name}[{idx_str}]" + table_data.append(create_table_row(element_name, non_frozen_idx)) + non_frozen_idx += 1 + + # Create and return the formatted table. + headers = [ + "Parameter", + "Initial", + "Nominal", + "Identified", + "Lower", + "Upper", + "Abs Err", + "Abs Del", + "Rel Err", + "Rel Del", + ] + + table = tabulate( + table_data, headers=headers, tablefmt="outline", disable_numparse=True + ) + + overall_rmse_val = "" if np.isnan(overall_rmse) else f"{overall_rmse:.4g}" + overall_rms_delta_val = f"{overall_rms_delta:.4g}" + + return f"{table}\nRMSE: {overall_rmse_val}\nRMS Delta: {overall_rms_delta_val}" + + def __str__(self) -> str: + """Return a string representation of all parameters in the dictionary.""" + if not self.parameters: + return f"{Fore.CYAN}ParameterDict{Style.RESET_ALL}(empty)" + + param_strings = [] + for name, param in self.parameters.items(): + if param.size == 1: + param_strings.append(f" {param}") + else: + # For multi-dimensional parameters, show each element on its own line + param_strings.append(f" {Fore.CYAN}{name}{Style.RESET_ALL}:") + if param.shape == (param.size,): # 1D array + for i in range(param.size): + param_strings.append( + f" [{i}]: {Fore.GREEN}{param.value[i]:.3g}{Style.RESET_ALL} " + f"∈ [{Fore.YELLOW}{param.min_value[i]:.3g}, " + f"{param.max_value[i]:.3g}{Style.RESET_ALL}]" + ) + else: # Multi-dimensional array + flat_idx = 0 + for idx in np.ndindex(param.shape): + idx_str = ",".join(str(x) for x in idx) + param_strings.append( + f" [{idx_str}]:" + f" {Fore.GREEN}{param.value[idx]:.3g}{Style.RESET_ALL} ∈" + f" [{Fore.YELLOW}{param.min_value.flat[flat_idx]:.3g}," + f" {param.max_value.flat[flat_idx]:.3g}{Style.RESET_ALL}]" + ) + flat_idx += 1 + + params_str = "\n".join(param_strings) + return f"{Fore.CYAN}ParameterDict{Style.RESET_ALL}(\n{params_str}\n)" + + def __repr__(self) -> str: + return self.__str__() + + def get_non_frozen_parameter_names(self) -> list[str]: + """Get the names of all non-frozen parameters, expanding multi-dimensional ones.""" + names = [] + for name, param in self.parameters.items(): + if not param.frozen: + if param.size == 1: + names.append(name) + else: + if param.shape == (param.size,): + for i in range(param.size): + names.append(f"{name}[{i}]") + else: + for idx in np.ndindex(param.shape): + idx_str = ",".join(map(str, idx)) + names.append(f"{name}[{idx_str}]") + return names + + def get_parameter_info(self) -> str: + """Get information about all parameters in the dictionary. + + Returns: + A formatted string with parameter information. + """ + if not self.parameters: + return "No parameters in dictionary." + + info = [] + info.append(f"{Fore.CYAN}Parameter Information:{Style.RESET_ALL}") + info.append( + f"{Fore.CYAN}{'Name':<20} {'Size':<10} {'Shape':<15} {'Frozen':<10}{Style.RESET_ALL}" + ) + info.append("-" * 60) + + for name, param in self.parameters.items(): + frozen_str = ( + f"{Fore.RED}Yes{Style.RESET_ALL}" + if param.frozen + else f"{Fore.GREEN}No{Style.RESET_ALL}" + ) + info.append( + f"{Fore.CYAN}{name:<20} {param.size:<10} {str(param.shape):<15} {frozen_str}{Style.RESET_ALL}" + ) + + return "\n".join(info) diff --git a/python/mujoco/sysid/_src/plotting.py b/python/mujoco/sysid/_src/plotting.py new file mode 100644 index 00000000..94b0a59c --- /dev/null +++ b/python/mujoco/sysid/_src/plotting.py @@ -0,0 +1,692 @@ +"""Plotting utilities.""" + +from __future__ import annotations + +from collections.abc import Sequence + +import matplotlib.pyplot as plt +import mujoco +import numpy as np +from matplotlib.lines import Line2D + +from mujoco.sysid._src import parameter + + +def plot_sensor_comparison( + model: mujoco.MjModel, + predicted_times: np.ndarray | None = None, + predicted_data: np.ndarray | None = None, + real_data: np.ndarray | None = None, + real_times: np.ndarray | None = None, + preid_data: np.ndarray | None = None, + preid_times: np.ndarray | None = None, + commanded_data: np.ndarray | None = None, + commanded_times: np.ndarray | None = None, + size_factor: float = 1.0, + title_prefix: str = "", + sensor_ids: list[int] | None = None, +): + """Plots sensor trajectories from simulation and real data. + + Args: + model: The model object providing sensor information. + predicted_times: Optional 1D array of timestamps corresponding to simulation data. + predicted_data: Optional 2D array of simulation sensor data with shape + (num_timesteps, sensor_data_dimension). + real_data: Optional 2D array of real sensor data with the same shape as + predicted_data. + real_times: A 1D array of timestamps corresponding to real data. + If None and real_data is provided, the first available timestamp array is used. + preid_data: Optional 2D array of pre-identification sensor data. + preid_times: A 1D array of timestamps for pre-identification data. + commanded_data: Optional 2D array of commanded sensor data. + commanded_times: A 1D array of timestamps for commanded data. + size_factor: A scaling factor for the figure size. + """ + # Define a more appealing color palette + predicted_color = "#1f77b4" # Steel blue + real_color = "#ff7f0e" # Safety orange + preid_color = "#2ca02c" # Forest green + commanded_color = "#9467bd" # Purple + + # Determine the reference time array to use + reference_times = None + if predicted_times is not None: + reference_times = predicted_times + elif real_times is not None: + reference_times = real_times + elif preid_times is not None: + reference_times = preid_times + elif commanded_times is not None: + reference_times = commanded_times + else: + raise ValueError("At least one time array must be provided") + + # Set times for data sources that don't have their own time arrays + if real_data is not None and real_times is None: + real_times = reference_times + if preid_data is not None and preid_times is None: + preid_times = reference_times + if commanded_data is not None and commanded_times is None: + commanded_times = reference_times + if predicted_data is not None and predicted_times is None: + predicted_times = reference_times + + if sensor_ids is None: + sensor_ids = list(range(model.nsensor)) + assert predicted_data is not None + n_plots = predicted_data.shape[1] + + fig, axes = plt.subplots( + n_plots, + 1, + figsize=(10 * size_factor, 2.5 * n_plots * size_factor), + sharex=True, + ) + if n_plots == 1: + axes = [axes] + axes = list(axes) # pyright: ignore[reportArgumentType] + + # Set an overall title for the figure. + fig.suptitle(title_prefix + " Sensors", fontsize=14) # , y=1.02) + + # Loop over each sensor. + plot_i = 0 + sensor_dim = 1 + j = 0 + dim_str = "" + for _i, sensor_id in enumerate(sensor_ids): + sensor = model.sensor(sensor_id) + sensor_name = sensor.name + sensor_dim = int(sensor.dim[0]) + sensor_addr = int(sensor.adr[0]) + + for j in range(sensor_dim): + ax = axes[plot_i] + plot_i += 1 + dim_str = "" if sensor_dim == 1 else f" {j}" + if predicted_data is not None: + assert predicted_times is not None + predicted_signal = predicted_data[:, sensor_addr : sensor_addr + sensor_dim] + ax.plot( + predicted_times, + predicted_signal[:, j], + lw=2, + color=predicted_color, + alpha=0.8, + label="Sim" + dim_str, + ) + if real_data is not None: + assert real_times is not None + real_signal = real_data[:, sensor_addr : sensor_addr + sensor_dim] + ax.plot( + real_times, + real_signal[:, j], + lw=2, + color=real_color, + linestyle="--", + alpha=0.7, + label="Real" + dim_str, + ) + if preid_data is not None: + assert preid_times is not None + preid_signal = preid_data[:, sensor_addr : sensor_addr + sensor_dim] + ax.plot( + preid_times, + preid_signal[:, j], + lw=2, + color=preid_color, + linestyle=":", + alpha=0.6, + label="Pre-ID" + dim_str, + ) + if commanded_data is not None: + assert commanded_times is not None + commanded_signal = commanded_data[:, sensor_addr : sensor_addr + sensor_dim] + ax.plot( + commanded_times, + commanded_signal[:, j], + lw=2, + color=commanded_color, + linestyle="-.", + alpha=0.6, + label="Commanded" + dim_str, + ) + # Place the sensor name in a white box in the top-left corner. + ax.text( + 0.02, + 0.9, + sensor_name + dim_str, + transform=ax.transAxes, + fontsize=10, + weight="bold", + verticalalignment="top", + horizontalalignment="left", + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + ) + + # Enable a dashed grid. + ax.grid(True, linestyle="--", alpha=0.7) + + # Loop over "extra" sensors from the user + for _ in range(plot_i, n_plots): + sensor_name = "user_sensor" + dim_str = "" if sensor_dim == 1 else f" {j}" + ax = axes[plot_i] + plot_i += 1 + if predicted_data is not None: + assert predicted_times is not None + predicted_signal = predicted_data[:, plot_i - 1] + ax.plot( + predicted_times, + predicted_signal, + lw=2, + color=predicted_color, + alpha=0.8, + label="Sim", + ) + if real_data is not None: + assert real_times is not None + real_signal = real_data[:, plot_i - 1] + ax.plot( + real_times, + real_signal, + lw=2, + color=real_color, + linestyle="--", + alpha=0.7, + label="Real", + ) + if preid_data is not None: + assert preid_times is not None + preid_signal = preid_data[:, plot_i - 1] + ax.plot( + preid_times, + preid_signal, + lw=2, + color=preid_color, + linestyle=":", + alpha=0.6, + label="Pre-ID", + ) + if commanded_data is not None: + assert commanded_times is not None + commanded_signal = commanded_data[:, plot_i - 1] + ax.plot( + commanded_times, + commanded_signal, + lw=2, + color=commanded_color, + linestyle="-.", + alpha=0.6, + label="Commanded", + ) + # Place the sensor name in a white box in the top-left corner. + ax.text( + 0.02, + 0.9, + sensor_name + dim_str, + transform=ax.transAxes, + fontsize=10, + weight="bold", + verticalalignment="top", + horizontalalignment="left", + bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"), + ) + + # Enable a dashed grid. + ax.grid(True, linestyle="--", alpha=0.7) + + # Add a unified, figure-level legend if any data is provided. + legend_handles = [] + if predicted_data is not None: + legend_handles.append( + Line2D([0], [0], color=predicted_color, lw=2, label="Simulation") + ) + if real_data is not None: + legend_handles.append( + Line2D([0], [0], color=real_color, lw=2, linestyle="--", label="Real") + ) + if preid_data is not None: + legend_handles.append( + Line2D([0], [0], color=preid_color, lw=2, linestyle=":", label="Pre-ID") + ) + if commanded_data is not None: + legend_handles.append( + Line2D( + [0], + [0], + color=commanded_color, + lw=2, + linestyle="-.", + label="Commanded", + ) + ) + + if legend_handles: + fig.legend( + handles=legend_handles, + loc="upper center", + bbox_to_anchor=(0.5, 0.935), + ncol=len(legend_handles), + fancybox=True, + shadow=True, + fontsize=10, + title="Data Source", + ) + + fig.supxlabel("Time (s)", fontsize=8) + plt.tight_layout(rect=(0, 0.03, 1, 0.9)) + + +def plot_objective( + objective: Sequence[float], + figsize: tuple[float, float] = (8, 5), +): + plt.figure(figsize=figsize) + plt.plot(objective, linewidth=2, marker="o", markersize=4) + final_value = objective[-1] + if abs(final_value) < 1e-3 or abs(final_value) > 1e3: + final_str = f"{final_value:.2e}" + else: + final_str = f"{final_value:.4f}" + plt.title(f"Objective Over Time (Final: {final_str})", fontsize=14, pad=10) + plt.grid(True, linestyle="--", alpha=0.6) + plt.xlabel("Iteration", fontsize=12) + plt.ylabel("Objective", fontsize=12) + plt.xticks(fontsize=10) + plt.yticks(fontsize=10) + plt.tight_layout() + + +def plot_candidate( + candidate: Sequence[np.ndarray], + bounds: tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray] + | None = None, + param_names: Sequence[str] | None = None, + figsize: tuple[float, float] = (12, 2.5), + dims_per_page: int = 6, + log_diff: bool = True, + bound_eps: float = 1e-3, +): + values = np.array(candidate) # shape: (n_iter, n_dim) + n_iter, n_dim = values.shape + diffs = np.diff(values, axis=0) + + mins = np.full(n_dim, -np.inf) + maxs = np.full(n_dim, np.inf) + if bounds is not None: + mins = np.array(bounds[0]) + maxs = np.array(bounds[1]) + assert mins.shape == (n_dim,) and maxs.shape == (n_dim,) + + if param_names is not None: + assert len(param_names) == n_dim + + # TODO support pages, they are currently broken because saving to disk overwrites the the pages + # n_pages = math.ceil(n_dim / dims_per_page) + n_pages = 1 + for _page in range(n_pages): + # start = page * dims_per_page + # end = min((page + 1) * dims_per_page, n_dim) + start = 0 + end = n_dim + dims_in_page = end - start + + fig, axes = plt.subplots( + dims_in_page, + 2, + figsize=(figsize[0], figsize[1] * dims_in_page), + sharex="col", + ) + if dims_in_page == 1: + axes = np.expand_dims(axes, 0) + + for i, dim in enumerate(range(start, end)): + label = param_names[dim] if param_names is not None else f"Dim {dim}" + ax_val, ax_diff = axes[i] + + vals = values[:, dim] + ax_val.set_ylabel(label, fontsize=10) + ax_val.grid(True, linestyle="--", alpha=0.6) + ax_val.tick_params(labelsize=9) + + if bounds is not None: + lower, upper = mins[dim], maxs[dim] + ax_val.axhspan(lower, upper, color="gray", alpha=0.08) + ax_val.plot( + [0, n_iter - 1], + [lower, lower], + color="gray", + linestyle="--", + alpha=0.3, + linewidth=1, + ) + ax_val.plot( + [0, n_iter - 1], + [upper, upper], + color="gray", + linestyle="--", + alpha=0.3, + linewidth=1, + ) + near_lower = np.abs(vals - lower) < bound_eps + near_upper = np.abs(vals - upper) < bound_eps + near_bound = near_lower | near_upper + for t in range(1, n_iter): + is_near_prev = near_bound[t - 1] + is_near_curr = near_bound[t] + color = "#d62728" if is_near_prev and is_near_curr else "#1f77b4" + ax_val.plot([t - 1, t], [vals[t - 1], vals[t]], color=color, linewidth=2) + ax_val.plot(t, vals[t], marker="o", markersize=3, color=color) + # Overlay triangle markers for near-bound points + for t in range(n_iter): + if near_lower[t]: + ax_val.plot(t, vals[t], marker="v", markersize=6, color="#d62728") + elif near_upper[t]: + ax_val.plot(t, vals[t], marker="^", markersize=6, color="#d62728") + else: + ax_val.plot(vals, linewidth=2, marker="o", markersize=3) + + # Annotate final value + final_val = vals[-1] + final_str = ( + f"{final_val:.2e}" + if abs(final_val) < 1e-3 or abs(final_val) > 1e3 + else f"{final_val:.4f}" + ) + ax_val.text( + n_iter - 1, + final_val, + final_str, + ha="right", + va="bottom", + fontsize=9, + color="blue", + ) + + # Annotate final value. + final_val = values[-1, dim] + final_str = ( + f"{final_val:.2e}" + if abs(final_val) < 1e-3 or abs(final_val) > 1e3 + else f"{final_val:.4f}" + ) + ax_val.text( + n_iter - 1, + final_val, + final_str, + ha="right", + va="bottom", + fontsize=9, + color="blue", + ) + + # Plot diffs + if log_diff: + eps = 1e-12 + ax_diff.plot( + np.log10(np.abs(diffs[:, dim]) + eps), + linewidth=2, + marker="x", + markersize=4, + color="tab:orange", + ) + ax_diff.set_ylabel("log Δ", fontsize=9) + else: + ax_diff.plot( + diffs[:, dim], + linewidth=2, + marker="x", + markersize=4, + color="tab:orange", + ) + + ax_diff.grid(True, linestyle="--", alpha=0.6) + ax_diff.tick_params(labelsize=9) + + # Set common labels/titles + axes[-1, 0].set_xlabel("Iteration", fontsize=12) + axes[-1, 1].set_xlabel("Iteration", fontsize=12) + axes[0, 0].set_title("Candidate Value", fontsize=12) + axes[0, 1].set_title("Δ Candidate (Diff)", fontsize=12) + + fig.suptitle(f"Candidate Values and Changes (Dims {start}-{end - 1})", fontsize=14) + fig.tight_layout(rect=(0, 0, 1, 0.96)) + + +def plot_candidate_heatmap( + candidate: Sequence[np.ndarray], + param_names: Sequence[str] | None = None, + bounds: tuple[Sequence[float] | np.ndarray, Sequence[float] | np.ndarray] + | None = None, + normalize: bool = True, + figsize: tuple[float, float] = (10, 6), + cmap: str = "RdBu", + show_colorbar: bool = True, + bound_eps: float = 1e-3, +): + data = np.array(candidate).T # shape: (n_dim, n_iter) + n_dim = data.shape[0] + + if normalize and bounds is not None: + min_bounds, max_bounds = bounds + assert len(min_bounds) == len(max_bounds) == n_dim + norm_data = np.empty_like(data) + for i in range(n_dim): + min_val = min_bounds[i] + max_val = max_bounds[i] + denom = max_val - min_val if max_val > min_val else 1.0 + norm_data[i] = (data[i] - min_val) / denom + else: + norm_data = data + + fig, ax = plt.subplots(figsize=figsize) + im = ax.imshow(norm_data, aspect="auto", cmap=cmap) + + ax.set_xlabel("Iteration", fontsize=12) + ax.set_ylabel("Parameter", fontsize=12) + + # Y-axis labels. + if param_names is not None: + assert len(param_names) == n_dim + ax.set_yticks(np.arange(n_dim)) + ax.set_yticklabels(param_names, fontsize=10) + else: + ax.set_yticks(np.arange(n_dim)) + ax.set_yticklabels([f"Dim {i}" for i in range(n_dim)], fontsize=10) + + # Plot Xs where values are at bounds. + if bounds is not None: + min_bounds, max_bounds = bounds + for dim in range(n_dim): + min_val = min_bounds[dim] + max_val = max_bounds[dim] + for iter_idx, val in enumerate(data[dim]): + if abs(val - min_val) < bound_eps or abs(val - max_val) < bound_eps: + ax.plot(iter_idx, dim, "kx", markersize=6, markeredgewidth=1.5) + + if show_colorbar: + cbar = fig.colorbar(im, ax=ax) + label = "Normalized Value" if normalize else "Value" + cbar.set_label(label, fontsize=12) + + ax.set_title("Candidate Heatmap", fontsize=14) + fig.tight_layout() + + +def parameter_confidence( + all_exp_names: Sequence[str], + all_params: Sequence[parameter.ParameterDict], + all_intervals: Sequence[np.ndarray], + cols: int = 5, + gt_params: parameter.ParameterDict | None = None, +): + named_estimates = {} + # Create an entry for every non-frozen parameter + for params in all_params: + param_names = params.get_non_frozen_parameter_names() + for name in param_names: + if name not in named_estimates: + named_estimates[name] = { + "x": [], + "intervals": [], + "min_bounds": [], + "max_bounds": [], + "plot_labels": [], + } + + for exp_name, params, intervals in zip( + all_exp_names, all_params, all_intervals, strict=True + ): + param_names = params.get_non_frozen_parameter_names() + xs = params.as_vector() + bounds = params.get_bounds() + assert xs.shape[0] == len(param_names) + if gt_params is not None: + for name in param_names: + if name in gt_params: + named_estimates[name]["xgt"] = gt_params[name].value[0] + else: + assert name[-1] == "]" + left_bracket_i = name[::-1].find("[") + index = int(name[-left_bracket_i:-1]) + named_estimates[name]["xgt"] = gt_params[name[: -left_bracket_i - 1]].value[ + index + ] + + for i, (name, x, interval) in enumerate( + zip(param_names, xs, intervals, strict=True) + ): + named_estimates[name]["x"].append(x) + named_estimates[name]["intervals"].append(interval) + named_estimates[name]["min_bounds"].append(bounds[0][i]) + named_estimates[name]["max_bounds"].append(bounds[1][i]) + named_estimates[name]["plot_labels"].append(exp_name) + + rows = len(named_estimates) // cols + 1 + fig, axs = plt.subplots( + rows, cols, figsize=(20, 2 * (len(named_estimates) // cols + 1)) + ) + if rows == 1: + axs = [axs] + + for i, name in enumerate(named_estimates): + x_list = named_estimates[name]["x"] + intervals = named_estimates[name]["intervals"] + plot_labels = named_estimates[name]["plot_labels"] + + row = i % rows + col = i // rows + + min_bound = np.min(named_estimates[name]["min_bounds"]) + max_bound = np.min(named_estimates[name]["max_bounds"]) + + for j, (x, interval, plot_label) in enumerate( + zip(x_list, intervals, plot_labels, strict=True) + ): + if not np.isfinite(interval) or 2.0 * interval > 2.0 * (max_bound - min_bound): + interval = 2.0 * (max_bound - min_bound) + eb = axs[row][col].errorbar(x, -j, xerr=interval) + eb[-1][0].set_linestyle("--") + else: + axs[row][col].errorbar(x, -j, xerr=interval) + axs[row][col].scatter(x, -j, marker="x", label=plot_label) + + axs[row][col].set_xlim([min_bound, max_bound]) + axs[row][col].yaxis.set_ticklabels([]) + axs[row][col].set_title(name) + axs[row][col].grid(True) + axs[row][col].legend(fontsize=5, loc="upper right", bbox_to_anchor=(1.4, 1.0)) + if gt_params is not None: + axs[row][col].axvline(named_estimates[name]["xgt"], color="b", ls="--") + + fig.tight_layout() + + +def render_rollout( + model: mujoco.MjModel | Sequence[mujoco.MjModel], + data: mujoco.MjData, + state: np.ndarray, + framerate: int, + camera: str | int = -1, + width: int = 640, + height: int = 480, + light_pos: Sequence[float] | None = None, +) -> list[np.ndarray]: + """Renders a rollout or batch of rollouts. + + Args: + model: Single model or list of models (one per batch). + data: MjData scratch object. + state: State array of shape (nbatch, nsteps, nstate). + framerate: Frames per second to render. + camera: Camera name or ID. + width: Image width. + height: Image height. + light_pos: Optional light position [x, y, z] to add a spotlight. + + Returns: + List of rendered frames (numpy arrays). + """ + nbatch = state.shape[0] + + if isinstance(model, mujoco.MjModel): + models_list = [model] * nbatch + else: + models_list = list(model) + if len(models_list) == 1: + models_list = models_list * nbatch + else: + assert len(models_list) == nbatch + + # Visual options + vopt = mujoco.MjvOption() + vopt.geomgroup[3] = 1 # Show visualization geoms + + pert = mujoco.MjvPerturb() + catmask = mujoco.mjtCatBit.mjCAT_DYNAMIC + + # Simulate and render. + frames = [] + + with mujoco.Renderer(models_list[0], height=height, width=width) as renderer: + for i in range(state.shape[1]): + # Check if we should capture this frame based on framerate + if len(frames) < i * models_list[0].opt.timestep * framerate: + for j in range(state.shape[0]): + # Set state + mujoco.mj_setState( + models_list[j], data, state[j, i, :], mujoco.mjtState.mjSTATE_FULLPHYSICS + ) + mujoco.mj_forward(models_list[j], data) + + # Use first model to make the scene, add subsequent models + if j == 0: + renderer.update_scene(data, camera, scene_option=vopt) + else: + mujoco.mjv_addGeoms( + models_list[j], data, vopt, pert, catmask, renderer.scene + ) + + # Add light, if requested + if light_pos is not None: + if renderer.scene.nlight < 100: # check limit + light = renderer.scene.lights[renderer.scene.nlight] + light.ambient = [0, 0, 0] + light.attenuation = [1, 0, 0] + light.castshadow = 1 + light.cutoff = 45 + light.diffuse = [0.8, 0.8, 0.8] + light.dir = [0, 0, -1] + light.type = mujoco.mjtLightType.mjLIGHT_SPOT + light.exponent = 10 + light.headlight = 0 + light.specular = [0.3, 0.3, 0.3] + light.pos = light_pos + renderer.scene.nlight += 1 + + # Render and add the frame. + pixels = renderer.render() + frames.append(pixels) + return frames diff --git a/python/mujoco/sysid/_src/residual.py b/python/mujoco/sysid/_src/residual.py new file mode 100644 index 00000000..2ce1816b --- /dev/null +++ b/python/mujoco/sysid/_src/residual.py @@ -0,0 +1,408 @@ +"""Residual computation for system identification.""" + +from __future__ import annotations + +import copy +import os +from collections.abc import Callable, Mapping, Sequence +from typing import TypeAlias + +import mujoco +import numpy as np + +from mujoco.sysid._src import ( + model_modifier, + parameter, + signal_modifier, + timeseries, +) +from mujoco.sysid._src.trajectory import ( + ModelSequences, + SystemTrajectory, + sysid_rollout, +) + +_NUM_CPUS: int = os.cpu_count() or 1 + +BuildModelFn: TypeAlias = Callable[ + [parameter.ParameterDict, mujoco.MjSpec], mujoco.MjModel +] + +CustomRolloutFn: TypeAlias = Callable[..., Sequence[SystemTrajectory]] +"""Replaces the default sysid_rollout. Called with keyword arguments: +models, datas, control_signal, initial_states, param_dicts, +rollout_signal_mapping, rollout_state_mapping, ctrl_mapping.""" + +ModifyResidualFn: TypeAlias = Callable[ + ..., tuple[np.ndarray, timeseries.TimeSeries, timeseries.TimeSeries] +] +"""Custom residual computation. Called as: +modify_residual(params, sensordata_predicted, sensordata_measured, +model, return_pred_all, state=..., sensor_weights=...).""" + + +def construct_ts_from_defaults( + state_ts: timeseries.TimeSeries, + pred_sensordata: timeseries.TimeSeries, + measured_sensordata: timeseries.TimeSeries, + enabled_observations: Sequence[tuple[str, timeseries.SignalType]] | None = None, +): + """Assemble predicted observations to match the measured signal layout. + + For each enabled observation, copies the predicted values from either + ``pred_sensordata`` (for MjSensor signals) or ``state_ts`` (for state + signals like qpos/qvel/act) into a new array whose columns align with + the measured data. + + Args: + state_ts: Predicted state TimeSeries (time column already stripped). + pred_sensordata: Raw predicted sensor TimeSeries from rollout. + measured_sensordata: Measured sensor TimeSeries (defines the target layout). + enabled_observations: Subset of observations to include. If None, all + observations in ``measured_sensordata`` are used. + + Returns: + A ``(measured, predicted)`` tuple of TimeSeries with matching signal + mappings, sliced to the enabled observations. + """ + assert measured_sensordata.signal_mapping is not None + + # Trim measured data enabled observations + if enabled_observations: + enabled_observations_names = [i[0] for i in enabled_observations] + enabled_observations_types = [i[1] for i in enabled_observations] + else: + enabled_observations_names = list(measured_sensordata.signal_mapping.keys()) + enabled_observations_types = [ + v[0] for v in measured_sensordata.signal_mapping.values() + ] + + selected_measured_sensordata = timeseries.TimeSeries.slice_by_name( + measured_sensordata, enabled_observations_names + ) + assert selected_measured_sensordata.signal_mapping is not None + selected_measured_signal_mapping = selected_measured_sensordata.signal_mapping + + shape = (pred_sensordata.data.shape[0], selected_measured_sensordata.data.shape[1]) + predicted_data_out = np.zeros(shape) + + measured_signal_mapping = measured_sensordata.signal_mapping + for enabled_obs_name, enabled_obs_type in zip( + enabled_observations_names, enabled_observations_types, strict=True + ): + assert state_ts.signal_mapping is not None + if ( + enabled_obs_name not in measured_signal_mapping + and enabled_obs_name not in state_ts.signal_mapping + ): + raise ValueError(f"{enabled_obs_name} is missing.") + + obs_type, indices = measured_signal_mapping[enabled_obs_name] + + if obs_type != enabled_obs_type: + raise ValueError( + f"Observation type error: {enabled_obs_name} is of type {obs_type} but declared as {enabled_obs_type}." + ) + + if obs_type == timeseries.SignalType.CustomObs: + raise ValueError( + f"You are attempting to use the default SysID's modify_residual with a custom observation of name {enabled_obs_name}. This is not supported. You must implement your own modify_residual. See documentation at ..." + ) + + elif obs_type == timeseries.SignalType.MjSensor: + target_indices = selected_measured_signal_mapping[enabled_obs_name][1] + predicted_data_out[:, ..., target_indices] = pred_sensordata.data[:, ..., indices] + + elif ( + obs_type == timeseries.SignalType.MjStateQPos + or obs_type == timeseries.SignalType.MjStateQVel + or obs_type == timeseries.SignalType.MjStateAct + ): + state_indices = state_ts.signal_mapping[enabled_obs_name][1] + + values = state_ts.data[:, ..., state_indices] + target_indices = selected_measured_signal_mapping[enabled_obs_name][1] + predicted_data_out[:, ..., target_indices] = values + + ts_predicted_data = timeseries.TimeSeries( + pred_sensordata.times, + predicted_data_out, + selected_measured_sensordata.signal_mapping, + ) + + return selected_measured_sensordata, ts_predicted_data + + +# Lowest level residual function, works on one model +def model_residual( + x: np.ndarray, + params: parameter.ParameterDict, + build_model: Callable[[parameter.ParameterDict], mujoco.MjModel], + traj_measured: Sequence[SystemTrajectory] | SystemTrajectory, + modify_residual: ModifyResidualFn | None = None, + custom_rollout: CustomRolloutFn | None = None, + n_threads: int = _NUM_CPUS, + return_pred_all: bool = False, + resample_true: bool = True, + sensor_weights: Mapping[str, float] | None = None, + enabled_observations: Sequence[tuple[str, timeseries.SignalType]] = (), +): + """Compute residuals for a single model against measured trajectories. + + Builds the model from *x*, rolls out each trajectory, and computes the + weighted difference between predicted and measured sensor data. + + Args: + x: Decision variable vector (flat, or 2-D for batched finite-difference). + params: Parameter dictionary — updated in-place from *x*. + build_model: ``(ParameterDict) -> MjModel`` factory. + traj_measured: Ground-truth trajectory or sequence of trajectories. + modify_residual: Optional custom residual callback (replaces the default + resampling / differencing logic). + custom_rollout: Optional replacement for :func:`sysid_rollout`. + n_threads: Number of ``MjData`` scratch objects for parallel rollout. + return_pred_all: If True, return full predicted/measured TimeSeries. + resample_true: Whether to resample the measured data at simulation + timesteps (ignored when *modify_residual* is provided). + sensor_weights: Per-sensor weights for the weighted diff. + enabled_observations: Subset of ``(name, SignalType)`` pairs to include. + + Returns: + A 3-tuple ``(residuals, pred_sensordatas, measured_sensordatas)``. + """ + # Convert single trajectory to list for consistent handling. + if isinstance(traj_measured, SystemTrajectory): + traj_measured = [traj_measured] + n_chunks = len(traj_measured) + + # Handle finite difference columns if present. + initial_ndim = x.ndim + n_fd = 1 + if x.ndim > 1: + n_fd = x.shape[1] + x_reshaped = x + else: + x_reshaped = x.reshape(-1, 1) + + # Process each finite difference column. + models = [] + models_x = [] + model_0 = None + for i in range(n_fd): + params.update_from_vector(x_reshaped[:, i]) + model = build_model(params) + if not model_0: + model_0 = model + models_x.extend([x_reshaped[:, i]] * n_chunks) + models.extend([model] * n_chunks) + + assert model_0 is not None + qpos_map, qvel_map, act_map, rollout_ctrl_map = ( + timeseries.TimeSeries.compute_all_state_mappings(model_0) + ) + rollout_state_mapping = qpos_map | qvel_map | act_map + rollout_signal_mapping = timeseries.TimeSeries.compute_all_sensor_mapping(model_0) + + # Create data objects for parallel computation. + datas = [mujoco.MjData(models[0]) for _ in range(n_threads)] + + # Interpolate control signal. + if resample_true: + control_chunks = [ + traj.control.resample(target_dt=models[0].opt.timestep) for traj in traj_measured + ] + else: + control_chunks = [traj.control for traj in traj_measured] + + # Rollout trajectories in parallel. + if custom_rollout is None: + pred_trajectories = sysid_rollout( + models=models[: n_fd * n_chunks], + datas=datas, + control_signal=[control for control in control_chunks] * n_fd, + initial_states=[chunk.initial_state for chunk in traj_measured] * n_fd, + rollout_signal_mapping=rollout_signal_mapping, + rollout_state_mapping=rollout_state_mapping, + ctrl_mapping=rollout_ctrl_map, + ) + else: + param_dicts = [copy.deepcopy(params) for i in range(x_reshaped.shape[1])] + [ + param_dicts[i].update_from_vector(x_reshaped[:, i]) + for i in range(x_reshaped.shape[1]) + ] + pred_trajectories = custom_rollout( + models=models[: n_fd * n_chunks], + datas=datas, + control_signal=[control for control in control_chunks] * n_fd, + initial_states=[chunk.initial_state for chunk in traj_measured] * n_fd, + param_dicts=param_dicts, + rollout_signal_mapping=rollout_signal_mapping, + rollout_state_mapping=rollout_state_mapping, + ctrl_mapping=rollout_ctrl_map, + ) + + # Compute residuals for each trajectory chunk. + all_residuals = [] + pred_sensordatas = [] + measured_sensordatas = [] + + for i in range(len(models)): + model = models[i] + pred_traj = pred_trajectories[i] + assert pred_traj.state is not None + pred_state = pred_traj.state.data + + rollout_state_ts = timeseries.TimeSeries( + times=pred_state[:, 0], + data=pred_state[:, 1:], + signal_mapping=rollout_state_mapping, + ) + + measuredidx = i % n_chunks + measuredtraj = traj_measured[measuredidx] + + pred_sensordata = pred_traj.sensordata + measured_sensordata = measuredtraj.sensordata + + # If the user passes a residual function allow them to handle all resampling, etc. + if modify_residual is not None: + params.update_from_vector(models_x[i]) + res, pred_sensordata, measured_sensordata = modify_residual( + params, + pred_sensordata, + measured_sensordata, + model, + return_pred_all, + state=pred_state, + ) + + # If the user does not pass a residual function, resample the ground truth data to + # match the sime times if requested. + else: + measured_sensordata, pred_sensordata = construct_ts_from_defaults( + rollout_state_ts, pred_sensordata, measured_sensordata, enabled_observations + ) + if resample_true: + # Window the true data so that times in it correspond to times spanned by + # predicted data. + measured_sensordata = signal_modifier.apply_delayed_ts_window( + measured_sensordata, pred_sensordata, 0.0, 0.0 + ) + # Sample the predicted signal at the true times. + pred_sensordata = pred_sensordata.resample(measured_sensordata.times) + + else: + # Do not include difference in first sensor outputs in residual vector. + # It corresponds to the initial condition and so provides little new + # information. Additionally the semantics of rollout make it difficult to + # simulate the sensor output corresponding to the initial condition. + measured_sensordata = timeseries.TimeSeries( + measured_sensordata.times[1:], + measured_sensordata.data[1:, :], + measured_sensordata.signal_mapping, + ) + + res = signal_modifier.weighted_diff( + predicted_data=pred_sensordata.data, + measured_data=measured_sensordata.data, + model=model, + sensor_weights=sensor_weights, + ) + res = signal_modifier.normalize_residual(res, measured_sensordata.data) + + if pred_sensordata.signal_mapping != measured_sensordata.signal_mapping: + raise ValueError( + "The observation mapping between the measured data and predicted rollout data" + " is not the same. You have not modified the observation data in TimeSeries" + " in modify_residual to correctly reflect the measured data." + ) + + all_residuals.append(res) + pred_sensordatas.append(pred_sensordata) + measured_sensordatas.append(measured_sensordata) + + res_array = np.stack(all_residuals, axis=0) + if initial_ndim == 1: + res_array = res_array.ravel() + else: + res_array = res_array.reshape(res_array.shape[0], -1) + + return res_array.T, pred_sensordatas, measured_sensordatas + + +def build_residual_fn(**captured_kwargs): + """Create a residual closure with pre-bound keyword arguments. + + Returns a function ``fn(x, params, **overrides)`` that calls + :func:`residual` with the captured kwargs merged in. This is the + recommended way to construct the callable passed to :func:`optimize`. + + Example:: + + residual_fn = build_residual_fn( + models_sequences=seqs, + signal_transform=transform, + ) + opt_params, result = optimize(params, residual_fn) + """ + + def built_residual_fn(x, params, **kwargs): + return residual( + x, + params, + **captured_kwargs, + **kwargs, + ) + + return built_residual_fn + + +def residual( + x: np.ndarray, + params: parameter.ParameterDict, + models_sequences: list[ModelSequences], + build_model: BuildModelFn = model_modifier.apply_param_modifiers, + modify_residual: ModifyResidualFn | None = None, + custom_rollout: CustomRolloutFn | None = None, + n_threads: int = _NUM_CPUS, + return_pred_all: bool = False, + resample_true: bool = True, + sensor_weights: Mapping[str, float] | None = None, + enabled_observations: Sequence[tuple[str, timeseries.SignalType]] = (), +): + """Top-level residual: iterate over all model-sequence groups. + + Calls :func:`model_residual` for every measured rollout in every + :class:`ModelSequences` entry and collects the results. + + Returns: + A 3-tuple ``(residuals, preds, records)`` — lists with one entry per + measured rollout across all groups. + """ + residuals = [] + preds = [] + records = [] + for model_sequences in models_sequences: + for measured_rollout in model_sequences.measured_rollout: + res = model_residual( + x, + params, + lambda p, _spec=model_sequences.spec: build_model(p, _spec), + measured_rollout, + modify_residual, + custom_rollout, + n_threads, + return_pred_all, + resample_true, + sensor_weights, + enabled_observations, + ) + if isinstance(res, np.ndarray): + residuals.append(res) + else: + residuals.append(res[0]) + preds.append(res[1]) + records.append(res[2]) + + return residuals, preds, records diff --git a/python/mujoco/sysid/_src/signal_modifier.py b/python/mujoco/sysid/_src/signal_modifier.py new file mode 100644 index 00000000..5728f6ce --- /dev/null +++ b/python/mujoco/sysid/_src/signal_modifier.py @@ -0,0 +1,244 @@ +"""Common signal modifiers.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import mujoco +import numpy as np + +from mujoco.sysid._src import parameter, timeseries + + +def _get_sensor_indices(model: mujoco.MjModel, sensor_name: str) -> list[int]: + sensor_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR.value, sensor_name) + if sensor_id == -1: + raise ValueError(f"Sensor not found in model: {sensor_name}") + + addr = model.sensor_adr[sensor_id] + dim = model.sensor_dim[sensor_id] + + return list(range(addr, addr + dim)) + + +def get_sensor_indices( + model: mujoco.MjModel, + sensor_name: str | list[str], + sort: bool = False, +) -> list[int]: + """Get sensor indices from a sensor configuration dictionary. + + Args: + model: MuJoCo model containing the sensors. + sensor_name: sensor name or list of names to return the indices for + """ + if isinstance(sensor_name, str): + return _get_sensor_indices(model, sensor_name) + all_indices = [] + for name in sensor_name: + all_indices.extend(_get_sensor_indices(model, name)) + if sort: + return sorted(all_indices) + return all_indices + + +def apply_bias( + ts: timeseries.TimeSeries, + sensor_name: str, + bias: parameter.Parameter, +) -> timeseries.TimeSeries: + indices = ts.get_indices(sensor_name)[1] + data_out = ts.data.copy() + data_out[..., indices] += bias.value + return timeseries.TimeSeries(ts.times, data_out, ts.signal_mapping) + + +def apply_gain( + ts: timeseries.TimeSeries, + sensor_name: str, + gain: parameter.Parameter, +) -> timeseries.TimeSeries: + indices = ts.get_indices(sensor_name)[1] + data_out = ts.data.copy() + data_out[..., indices] *= gain.value + return timeseries.TimeSeries(ts.times, data_out, ts.signal_mapping) + + +def apply_delay( + ts: timeseries.TimeSeries, + sensor_name: str, + delay: parameter.Parameter, +) -> timeseries.TimeSeries: + indices = ts.get_indices(sensor_name)[1] + + ts_sensor = timeseries.TimeSeries(ts.times, ts.data[:, indices], ts.signal_mapping) + ts_sensor_delayed = ts_sensor.resample(ts.times - delay.value) + + ts_delayed = timeseries.TimeSeries(ts.times, ts.data, ts.signal_mapping) + ts_delayed.data[:, indices] = ts_sensor_delayed.data + + return ts_delayed + + +def apply_time_window( + ts: timeseries.TimeSeries, + min_t: float, + max_t: float, +) -> timeseries.TimeSeries: + """Select a subset of a timeseries whose timestamps plus the max delay can be + sampled from ts_sample.""" + min_i = np.searchsorted(ts.times, min_t, side="left") + max_i = np.searchsorted(ts.times, max_t, side="right") + return timeseries.TimeSeries( + ts.times[min_i:max_i], ts.data[min_i:max_i], ts.signal_mapping + ) + + +def apply_delayed_ts_window( + ts: timeseries.TimeSeries, + ts_delayed: timeseries.TimeSeries, + min_delay: float, + max_delay: float, +) -> timeseries.TimeSeries: + """Window a timeseries so that the included timestamps lay within the bounds of a + timeseries that may be delayed between min_delay and max_delay. + + Args: + ts: The timeseries to window. + ts_delayed: The timeseries to use as the bounds. + min_delay: The minimum delay. May be negative. + max_delay: The maximum delay. + + Returns: + A new timeseries with the timestamps windowed. + """ + if min_delay > max_delay: + raise ValueError( + "min_delay must be less than or equal to max_delay, " + f"received {min_delay} and {max_delay}" + ) + return apply_time_window( + ts, ts_delayed.times[0] - min_delay, ts_delayed.times[-1] - max_delay + ) + + +def _build_per_column_delays( + ts: timeseries.TimeSeries, + default_delay: float, + sensor_delays: dict[str, float] | None, + predicted_data: bool, +) -> list[float]: + """Build a per-column delay list, shared by both implementations.""" + delays = [default_delay] * ts.data.shape[1] + if sensor_delays is None: + sensor_delays = {} + for name, delay in sensor_delays.items(): + sensor_indices = ts.get_indices(name)[1] + for i in sensor_indices: + delays[i] = delay + if predicted_data: + delays = [-d for d in delays] + return delays + + +def _apply_resample_and_delay_columnwise( + ts: timeseries.TimeSeries, + times: np.ndarray, + delays: list[float], +) -> np.ndarray: + """Reference implementation: resample each column independently.""" + resampled_ts = [] + for i, d in enumerate(delays): + ts_sliced = timeseries.TimeSeries( + ts.times, ts.data[:, i : i + 1], ts.signal_mapping + ) + ts_sliced_resampled = ts_sliced.resample(times + d) + resampled_ts.append(ts_sliced_resampled) + return np.concatenate([t.data for t in resampled_ts], axis=1) + + +_VERIFY_RESAMPLE_GROUPING = False + + +def apply_resample_and_delay( + ts: timeseries.TimeSeries, + times: np.ndarray, + default_delay: float, + sensor_delays: dict[str, float] | None = None, + predicted_data: bool = True, +) -> timeseries.TimeSeries: + delays = _build_per_column_delays(ts, default_delay, sensor_delays, predicted_data) + + # Group columns by delay value to minimize interpolation calls. + delay_to_cols: dict[float, list[int]] = {} + for i, d in enumerate(delays): + delay_to_cols.setdefault(d, []).append(i) + + data_out = np.empty((len(times), ts.data.shape[1])) + for d, cols in delay_to_cols.items(): + group_data = ts.data[:, cols] + group_ts = timeseries.TimeSeries(ts.times, group_data, ts.signal_mapping) + resampled = group_ts.resample(times + d) + data_out[:, cols] = resampled.data + + if _VERIFY_RESAMPLE_GROUPING: + reference = _apply_resample_and_delay_columnwise(ts, times, delays) + np.testing.assert_array_equal(data_out, reference) + + return timeseries.TimeSeries(times, data_out, ts.signal_mapping) + + +def prepare_sensor_weights( + sensor_weights: Mapping[str, float] | np.ndarray, + n_sensors: int, + model: mujoco.MjModel, +) -> np.ndarray: + if isinstance(sensor_weights, np.ndarray): + if sensor_weights.ndim != 1 or sensor_weights.shape[0] != n_sensors: + raise ValueError( + "Expected sensor_weights to be a numpy array of shape (n_sensors,), " + f"received {sensor_weights.shape}" + ) + return sensor_weights + else: + weights = np.ones(n_sensors) + ids = get_sensor_indices(model, list(sensor_weights.keys())) + for i, w in zip(ids, sensor_weights.values(), strict=True): + weights[i] = w + return weights + + +def weighted_diff( + predicted_data: np.ndarray, + measured_data: np.ndarray, + model: mujoco.MjModel | None = None, + sensor_weights: Mapping[str, float] | np.ndarray | None = None, +) -> np.ndarray: + """Compute the difference `measured_data - predicted_data`, optionally scaled by + sensor weights. + + Args: + predicted_data: The predicted data, of shape (n_timesteps, n_sensors). + measured_data: The measured data, of shape (n_timesteps, n_sensors). + sensor_weights: An optional dict mapping sensor name to weight. Unspecified sensors + are assumed to have a weight of 1. + model: Optional mujoco model. This argument is required if sensor_weights is not + None. + + Returns: + A numpy array of the weighted difference. + """ + res = measured_data - predicted_data + if sensor_weights is None: + return res + if model is None: + raise ValueError("model is required if sensor_weights is provided") + return res * prepare_sensor_weights(sensor_weights, res.shape[-1], model) + + +def normalize_residual( + residual: np.ndarray, + measured_data: np.ndarray, +) -> np.ndarray: + """Normalize the residual by the standard deviation of the measured data.""" + return residual / (np.linalg.norm(measured_data, axis=0) / np.sqrt(2)) diff --git a/python/mujoco/sysid/_src/signal_transform.py b/python/mujoco/sysid/_src/signal_transform.py new file mode 100644 index 00000000..0ce7e809 --- /dev/null +++ b/python/mujoco/sysid/_src/signal_transform.py @@ -0,0 +1,257 @@ +"""Declarative signal transformation for system identification residuals.""" + +from __future__ import annotations + +from collections.abc import Mapping +from fnmatch import fnmatch + +import mujoco +import numpy as np + +from mujoco.sysid._src import parameter, signal_modifier, timeseries + + +class SignalTransform: + """Declarative signal transformation replacing boilerplate modify_residual callbacks. + + Usage:: + + transform = SignalTransform() + transform.delay("*_pos", params["delay_pos"]) + transform.delay("*_torque", params["delay_torque"]) + transform.gain("*_torque", params["torque_scale"], target="predicted") + transform.enable_sensors(cfg.sensors_enabled) + + The ``apply`` method has the same signature as ``ModifyResidualFn`` and can + be passed directly to ``build_residual_fn(signal_transform=transform)``. + """ + + def __init__(self, normalize: bool = True): + self._delays: list[tuple[str, str, parameter.Parameter]] = [] + self._gains: list[tuple[str, str, str]] = [] + self._biases: list[tuple[str, str, str]] = [] + self._enabled_sensors: list[str] | None = None + self._sensor_weights: Mapping[str, float] | None = None + self.normalize = normalize + + def delay(self, pattern: str, param: parameter.Parameter) -> None: + """Register a delay for sensors matching *pattern* (fnmatch).""" + self._delays.append((pattern, param.name, param)) + + def gain( + self, + pattern: str, + param: parameter.Parameter, + target: str = "both", + ) -> None: + """Register a multiplicative gain for sensors matching *pattern*. + + Args: + pattern: fnmatch pattern matched against sensor names. + param: Parameter whose ``.value`` is the gain factor. + target: One of ``"predicted"``, ``"measured"``, or ``"both"``. + """ + if target not in ("predicted", "measured", "both"): + raise ValueError( + f"target must be 'predicted', 'measured', or 'both', got {target!r}" + ) + self._gains.append((pattern, param.name, target)) + + def bias( + self, + pattern: str, + param: parameter.Parameter, + target: str = "both", + ) -> None: + """Register an additive bias for sensors matching *pattern*.""" + if target not in ("predicted", "measured", "both"): + raise ValueError( + f"target must be 'predicted', 'measured', or 'both', got {target!r}" + ) + self._biases.append((pattern, param.name, target)) + + def enable_sensors(self, sensor_names: list[str]) -> None: + """Only include these sensors in the returned residual/timeseries.""" + self._enabled_sensors = list(sensor_names) + + def set_sensor_weights(self, weights: Mapping[str, float]) -> None: + """Set per-sensor weights for the weighted diff.""" + self._sensor_weights = weights + + # Private methods. + + def _resolve_delays( + self, + sensor_names: list[str], + params: parameter.ParameterDict, + ) -> dict[str, float]: + """Resolve delay patterns to concrete sensor name -> delay value (last match wins).""" + resolved: dict[str, float] = {} + for pattern, param_name, _ in self._delays: + delay_value = params[param_name].value[0] + for name in sensor_names: + if fnmatch(name, pattern): + resolved[name] = delay_value + return resolved + + def _compute_delay_bounds(self) -> tuple[float, float]: + """Compute min/max delay across all registered delay params (deduplicated by name).""" + if not self._delays: + return 0.0, 0.0 + seen: set[str] = set() + min_vals: list[float] = [] + max_vals: list[float] = [] + for _, param_name, param in self._delays: + if param_name in seen: + continue + seen.add(param_name) + min_vals.append(float(param.min_value[0])) + max_vals.append(float(param.max_value[0])) + return min(min_vals), max(max_vals) + + def _get_sensor_names(self, ts: timeseries.TimeSeries) -> list[str]: + """Extract sensor names from a TimeSeries signal_mapping.""" + if ts.signal_mapping is None: + return [] + return list(ts.signal_mapping.keys()) + + def _apply_gains_biases_reference( + self, + ts: timeseries.TimeSeries, + target_label: str, + params: parameter.ParameterDict, + ) -> timeseries.TimeSeries: + """Reference implementation: one full copy per gain/bias application.""" + sensor_names = self._get_sensor_names(ts) + for pattern, param_name, target in self._gains: + if target != target_label and target != "both": + continue + for name in sensor_names: + if fnmatch(name, pattern): + ts = signal_modifier.apply_gain(ts, name, params[param_name]) + for pattern, param_name, target in self._biases: + if target != target_label and target != "both": + continue + for name in sensor_names: + if fnmatch(name, pattern): + ts = signal_modifier.apply_bias(ts, name, params[param_name]) + return ts + + _VERIFY_GAINS_BIASES = False + + def _apply_gains_biases( + self, + ts: timeseries.TimeSeries, + target_label: str, + params: parameter.ParameterDict, + ) -> timeseries.TimeSeries: + """Apply matching gains and biases to a timeseries for the given target label.""" + sensor_names = self._get_sensor_names(ts) + data = ts.data.copy() + + for pattern, param_name, target in self._gains: + if target != target_label and target != "both": + continue + for name in sensor_names: + if fnmatch(name, pattern): + indices = ts.get_indices(name)[1] + data[..., indices] *= params[param_name].value + + for pattern, param_name, target in self._biases: + if target != target_label and target != "both": + continue + for name in sensor_names: + if fnmatch(name, pattern): + indices = ts.get_indices(name)[1] + data[..., indices] += params[param_name].value + + result = timeseries.TimeSeries(ts.times, data, ts.signal_mapping) + + if self._VERIFY_GAINS_BIASES: + import numpy as _np + + ref = self._apply_gains_biases_reference(ts, target_label, params) + _np.testing.assert_array_equal(result.data, ref.data) + + return result + + def apply( + self, + params: parameter.ParameterDict, + sensordata_predicted: timeseries.TimeSeries, + sensordata_measured: timeseries.TimeSeries, + model: mujoco.MjModel, + return_pred_all: bool, + state: np.ndarray | None = None, + sensor_weights: Mapping[str, float] | None = None, + ) -> tuple[np.ndarray, timeseries.TimeSeries, timeseries.TimeSeries]: + """Apply all registered transforms and compute the residual. + + Signature matches :data:`ModifyResidualFn` so this method can be passed + directly as ``modify_residual`` to :func:`model_residual`. + + Pipeline: window measured data, resample + delay predicted data, apply + gains/biases, weighted diff, normalise, slice to enabled sensors. + + Returns: + ``(residual_array, predicted_ts, measured_ts)`` — the residual matrix + and the (possibly sliced) predicted/measured TimeSeries. + """ + del state # Part of ModifyResidualFn signature but unused here. + sensor_names = self._get_sensor_names(sensordata_predicted) + + # 1. Resolve delays and compute bounds. + sensor_delays = self._resolve_delays(sensor_names, params) + min_delay, max_delay = self._compute_delay_bounds() + + # 2. Window measured data. + sensordata_measured = signal_modifier.apply_delayed_ts_window( + sensordata_measured, sensordata_predicted, min_delay, max_delay + ) + + # 3. Resample and delay predicted data. + if sensor_delays: + sensordata_predicted = signal_modifier.apply_resample_and_delay( + sensordata_predicted, + sensordata_measured.times, + 0.0, + sensor_delays=sensor_delays, + ) + else: + sensordata_predicted = sensordata_predicted.resample(sensordata_measured.times) + + # 4. Apply gains and biases. + sensordata_predicted = self._apply_gains_biases( + sensordata_predicted, "predicted", params + ) + sensordata_measured = self._apply_gains_biases( + sensordata_measured, "measured", params + ) + + # 5. Weighted diff. + weights = sensor_weights or self._sensor_weights + res = signal_modifier.weighted_diff( + predicted_data=sensordata_predicted.data, + measured_data=sensordata_measured.data, + model=model, + sensor_weights=weights, + ) + + # 6. Normalize. + if self.normalize: + res = signal_modifier.normalize_residual(res, sensordata_measured.data) + + # 7. Slice to enabled sensors. + if not return_pred_all and self._enabled_sensors is not None: + indices = signal_modifier.get_sensor_indices(model, self._enabled_sensors) + sensordata_predicted = timeseries.TimeSeries( + sensordata_predicted.times, + sensordata_predicted.data[:, indices], + ) + sensordata_measured = timeseries.TimeSeries( + sensordata_measured.times, + sensordata_measured.data[:, indices], + ) + res = res[:, indices] + + return res, sensordata_predicted, sensordata_measured diff --git a/python/mujoco/sysid/_src/timeseries.py b/python/mujoco/sysid/_src/timeseries.py new file mode 100644 index 00000000..c46e0ea6 --- /dev/null +++ b/python/mujoco/sysid/_src/timeseries.py @@ -0,0 +1,700 @@ +"""Time series utilities.""" + +from __future__ import annotations + +import pathlib +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Literal, TypeAlias + +import mujoco +import numpy as np +import scipy.interpolate + + +class SignalType(Enum): + MjSensor = 0 + CustomObs = 1 + MjStateQPos = 2 + MjStateQVel = 3 + MjStateAct = 4 + MjCtrl = 5 + + +SignalMappingType: TypeAlias = dict[str, tuple[SignalType, np.ndarray]] + +InterpolationMethod = Literal[ + "linear", "cubic", "quadratic", "quintic", "zero_order_hold", "zoh" +] + + +def _resolve_signals( + model: mujoco.MjModel, + names: Sequence[str | tuple[str, SignalType]], + allowed_types: set[SignalType], +) -> SignalMappingType: + """Resolves signal names to (canonical_name, type, indices) mappings. + + Each name can be a string or (name, SignalType) tuple for disambiguation. + """ + result: SignalMappingType = {} + idx = 0 + + for item in names: + name, hint = item if isinstance(item, tuple) else (item, None) + resolved = _resolve_one(model, name, hint, allowed_types) + + if resolved is None: + if hint is not None and hint not in allowed_types: + raise ValueError(f"Signal '{name}' has type {hint.name} which is not allowed.") + raise ValueError( + f"Could not resolve signal '{item}' with allowed types {[t.name for t in allowed_types]}." + ) + + canon_name, sig_type, width = resolved + result[canon_name] = (sig_type, np.arange(idx, idx + width)) + idx += width + + return result + + +# Suffix conventions for state/control signals +_SUFFIXES = { + SignalType.MjStateQPos: "_qpos", + SignalType.MjStateQVel: "_qvel", + SignalType.MjStateAct: "_act", + SignalType.MjCtrl: "_ctrl", +} + + +def _strip_suffix(name: str, suffix: str) -> str: + """Strip suffix from name if present.""" + return name[: -len(suffix)] if name.endswith(suffix) else name + + +def _resolve_one( + model: mujoco.MjModel, + name: str, + hint: SignalType | None, + allowed: set[SignalType], +) -> tuple[str, SignalType, int] | None: + """Resolve a single signal name to (canonical_name, type, width).""" + + # 1. Sensor + if _type_allowed(hint, SignalType.MjSensor, allowed): + sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name) + if sid >= 0: + return (name, SignalType.MjSensor, model.sensor_dim[sid]) + + # 2. Control + if _type_allowed(hint, SignalType.MjCtrl, allowed): + base = _strip_suffix(name, "_ctrl") + aid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, base) + if aid >= 0: + return (base + "_ctrl", SignalType.MjCtrl, 1) + + # 3. State (qpos/qvel) + for sig_type in (SignalType.MjStateQPos, SignalType.MjStateQVel): + if _type_allowed(hint, sig_type, allowed): + base = _strip_suffix(name, _SUFFIXES[sig_type]) + width = _joint_or_body_width(model, base, sig_type) + if width > 0: + return (base + _SUFFIXES[sig_type], sig_type, width) + + # 4. Actuator state (act) + if _type_allowed(hint, SignalType.MjStateAct, allowed): + base = _strip_suffix(name, "_act") + aid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, base) + if aid >= 0 and model.actuator_actnum[aid] > 0: + return (base + "_act", SignalType.MjStateAct, model.actuator_actnum[aid]) + + return None + + +def _type_allowed( + hint: SignalType | None, target: SignalType, allowed: set[SignalType] +) -> bool: + """Check if target type is allowed given hint and allowed set.""" + return (hint is None or hint == target) and target in allowed + + +def _joint_or_body_width(model: mujoco.MjModel, name: str, sig_type: SignalType) -> int: + """Get state width for a joint or free body.""" + # Joint widths by type + QPOS_WIDTHS = {mujoco.mjtJoint.mjJNT_FREE: 7, mujoco.mjtJoint.mjJNT_BALL: 4} + QVEL_WIDTHS = {mujoco.mjtJoint.mjJNT_FREE: 6, mujoco.mjtJoint.mjJNT_BALL: 3} + + jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) + if jid >= 0: + jtype = model.jnt_type[jid] + widths = QPOS_WIDTHS if sig_type == SignalType.MjStateQPos else QVEL_WIDTHS + return widths.get(jtype, 1) + + # Free body + bid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name) + if bid >= 0 and model.body_dofnum[bid] == 6: + return 7 if sig_type == SignalType.MjStateQPos else 6 + + return 0 + + +@dataclass(frozen=True) +class TimeSeries: + """A utility class for working with time-series data. + + Attributes: + times: 1D array of timestamps. + data: Array of signal data. The first axis corresponds to time. + signal_mapping: Dict of tuples that maps the signal type and its + signal fields in data + """ + + times: np.ndarray + data: np.ndarray + signal_mapping: SignalMappingType | None = None + + @staticmethod + def compute_all_sensor_mapping(model: mujoco.MjModel) -> SignalMappingType: + """Computes mapping for all sensors in the model.""" + signal_mapping = {} + for sensor_id in range(model.nsensor): + addr = model.sensor_adr[sensor_id] + dim = model.sensor_dim[sensor_id] + name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_id) + indices = np.arange(addr, addr + dim) + signal_mapping[name] = (SignalType.MjSensor, indices) + return signal_mapping + + @staticmethod + def compute_all_control_mapping(model: mujoco.MjModel) -> SignalMappingType: + """Computes mapping for all controls (actuators) in the model.""" + ctrl_map: SignalMappingType = {} + for act_id in range(model.nu): + act_name = model.actuator(act_id).name + ctrl_indices = np.arange(act_id, act_id + 1) + ctrl_map[f"{act_name}_ctrl"] = (SignalType.MjCtrl, ctrl_indices) + return ctrl_map + + @staticmethod + def compute_all_state_mappings( + model: mujoco.MjModel, + ) -> tuple[ + SignalMappingType, SignalMappingType, SignalMappingType, SignalMappingType + ]: + """Computes mappings for all state components (qpos, qvel, act) + ctrl.""" + qpos_map: SignalMappingType = {} + qvel_map: SignalMappingType = {} + act_map: SignalMappingType = {} + + nq = model.nq + nv = model.nv + + # Bodies + for body_id in range(model.nbody): + b = model.body(body_id) + body_name = b.name + start_index = model.body_dofadr[body_id] + + if start_index >= 0 and b.dofnum[0] == 6: + qpos_indices = np.arange(start_index, start_index + 7) + qpos_map[f"{body_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices) + qvel_indices = np.arange(start_index + nq, start_index + nq + 6) + qvel_map[f"{body_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices) + + # Joints + for jnt_id in range(model.njnt): + jnt_name = model.joint(jnt_id).name + start_index = model.jnt_qposadr[jnt_id] + jnt_type = model.jnt_type[jnt_id] + + qpos_width = 1 + qvel_width = 1 + if jnt_type == mujoco.mjtJoint.mjJNT_BALL: + qpos_width = 4 + qvel_width = 3 + elif jnt_type == mujoco.mjtJoint.mjJNT_FREE: + continue + + qpos_indices = np.arange(start_index, start_index + qpos_width) + qpos_map[f"{jnt_name}_qpos"] = (SignalType.MjStateQPos, qpos_indices) + qvel_indices = np.arange(start_index + nq, start_index + nq + qvel_width) + qvel_map[f"{jnt_name}_qvel"] = (SignalType.MjStateQVel, qvel_indices) + + # Actuators + for act_id in range(model.nu): + act_name = model.actuator(act_id).name + start_index = model.actuator_actadr[act_id] + num_vals = model.actuator_actnum[act_id] + # if index is -1, the actuator is stateless. + if start_index != -1: + indices = np.arange(start_index + nq + nv, start_index + nq + nv + num_vals) + act_map[f"{act_name}_act"] = (SignalType.MjStateAct, indices) + + ctrl_map = TimeSeries.compute_all_control_mapping(model) + + return qpos_map, qvel_map, act_map, ctrl_map + + @classmethod + def from_custom_map( + cls, + times: np.ndarray, + data: np.ndarray, + signals: Sequence[str | tuple[str, int, SignalType]], + ) -> TimeSeries: + """Construct a TimeSeries from custom data with explicit signal definitions. + + Use this when you have custom signal types (e.g., from a custom modify_residual + function) that are not auto-resolved from a MuJoCo model. You must explicitly + specify the signal names, widths, and types. + + Args: + times: 1-D timestamp array of length N. + data: 2-D array of shape ``(N, D)``. + signals: Defines the layout of the columns in `data`. + - If a list of strings: Each string is a signal name with width 1 and type ``CustomObs``. + - If a list of tuples: Each tuple is ``(name, width, type)``. + + Returns: + A TimeSeries object with the constructed signal mapping. + """ + if data.ndim != 2: + raise ValueError("The 'data' array must be 2-dimensional (Time x Features).") + + signal_mapping_dict: SignalMappingType = {} + current_index = 0 + total_width = 0 + + for item in signals: + if isinstance(item, str): + name = item + width = 1 + sig_type = SignalType.CustomObs + else: + name, width, sig_type = item + + if width <= 0: + raise ValueError(f"Signal '{name}' must have positive width, got {width}.") + + indices = np.arange(current_index, current_index + width) + signal_mapping_dict[name] = (sig_type, indices) + current_index += width + total_width += width + + if total_width != data.shape[1]: + raise ValueError( + f"Total width of signals ({total_width}) does not match " + f"data columns ({data.shape[1]})." + ) + + return cls(times=times, data=data, signal_mapping=signal_mapping_dict) + + @classmethod + def from_names( + cls, + times: np.ndarray, + data: np.ndarray, + model: mujoco.MjModel, + names: Sequence[str | tuple[str, SignalType]] | None = None, + ) -> TimeSeries: + """Construct a TimeSeries for observations (sensors or state) from the model. + + This method automatically resolves signal names (sensors, qpos, qvel, act) from + the MuJoCo model, determining their types and data layout. Use this for standard + observation signals that are defined in the model. + + Args: + times: 1-D timestamps of length N. + data: 2-D array of shape (N, D). + model: MuJoCo model used to auto-resolve signal names and types. + names: Signal names to map. Can be strings or (name, SignalType) tuples. + If None, maps ALL model sensors in sensor address order. + + Warning: + When names=None, data columns MUST match the model's sensor layout + (i.e., data[:, i] corresponds to model.sensordata[i] during simulation). + If your data is in a different order, pass explicit names. + + Raises: + ValueError: If MjCtrl signals are passed (use from_control_names). + """ + if data.ndim != 2: + raise ValueError("The 'data' array must be 2-dimensional (Time x Features).") + + if names is None: + signal_mapping = cls.compute_all_sensor_mapping(model) + # Verify width: assumes data contains ALL sensors in sensor_adr order + if model.nsensordata != data.shape[1]: + raise ValueError( + f"Data columns ({data.shape[1]}) do not match model sensors dim ({model.nsensordata})." + ) + else: + signal_mapping = _resolve_signals( + model, + names, + allowed_types={ + SignalType.MjSensor, + SignalType.MjStateQPos, + SignalType.MjStateQVel, + SignalType.MjStateAct, + }, + ) + # Verify total resolved width + max_idx = 0 + for _, indices in signal_mapping.values(): + if len(indices) > 0: + max_idx = max(max_idx, indices[-1] + 1) + if max_idx != data.shape[1]: + raise ValueError( + f"Resolved signal width ({max_idx}) does not match data columns ({data.shape[1]})." + ) + + return cls(times=times, data=data, signal_mapping=signal_mapping) + + @classmethod + def from_control_names( + cls, + times: np.ndarray, + data: np.ndarray, + model: mujoco.MjModel, + names: Sequence[str | tuple[str, SignalType]] | None = None, + ) -> TimeSeries: + """Construct a TimeSeries for control signals from the model. + + This method automatically resolves control/actuator names from the MuJoCo model, + determining their layout. Use this for control signals (MjCtrl type). + + Args: + times: 1-D timestamps of length N. + data: 2-D array of shape (N, model.nu). + model: MuJoCo model used to auto-resolve actuator names. + names: Actuator names to map. If None, maps ALL actuators in order. + + Warning: + When names=None, data columns MUST match actuator order in the model + (i.e., data[:, i] corresponds to actuator i). Pass explicit names + if your data is in a different order. + """ + if data.ndim != 2: + raise ValueError("The 'data' array must be 2-dimensional (Time x Features).") + + if names is None: + signal_mapping = cls.compute_all_control_mapping(model) + if model.nu != data.shape[1]: + raise ValueError( + f"Data columns ({data.shape[1]}) do not match model controls ({model.nu})." + ) + else: + signal_mapping = _resolve_signals(model, names, allowed_types={SignalType.MjCtrl}) + max_idx = 0 + for _, indices in signal_mapping.values(): + if len(indices) > 0: + max_idx = max(max_idx, indices[-1] + 1) + if max_idx != data.shape[1]: + raise ValueError( + f"Resolved signal width ({max_idx}) does not match data columns ({data.shape[1]})." + ) + + return cls(times=times, data=data, signal_mapping=signal_mapping) + + def get_indices(self, obs_name: str) -> tuple[SignalType, np.ndarray]: + """Look up the signal type and column indices for a named observation.""" + assert self.signal_mapping is not None + if obs_name not in self.signal_mapping: + raise ValueError(f"{obs_name} observation is not in the observation name map.") + return self.signal_mapping[obs_name] + + @classmethod + def create( + cls, + times: np.ndarray, + data: np.ndarray, + signal_mapping: dict[str, tuple[SignalType, np.ndarray | list | int]], + ) -> TimeSeries: + """Construct a TimeSeries, normalising index entries to ``np.ndarray``.""" + normalized: SignalMappingType = {} + for key in signal_mapping: + signal_type, indices = signal_mapping[key] + normalized[key] = (signal_type, np.atleast_1d(indices)) + + return cls(times, data, normalized) + + @classmethod + def slice_by_name(cls, ts: TimeSeries, enabled_sensors: list[str]) -> TimeSeries: + """Return a new TimeSeries containing only the named signals. + + Columns are re-indexed so the resulting ``signal_mapping`` has contiguous + indices starting from 0. + """ + if not ts.signal_mapping: + return ts + + original_indices_to_keep = [] + original_to_new_index_map = {} + all_original_indices = [] + for name in ts.signal_mapping: + all_original_indices.extend(ts.signal_mapping[name][1]) + + # Build a set of indices to keep for quick lookups + kept_indices_set = set() + for name in enabled_sensors: + if name not in ts.signal_mapping: + raise ValueError( + f"Attemping to slice TimeSeries failed. {name} is not in {ts.signal_mapping}." + ) + kept_indices_set.update(ts.signal_mapping[name][1]) + new_index_counter = 0 + for original_index in all_original_indices: + if original_index in kept_indices_set: + original_to_new_index_map[original_index] = new_index_counter + new_index_counter += 1 + original_indices_to_keep.append(original_index) + + data = ts.data[..., original_indices_to_keep] + + trimmed_signal_mapping = {} + for name in enabled_sensors: + metadata, original_indices = ts.signal_mapping[name] + + new_indices = [] + for original_index in original_indices: + new_indices.append(original_to_new_index_map[original_index]) + + trimmed_signal_mapping[name] = (metadata, np.asarray(new_indices)) + + return cls(ts.times, data, trimmed_signal_mapping) + + def __post_init__(self): + """Validate the time series data after initialization. + + Raises: + ValueError: If times is not 1D, if lengths don't match, if times + is not strictly increasing, or if arrays are empty. + """ + if self.times.size == 0: + raise ValueError("Empty arrays are not allowed in TimeSeries") + if self.times.ndim != 1: + raise ValueError(f"times must be a 1D array, got {self.times.ndim}D array") + if len(self.times) != len(self.data): + raise ValueError( + f"Length of times ({len(self.times)}) and data ({len(self.data)}) must match" + ) + if not np.all(np.diff(self.times) > 0): + raise ValueError("times must be strictly increasing") + + def __len__(self) -> int: + return len(self.data) + + def save_to_disk(self, path: str | pathlib.Path) -> None: + """Save the time series data to disk. + + Args: + path: Path where the data will be saved. + """ + np.savez( + path, + times=self.times, + data=self.data, + signal_mapping=np.array(self.signal_mapping, dtype=object), + ) + + def save_to_csv(self, path: str | pathlib.Path) -> None: + """Save the time series data to a CSV file.""" + np.savetxt( + path, + np.concatenate([self.times[:, None], self.data], axis=1), + delimiter=",", + ) + + @classmethod + def load_from_disk(cls, path: str | pathlib.Path) -> TimeSeries: + """Load time series data from disk. + + Args: + path: Path to the saved data. + + Returns: + A new TimeSeries object. + """ + with np.load(path, allow_pickle=True) as npz: + times = npz["times"] + data = npz["data"] + if "signal_mapping" in npz: + signal_mapping = npz["signal_mapping"].item() + else: + signal_mapping = None + + return cls(times=times, data=data, signal_mapping=signal_mapping) + + def interpolate( + self, t: float | np.ndarray, method: InterpolationMethod = "linear" + ) -> np.ndarray: + """Interpolate data at specified time(s). + + This is the core interpolation function used by both get() and resample(). + + Args: + t: Time point(s) at which to interpolate data. + method: Interpolation method to use. + + Returns: + Interpolated data values. + """ + t = np.atleast_1d(np.asarray(t)) + + if method in ("zero_order_hold", "zoh"): + indices = np.searchsorted(self.times, t, side="right") - 1 + indices = np.clip(indices, 0, len(self.times) - 1) + return self.data[indices] + + return scipy.interpolate.interp1d( + self.times, + self.data, + kind=method, + axis=0, + bounds_error=False, + fill_value=(self.data[0], self.data[-1]), # pyright: ignore[reportArgumentType] + assume_sorted=True, + )(t) + + def get( + self, t: float | np.ndarray, method: InterpolationMethod = "linear" + ) -> tuple[np.ndarray, np.ndarray]: + """Get interpolated data at specified time(s). + + This method is useful for querying data at specific timestamps without + creating a new TimeSeries object. + + Args: + t: Time point(s) at which to get data. + method: Interpolation method to use. + + Returns: + Tuple of (times, interpolated_data). + """ + t_orig = np.asarray(t) + t_shape = t_orig.shape + result = self.interpolate(t_orig, method=method) + if t_shape == (): + result = result.squeeze(axis=0) + return t_orig, result + + def resample( + self, + new_times: np.ndarray | None = None, + target_dt: float | None = None, + method: InterpolationMethod = "linear", + ) -> TimeSeries: + """Resample the time series to new timestamps or a specific time interval. + + This method creates a new TimeSeries object with data interpolated at the + specified timestamps. + + Args: + new_times: Optional array of new timestamps. If provided, target_dt is + ignored. + target_dt: Optional time interval for regular resampling. Only used if + new_times is None. + method: Interpolation method to use. + + Returns: + A new TimeSeries object with resampled data. + + Raises: + ValueError: If neither new_times nor target_dt is provided, or if + new_times is not strictly increasing. + """ + # Generate new times if target_dt is provided. + if new_times is None: + if target_dt is None: + raise ValueError("Either new_times or target_dt must be provided") + if target_dt <= 0: + raise ValueError("target_dt must be a positive float") + + # Create evenly spaced timestamps. + new_nsteps = int(np.ceil((self.times[-1] - self.times[0]) / target_dt)) + 1 + new_times = np.linspace(self.times[0], self.times[-1], new_nsteps, endpoint=True) + else: + # Make sure new_times is valid. + if new_times.ndim != 1: + raise ValueError("new_times must be a 1D array") + if not np.all(np.diff(new_times) > 0): + raise ValueError("new_times must be strictly increasing") + + assert new_times is not None + new_data = self.interpolate(new_times, method=method) + return TimeSeries( + times=new_times, data=new_data, signal_mapping=self.signal_mapping + ) + + def remove_from_beginning(self, time_to_remove_s: float) -> TimeSeries: + """Remove time from the beginning of the time series. + + Args: + time_to_remove_s: Time to remove from the beginning of the time series. + + Returns: + A new TimeSeries object with the specified time removed. + """ + if time_to_remove_s < 0: + raise ValueError("time_to_remove_s must be non-negative") + if time_to_remove_s > self.times[-1]: + raise ValueError( + "time_to_remove_s is greater than the duration of the time series" + ) + idx = np.searchsorted(self.times, time_to_remove_s) + times_shifted = self.times[idx:] - self.times[idx] + return TimeSeries( + times=times_shifted, data=self.data[idx:], signal_mapping=self.signal_mapping + ) + + def dt_statistics(self) -> dict[str, float]: + """Calculate statistics about the time intervals. + + Returns: + Dictionary with mean, median, std, min, and max of time intervals. + + Raises: + ValueError: If there are fewer than two timestamps. + """ + if self.times.size < 2: + raise ValueError("Must have at least two timestamps to compute dt statistics.") + dt_values = np.diff(self.times) + stats = {} + for fn in ["mean", "median", "std", "min", "max"]: + stats[fn] = float(getattr(np, fn)(dt_values)) + return stats + + def __repr__(self) -> str: + """Return a string representation of the TimeSeries object.""" + t_start, t_end = self.times[0], self.times[-1] + duration = t_end - t_start + + data_shape = self.data.shape + n_samples = len(self) + + dt_stats = self.dt_statistics() + mean_dt = dt_stats["mean"] + min_dt = dt_stats["min"] + max_dt = dt_stats["max"] + + is_uniform = dt_stats["std"] / mean_dt < 0.01 # Less than 1% variation. + + # Calculate data range (min/max values). + data_min = np.min(self.data) + data_max = np.max(self.data) + data_range = f"[{data_min:.3g}, {data_max:.3g}]" + + parts = [ + "TimeSeries(", + f" samples={n_samples}", + f" shape={data_shape}", + f" time_range=[{t_start:.3g}, {t_end:.3g}] (duration={duration:.3g})", + f" dt={mean_dt:.3g}" + + (" (uniform)" if is_uniform else f" (min={min_dt:.3g}, max={max_dt:.3g})"), + f" data_range={data_range}", + f" signal_mapping={self.signal_mapping}", + ")", + ] + + return "\n".join(parts) diff --git a/python/mujoco/sysid/_src/trajectory.py b/python/mujoco/sysid/_src/trajectory.py new file mode 100644 index 00000000..f8e4ea77 --- /dev/null +++ b/python/mujoco/sysid/_src/trajectory.py @@ -0,0 +1,501 @@ +"""Trajectory data containers for system identification.""" + +from __future__ import annotations + +import dataclasses +import pathlib +from collections.abc import Sequence + +import mujoco +import mujoco.rollout as mj_rollout +import numpy as np +from absl import logging + +from mujoco.sysid._src import timeseries + + +@dataclasses.dataclass(frozen=True) +class SystemTrajectory: + """Encapsulates a trajectory rolled out from a system. + + Attributes: + model: MuJoCo model used to simulate the trajectory. + control: A TimeSeries instance containing control signals. + sensordata: A TimeSeries instance containing sensor data. + initial_state: Initial state of the simulation. Shape (n_state,). + state: Simulation states over time. Shape (n_steps, n_state). Optional for + real robot trajectories. + """ + + model: mujoco.MjModel + control: timeseries.TimeSeries + sensordata: timeseries.TimeSeries + initial_state: np.ndarray + state: timeseries.TimeSeries | None + + def replace(self, **kwargs) -> SystemTrajectory: + """Return a copy with the specified fields replaced.""" + return dataclasses.replace(self, **kwargs) + + def get_sensordata_slice(self, sensor: str = "joint_pos") -> np.ndarray: + """Extract contiguous sensor columns by type. + + Args: + sensor: One of ``"joint_pos"``, ``"joint_vel"``, or ``"joint_torque"``. + + Returns: + 2-D array of shape ``(n_steps, total_sensor_dim)``. + """ + if sensor == "joint_pos": + sensor_type = mujoco.mjtSensor.mjSENS_JOINTPOS + elif sensor == "joint_vel": + sensor_type = mujoco.mjtSensor.mjSENS_JOINTVEL + elif sensor == "joint_torque": + sensor_type = mujoco.mjtSensor.mjSENS_JOINTACTFRC + else: + raise ValueError(f"Unsupported sensor type: {sensor}") + adr = [] + dims = [] + for i in range(self.model.nsensor): + if self.model.sensor(i).type == sensor_type: + sensor_id = self.model.sensor(i).id + adr.append(self.model.sensor_adr[sensor_id]) + dims.append(self.model.sensor_dim[sensor_id]) + sensors = sorted(zip(adr, dims, strict=True), key=lambda x: x[0]) + start = sensors[0][0] + total_dim = sum(d for _, d in sensors) + end = start + total_dim + return self.sensordata.data[:, start:end] + + @property + def sensordim(self) -> int: + """Total number of scalar sensor outputs in the model.""" + return self.model.nsensordata + + def __len__(self) -> int: + """Number of time steps in the trajectory.""" + return len(self.sensordata) + + def save_to_disk(self, path: pathlib.Path) -> None: + save_dict = { + "control_times": self.control.times, + "control_data": self.control.data, + "sensordata_times": self.sensordata.times, + "sensordata_data": self.sensordata.data, + "initial_state": self.initial_state, + } + if self.state is not None: + save_dict["state_times"] = self.state.times + save_dict["state_data"] = self.state.data + save_dict["state_signal_mapping"] = np.array( + self.state.signal_mapping, dtype=object + ) + + if self.control.signal_mapping: + save_dict["control_signal_mapping"] = np.array( + self.control.signal_mapping, dtype=object + ) + + if self.sensordata.signal_mapping: + save_dict["sensordata_signal_mapping"] = np.array( + self.sensordata.signal_mapping, dtype=object + ) + + np.savez(path, **save_dict) # type: ignore + + @classmethod + def load_from_disk( + cls, + path: pathlib.Path, + model: mujoco.MjModel, + allow_missing_sensors: bool = False, + ) -> SystemTrajectory: + with np.load(path, allow_pickle=True) as npz: + control_times = npz["control_times"] + control_data = npz["control_data"] + sensordata_times = npz["sensordata_times"] + sensordata_data = npz["sensordata_data"] + initial_state = npz["initial_state"] + state_times = npz.get("state_times", None) + state_data = npz.get("state_data", None) + + control_signal_mapping = None + if "control_signal_mapping" in npz: + control_signal_mapping = npz["control_signal_mapping"].item() + + sensordata_signal_mapping = None + if "sensordata_signal_mapping" in npz: + sensordata_signal_mapping = npz["sensordata_signal_mapping"].item() + + state_signal_mapping = None + if "state_signal_mapping" in npz: + state_signal_mapping = npz["state_signal_mapping"].item() + + predicted_rollout = cls( + model=model, + control=timeseries.TimeSeries( + control_times, control_data, signal_mapping=control_signal_mapping + ), + sensordata=timeseries.TimeSeries( + sensordata_times, sensordata_data, signal_mapping=sensordata_signal_mapping + ), + initial_state=initial_state, + state=timeseries.TimeSeries( + state_times, state_data, signal_mapping=state_signal_mapping + ) + if state_times is not None + else None, + ) + predicted_rollout.check_compatible(allow_missing_sensors) + return predicted_rollout + + def check_compatible(self, allow_missing_sensors: bool = False) -> None: + """Validate that data dimensions match the model. + + Checks sensor, control, state, and initial-state dimensions. + + Args: + allow_missing_sensors: If True, a sensor dimension mismatch is logged + as a warning instead of raising. + """ + if self.sensordata.data.shape[1] != self.model.nsensordata: + if not allow_missing_sensors: + raise ValueError( + f"Sensor data dimension {self.sensordata.data.shape[1]} does not" + f" match model sensor dimension {self.model.nsensordata}" + ) + else: + logging.warning( + f"Sensor data dimension {self.sensordata.data.shape[1]} does not" + f" match model sensor dimension {self.model.nsensordata}" + ) + + if self.control.data.shape[1] != self.model.nu: + raise ValueError( + f"Control data dimension {self.control.data.shape[1]} does not" + f" match model control dimension {self.model.nu}" + ) + + state_spec = mujoco.mjtState.mjSTATE_FULLPHYSICS + state_size = mujoco.mj_stateSize(self.model, state_spec.value) + if self.state is not None: + if self.state.data.shape[1] != state_size: + raise ValueError( + f"State dimension {self.state.data.shape[1]} does not match " + f"model state dimension {state_size}" + ) + if self.initial_state.shape[0] != state_size: + raise ValueError( + f"Initial state dimension {self.initial_state.shape[0]} does not" + f" match model state dimension {state_size}" + ) + + def split(self, chunk_size: int) -> list[SystemTrajectory]: + """Split into consecutive non-overlapping chunks of *chunk_size* steps. + + Incomplete trailing steps are discarded. Requires ``state`` to be set + (needed to extract the initial state for each chunk). + """ + if self.state is None: + raise ValueError("Cannot split rollout with missing state field.") + steps = len(self.sensordata.times) + n_complete_chunks = steps // chunk_size + control_times = self.control.times + control_data = self.control.data + sensordata_times = self.sensordata.times + sensordata_data = self.sensordata.data + trajectories = [] + for i in range(n_complete_chunks): + start_idx = i * chunk_size + end_idx = start_idx + chunk_size + initial_state = ( + self.initial_state if start_idx == 0 else self.state.data[start_idx - 1] + ) + control_times_chunk = control_times[start_idx:end_idx] + control_data_chunk = control_data[start_idx:end_idx] + sensordata_times_chunk = sensordata_times[start_idx:end_idx] + sensordata_data_chunk = sensordata_data[start_idx:end_idx] + trajectories.append( + SystemTrajectory( + model=self.model, + control=timeseries.TimeSeries(control_times_chunk, control_data_chunk), + sensordata=timeseries.TimeSeries( + sensordata_times_chunk, sensordata_data_chunk + ), + initial_state=initial_state, + state=timeseries.TimeSeries( + times=self.state.times[start_idx:end_idx], + data=self.state.data[start_idx:end_idx], + signal_mapping=self.state.signal_mapping, + ), + ) + ) + return trajectories + + def render( + self, + height: int = 240, + width: int = 320, + camera: str | int = -1, + fps: int = 30, + ) -> list[np.ndarray]: + """Render this trajectory to a list of RGB frames. + + Requires ``state`` to be set. Delegates to + :func:`~mujoco_sysid._src.plotting.render_rollout`. + """ + if self.state is None: + raise ValueError("Cannot render rollout with missing state field.") + + from mujoco.sysid._src.plotting import render_rollout + + # Adapt state to batch format (nbatch=1, nsteps, nstate) + state_batch = self.state.data[np.newaxis, :, :] + + data = mujoco.MjData(self.model) + + return render_rollout( + model=self.model, + data=data, + state=state_batch, + framerate=fps, + camera=camera, + width=width, + height=height, + ) + + +def create_initial_state( + model: mujoco.MjModel, + qpos: np.ndarray, + qvel: np.ndarray | None = None, + act: np.ndarray | None = None, +) -> np.ndarray: + """Build a ``mjSTATE_FULLPHYSICS`` initial-state vector from components. + + Args: + model: MuJoCo model. + qpos: Joint positions, shape ``(nq,)``. + qvel: Joint velocities, shape ``(nv,)``. Defaults to zero. + act: Actuator activations, shape ``(na,)``. Defaults to zero. + + Returns: + Flat state vector suitable for ``mujoco.rollout``. + """ + data = mujoco.MjData(model) + initial_state = np.empty( + (mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS),) + ) + if qpos.shape[0] != model.nq: + raise ValueError(f"Expected qpos to have shape {model.nq}, got {qpos.shape[0]}.") + data.qpos[:] = qpos + if qvel is not None: + if qvel.shape[0] != model.nv: + raise ValueError(f"Expected qvel to have shape {model.nv}, got {qvel.shape[0]}.") + data.qvel[:] = qvel + if act is not None: + if act.shape[0] != model.na: + raise ValueError(f"Expected act to have shape {model.na}, got {act.shape[0]}.") + data.act[:] = act + mujoco.mj_getState(model, data, initial_state, mujoco.mjtState.mjSTATE_FULLPHYSICS) + return initial_state + + +class ModelSequences: + """A model spec paired with one or more measured trajectory sequences. + + Groups a single ``MjSpec`` (the model to be identified) with the + corresponding measured data so that the residual pipeline can iterate + over all sequences for that model. + + Args: + name: Identifier for this model group (used for file-naming on save). + spec: MjSpec that will be recompiled with candidate parameters. + sequence_name: Name(s) identifying each measured sequence. + initial_state: Initial state(s) for each sequence. + control: Measured control TimeSeries for each sequence. + sensordata: Measured sensor TimeSeries for each sequence. + allow_missing_sensors: Passed through to + :meth:`SystemTrajectory.check_compatible`. + """ + + def __init__( + self, + name: str, + spec: mujoco.MjSpec, + sequence_name: str | Sequence[str], + initial_state: np.ndarray | Sequence[np.ndarray], + control: timeseries.TimeSeries | Sequence[timeseries.TimeSeries], + sensordata: timeseries.TimeSeries | Sequence[timeseries.TimeSeries], + allow_missing_sensors: bool = False, + ): + self.name = name + self.spec = spec + self.allow_missing_sensors = allow_missing_sensors + + self.gt_model = self.spec.compile() + + self.sequence_name: list[str] = ( + [sequence_name] if isinstance(sequence_name, str) else list(sequence_name) + ) + self.initial_state: list[np.ndarray] = ( + [initial_state] if isinstance(initial_state, np.ndarray) else list(initial_state) + ) + self.control: list[timeseries.TimeSeries] = ( + [control] if isinstance(control, timeseries.TimeSeries) else list(control) + ) + self.sensordata: list[timeseries.TimeSeries] = ( + [sensordata] + if isinstance(sensordata, timeseries.TimeSeries) + else list(sensordata) + ) + + self.measured_rollout: list[SystemTrajectory] = [] + for initial_state_, control_, sensordata_ in zip( + self.initial_state, self.control, self.sensordata, strict=True + ): + measured_rollout_ = SystemTrajectory( + model=self.gt_model, + control=control_, + sensordata=sensordata_, + initial_state=initial_state_, + state=None, + ) + measured_rollout_.check_compatible(allow_missing_sensors=allow_missing_sensors) + self.measured_rollout.append(measured_rollout_) + + def __getitem__(self, key): + return ModelSequences( + self.name, + self.spec, + self.sequence_name[key], + self.initial_state[key], + self.control[key], + self.sensordata[key], + self.allow_missing_sensors, + ) + + +def timeseries2array( + control_signal: timeseries.TimeSeries | Sequence[timeseries.TimeSeries], +) -> tuple[np.ndarray, np.ndarray]: + if isinstance(control_signal, timeseries.TimeSeries): + control = control_signal.data + control_times = control_signal.times + else: + control = np.stack([ts.data for ts in control_signal], axis=0) + control_times = np.stack([ts.times for ts in control_signal], axis=0) + # The measured data has N sensor measurements and N controls, where the first sensor + # measurement corresponds to the initial condition. Thus we don't have ground truth + # for the N+1'th state produced by the N'th control and so there is no point in + # simulating it. + if control.ndim == 3: + control_applied_times = control_times[:, :-1] + control_applied = control[:, :-1, :] + else: + control_applied_times = control_times[:-1] + control_applied = control[:-1, :] + return control_applied, control_applied_times + + +def sequence2array( + initial_states: np.ndarray | Sequence[np.ndarray], +) -> np.ndarray: + if isinstance(initial_states, np.ndarray): + return initial_states + return np.stack(initial_states, axis=0) + + +def arrays2traj( + models: mujoco.MjModel | Sequence[mujoco.MjModel], + initial_states: np.ndarray | Sequence[np.ndarray], + control: np.ndarray, + control_times: np.ndarray, + state: np.ndarray, + sensordata: np.ndarray, + signal_mapping: timeseries.SignalMappingType, + state_mapping: timeseries.SignalMappingType, + ctrl_mapping: timeseries.SignalMappingType, +) -> Sequence[SystemTrajectory]: + nbatch = state.shape[0] + # TODO(kevin): When is np.tile necessary? + # initial_states = np.tile(initial_states, (nbatch, 1)) + # control = np.tile(control, (nbatch, 1, 1)) + # control_times = np.tile(control_times, (nbatch, 1)) + + if isinstance(models, mujoco.MjModel): + models_list = [models] * nbatch + else: + models_list = list(models) + + return [ + SystemTrajectory( + model=models_list[i], + control=timeseries.TimeSeries( + control_times[i], control[i], signal_mapping=ctrl_mapping + ), + # NOTE(kevin): When using mjSTATE_FULLPHYSICS, the first element of + # the state corresponds to the simulation time. The reason we do not + # use control_times[i] is because sensordata times are shifted by + # one time step. + sensordata=timeseries.TimeSeries(state[i][:, 0], sensordata[i], signal_mapping), + initial_state=initial_states[i], + state=timeseries.TimeSeries( + times=state[i][:, 0], data=state[i], signal_mapping=state_mapping + ), + ) + for i in range(nbatch) + ] + + +def sysid_rollout( + models: mujoco.MjModel | Sequence[mujoco.MjModel], + datas: mujoco.MjData | Sequence[mujoco.MjData], + control_signal: Sequence[timeseries.TimeSeries] | timeseries.TimeSeries, + initial_states: np.ndarray | Sequence[np.ndarray], + rollout_signal_mapping: timeseries.SignalMappingType | None = None, + rollout_state_mapping: timeseries.SignalMappingType | None = None, + ctrl_mapping: timeseries.SignalMappingType | None = None, +) -> Sequence[SystemTrajectory]: + """Rollout trajectories in parallel for the given models and controls. + + Args: + models: MuJoCo model or sequence of models. + datas: MuJoCo data or sequence of data. + control_signal: Control signals as TimeSeries or sequence of TimeSeries. + initial_states: Initial states of the simulation. Shape (n_state,) or + (n_batch, n_state). + + Returns: + Sequence of SystemTrajectory instances containing the simulation results. + """ + + # if the user does not supply it, we create it. Note that this will impact perf. + if not rollout_signal_mapping or not rollout_state_mapping or not ctrl_mapping: + if isinstance(models, mujoco.MjModel): + model0 = models + else: + model0 = models[0] + qpos_map, qvel_map, act_map, ctrl_mapping = ( + timeseries.TimeSeries.compute_all_state_mappings(model0) + ) + rollout_state_mapping = qpos_map | qvel_map | act_map + rollout_signal_mapping = timeseries.TimeSeries.compute_all_sensor_mapping(model0) + + control, control_times = timeseries2array(control_signal) + initial_states = sequence2array(initial_states) + state, sensordata = mj_rollout.rollout(models, datas, initial_states, control) + assert isinstance(state, np.ndarray) + assert isinstance(sensordata, np.ndarray) + + return arrays2traj( + models, + initial_states, + control, + control_times, + state, + sensordata, + rollout_signal_mapping, + rollout_state_mapping, + ctrl_mapping, + ) diff --git a/python/mujoco/sysid/py.typed b/python/mujoco/sysid/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/python/mujoco/sysid/report/builder.py b/python/mujoco/sysid/report/builder.py new file mode 100644 index 00000000..131af433 --- /dev/null +++ b/python/mujoco/sysid/report/builder.py @@ -0,0 +1,104 @@ +# report/builder.py +import os +from typing import Any + +import jinja2 + +from mujoco.sysid.report.sections.base import ReportSection + +# Path to the templates directory +TEMPLATE_DIR = os.path.join(os.path.dirname(__file__), "templates") + + +class ReportBuilder: + def __init__(self, title: str, global_context: dict[str, Any] | None = None): + self._title = title + self._sections: list[ReportSection] = [] + self._global_context = global_context or {} + + # Setup Jinja to load from files + self._env = jinja2.Environment( + loader=jinja2.FileSystemLoader(TEMPLATE_DIR), + autoescape=jinja2.select_autoescape(["html", "xml"]), + ) + + def add_section(self, section: ReportSection): + self._sections.append(section) + + def build(self) -> str: + # Load the main layout + layout_template = self._env.get_template("layout.html") + + # Render each section individually + rendered_sections = [] + all_header_includes = set() + + for section in self._sections: + # Get the specific template for this section + sec_template = self._env.get_template(section.template_filename) + + # Check if this is a GroupSection (has 'sections' attribute) + extra_context = {} + child_sections: list[ReportSection] = getattr(section, "sections", []) + if child_sections: + child_sections_content = [] + for child in child_sections: + child_template = self._env.get_template(child.template_filename) + child_html = child_template.render(child.get_context()) + all_header_includes.update(child.header_includes()) + + # Fallback anchor for child + child_anchor = child.anchor + if not child_anchor: + child_anchor = ( + child.title.lower() + .replace(" ", "-") + .replace("[", "") + .replace("]", "") + .replace("(", "") + .replace(")", "") + ) + + child_sections_content.append( + {"title": child.title, "content": child_html, "anchor": child_anchor} + ) + extra_context["child_sections"] = child_sections_content + + # Render the section HTML (main wrapper) + html_content = sec_template.render(section.get_context() | extra_context) + # Collect header requirements (scripts/css) + all_header_includes.update(section.header_includes()) + + # Fallback anchor generation + anchor = section.anchor + if not anchor: + anchor = ( + section.title.lower() + .replace(" ", "-") + .replace("[", "") + .replace("]", "") + .replace("(", "") + .replace(")", "") + ) + + rendered_sections.append( + { + "title": section.title, + "anchor": anchor, + "collapsible": section.collapsible, + "is_open": section.is_open, + "content": html_content, + } + ) + + # Render final report + return layout_template.render( + report_title=self._title, + sections=rendered_sections, + header_includes=all_header_includes, + **self._global_context, + ) + + def save(self, path: str): + with open(path, "w", encoding="utf-8") as f: + f.write(self.build()) diff --git a/python/mujoco/sysid/report/defaults.py b/python/mujoco/sysid/report/defaults.py new file mode 100644 index 00000000..8882f389 --- /dev/null +++ b/python/mujoco/sysid/report/defaults.py @@ -0,0 +1,376 @@ +# Copyright 2026 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Default report generation for system identification results.""" + +import os +import pathlib +from collections.abc import Sequence + +import matplotlib.pyplot as plt +import numpy as np +import scipy.optimize as scipy_optimize + +from mujoco.sysid._src import model_modifier, parameter, plotting +from mujoco.sysid._src.optimize import calculate_intervals +from mujoco.sysid._src.residual import BuildModelFn +from mujoco.sysid._src.trajectory import ModelSequences +from mujoco.sysid.report.builder import ReportBuilder +from mujoco.sysid.report.sections.covariance import Covariance +from mujoco.sysid.report.sections.optimization_trace import OptimizationTrace +from mujoco.sysid.report.sections.parameters import ParametersTable +from mujoco.sysid.report.sections.signals import SignalReport + + +def default_report( + models_sequences: Sequence[ModelSequences], + initial_params: parameter.ParameterDict, + opt_params: parameter.ParameterDict, + residual_fn, + opt_result: scipy_optimize.OptimizeResult, + title="SysID", + save_path=None, + build_model: BuildModelFn | None = model_modifier.apply_param_modifiers, + generate_videos=True, +) -> ReportBuilder: + """Returns a ReportBuilder containing experiment results. + + Users needing a custom report can copy and modify this code. + """ + from mujoco.sysid.report.sections.group import GroupSection + from mujoco.sysid.report.sections.insights import AutomatedInsights + from mujoco.sysid.report.sections.parameter_distribution import ParameterDistribution + from mujoco.sysid.report.sections.row import RowSection + from mujoco.sysid.report.sections.video import ( + VideoPlayer, + generate_video_from_trajectories, + ) + + #################################### + # Build report + # Sections: + # Fit + # Parameter tables + # Confidence intervals + # Extras: Optimization trace + #################################### + rb = ReportBuilder(title) + + if generate_videos: + # 1. Video Player + if save_path is None: + raise ValueError("save_path is required when generate_videos=True") + if build_model is None: + raise ValueError("build_model is required when generate_videos=True") + + # Collect ALL trajectories from all model sequences + all_trajectories = [] + for model_sequences in models_sequences: + for traj in model_sequences.measured_rollout: + all_trajectories.append(traj) + + # Use first model's spec for rendering + model_spec_to_render = models_sequences[0].spec + + video_dir = pathlib.Path(save_path) + video_dir.mkdir(parents=True, exist_ok=True) + + # Video 1: All (Initial + Nominal + Optimized) - all trajectories concatenated + video_all_path = video_dir / "video_all.mp4" + generate_video_from_trajectories( + initial_params=initial_params, + opt_params=opt_params, + build_model=build_model, + trajectories=all_trajectories, + model_spec=model_spec_to_render, + output_filepath=video_all_path, + fps=60, + ) + + # Video 2: Initial + Nominal (no optimized) + video_init_path = video_dir / "video_init.mp4" + generate_video_from_trajectories( + initial_params=initial_params, + opt_params=opt_params, + build_model=build_model, + trajectories=all_trajectories, + model_spec=model_spec_to_render, + output_filepath=video_init_path, + render_opt=False, + fps=60, + ) + + # Video 3: Optimized + Nominal (no initial) + video_opt_path = video_dir / "video_opt.mp4" + generate_video_from_trajectories( + initial_params=initial_params, + opt_params=opt_params, + build_model=build_model, + trajectories=all_trajectories, + model_spec=model_spec_to_render, + output_filepath=video_opt_path, + render_initial=False, + fps=60, + ) + + video_all_section = VideoPlayer( + title="All Models", + video_filepath=video_all_path, + anchor="visual_run_all", + autoplay=True, + muted=True, + width="100%", + height=None, + caption="Initial, Nominal, Optimized", + ) + + video_init_section = VideoPlayer( + title="Initial vs Nominal", + video_filepath=video_init_path, + anchor="visual_run_init", + autoplay=True, + muted=True, + width="100%", + height=None, + caption="Initial, Nominal", + ) + + video_opt_section = VideoPlayer( + title="Optimized vs Nominal", + video_filepath=video_opt_path, + anchor="visual_run_opt", + autoplay=True, + muted=True, + width="100%", + height=None, + caption="Nominal, Optimized", + ) + + rb.add_section( + RowSection( + title="Visual Comparison", + sections=[video_all_section, video_init_section, video_opt_section], + anchor="visual_comparison", + description="Visual comparison of the system identification results. The nominal model is shown in green, the initial model in red, and the optimized model in blue.", + ) + ) + + # 2. Automated Insights (Logs) + rb.add_section(AutomatedInsights("Automated Insights", opt_params)) + + # 3. Parameters Table (Unified) + rb.add_section( + ParametersTable("Parameters", opt_params, initial_params, anchor="Parameters") + ) + + # 4. Control Signals (per sequence, grouped like observations) + # Get predictions for initial solution. + names = [ + f"{model_sequences.name}\n{sequence}" + for model_sequences in models_sequences + for sequence in model_sequences.sequence_name + ] + _, pred0s, _ = residual_fn( + initial_params.as_vector(), initial_params, return_pred_all=True + ) + + residuals_star, preds_star, records_star = residual_fn( + opt_params.as_vector(), opt_params, return_pred_all=True + ) + + assert build_model is not None + model_hat = build_model(initial_params, models_sequences[0].spec) + + # Build control signal reports for each sequence + control_reports = [] + seq_idx = 0 + for model_sequences in models_sequences: + for i, seq_name in enumerate(model_sequences.sequence_name): + ctrl_ts = model_sequences.control[i] + name = f"{model_sequences.name}\n{seq_name}" + control_reports.append( + SignalReport( + f"Sequence: {name}", + model_hat, + title_prefix="", + ts_dict={"control": ctrl_ts}, + collapsible=True, + ) + ) + seq_idx += 1 + + rb.add_section( + GroupSection("Control Signals", control_reports, anchor="control_signals") + ) + + # 5. Observation Signals + observation_reports = [] + for name, pred, record, pred0 in zip( + names, preds_star, records_star, pred0s, strict=True + ): + obs_dict = {"initial": pred0[0], "nominal": record[0], "fitted": pred[0]} + observation_reports.append( + SignalReport( + f"Sequence: {name}", + model_hat, + title_prefix="", + ts_dict=obs_dict, + collapsible=True, + ) + ) + + rb.add_section( + GroupSection("Observation Signals", observation_reports, anchor="observations") + ) + + covariance, intervals = calculate_intervals(residuals_star, opt_result.jac) + + # 6. Parameter Distribution + rb.add_section( + ParameterDistribution( + title="Parameter Distribution", + opt_params=opt_params, + initial_params=initial_params, + confidence_intervals=intervals, + anchor="param_dist", + ) + ) + + rb.add_section( + Covariance( + title="Covariance and Correlation", + anchor="cov", + covariance=covariance, + parameter_dict=opt_params, + ) + ) + + # Add diagnostic optimization trace plots. + if "extras" in opt_result: + # Add to the report. + rb.add_section( + OptimizationTrace( + title="Optimization Trace", + anchor="opt", + objective=opt_result.extras.get("objective"), + candidate=opt_result.extras.get("candidate"), + bounds=opt_params.get_bounds(), + param_names=opt_params.get_non_frozen_parameter_names(), + ) + ) + + rb.build() + if save_path: + rb.save(save_path / "report.html") + return rb + + +# TODO(nimrod): Consider deleting this function, given we can export plots from +# plotly either on the web or with fig.write_image. +def default_report_matplotlib( + experiment_results_folder: os.PathLike, + models_sequences: Sequence[ModelSequences], + params: parameter.ParameterDict, + sysid_residual, + x0: np.ndarray, + opt_result: scipy_optimize.OptimizeResult, + build_model: BuildModelFn | None = model_modifier.apply_param_modifiers, +): + """Outputs PNG plots to the experiment results folder.""" + experiment_results_folder = pathlib.Path(experiment_results_folder) + if not experiment_results_folder.exists(): + experiment_results_folder.mkdir(parents=True, exist_ok=True) + + x_hat = opt_result.x + params.update_from_vector(x_hat) + + # Save the ID'd models out + assert build_model is not None + model_hat = None + for model_sequences in models_sequences: + model_hat = build_model(params, model_sequences.spec) + assert model_hat is not None + + # Get predictions for initial solution. + params.update_from_vector(x0) + names = [ + f"{model_sequences.name}\n{sequence}" + for model_sequences in models_sequences + for sequence in model_sequences.sequence_name + ] + _, pred0s, record0s = sysid_residual(x0, return_pred_all=True) + + for name, pred0, record0 in zip(names, pred0s, record0s, strict=True): + plotting.plot_sensor_comparison( + model_hat, + predicted_times=pred0[0].times, + predicted_data=pred0[0].data, + real_times=record0[0].times, + real_data=record0[0].data, + title_prefix=f"x0 {name}", + size_factor=0.5, + ) + name_fig = name.replace("/", " ") + name_fig = name_fig.replace("\n", " ") + plt.savefig(os.path.join(experiment_results_folder, f"x0 {name_fig}.png")) + + residuals_star, preds_star, records_star = sysid_residual(x_hat, return_pred_all=True) + for name, pred, record, _pred0 in zip( + names, preds_star, records_star, pred0s, strict=True + ): + plotting.plot_sensor_comparison( + model_hat, + predicted_times=pred[0].times, + predicted_data=pred[0].data, + real_times=record[0].times, + real_data=record[0].data, + title_prefix=f"x* {name}", + size_factor=0.5, + ) + name_fig = name.replace("/", " ") + name_fig = name_fig.replace("\n", " ") + plt.savefig(experiment_results_folder / f"xstar {name_fig}.png") + + # Add diagnostic optimization trace plots. + if "extras" in opt_result: + # Objective value over iterations. + objective = opt_result.extras["objective"] + plotting.plot_objective(objective) + plt.savefig(experiment_results_folder / "loss.png", dpi=300) + + # Candidate parameter values over iterations. + candidate = opt_result.extras["candidate"] + + # Candidate parameter values over iterations. + # Candidate heatmap over iterations. + plotting.plot_candidate_heatmap( + candidate, + param_names=params.get_non_frozen_parameter_names(), + bounds=params.get_bounds(), + ) + plt.savefig(experiment_results_folder / "candidate_heatmap.png", dpi=300) + + plotting.plot_candidate( + candidate, + bounds=params.get_bounds(), + param_names=params.get_non_frozen_parameter_names(), + ) + plt.savefig(experiment_results_folder / "candidate.png", dpi=300) + + _, intervals = calculate_intervals(residuals_star, opt_result.jac) + plotting.parameter_confidence( + all_exp_names=["trial"], all_params=[params], all_intervals=[intervals] + ) + # plotting.parameter_confidence(["trial"], [params], [x_hat], [intervals]) + plt.savefig(experiment_results_folder / "params.png") diff --git a/python/mujoco/sysid/report/sections/__init__.py b/python/mujoco/sysid/report/sections/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python/mujoco/sysid/report/sections/base.py b/python/mujoco/sysid/report/sections/base.py new file mode 100644 index 00000000..445cb0cb --- /dev/null +++ b/python/mujoco/sysid/report/sections/base.py @@ -0,0 +1,47 @@ +import abc +from collections.abc import Iterable +from typing import Any + + +class ReportSection(abc.ABC): + """Abstract base class for all report sections.""" + + @property + @abc.abstractmethod + def template_filename(self) -> str: + """The filename of the Jinja2 template (e.g., 'parameters.html').""" + pass + + @abc.abstractmethod + def get_context(self) -> dict[str, Any]: + """Returns data needed by the template.""" + pass + + def __init__(self, collapsible: bool = True, is_open: bool = True): + self._collapsible = collapsible + self._is_open = is_open + self._anchor = "" + + @property + def title(self) -> str: + return "" + + @property + def anchor(self) -> str: + """Returns a unique HTML anchor string.""" + # Auto-generate a safe anchor from title if not provided + if not hasattr(self, "_anchor") or not self._anchor: + return self.title.lower().replace(" ", "-") + return self._anchor + + @property + def collapsible(self) -> bool: + return self._collapsible + + @property + def is_open(self) -> bool: + return self._is_open + + def header_includes(self) -> Iterable[str]: + """Returns strings (like + + + + \ No newline at end of file diff --git a/python/mujoco/sysid/report/templates/optimization_trace.html b/python/mujoco/sysid/report/templates/optimization_trace.html new file mode 100644 index 00000000..df228348 --- /dev/null +++ b/python/mujoco/sysid/report/templates/optimization_trace.html @@ -0,0 +1,20 @@ +{% if objective_plot_html %} +

Objective Function

+
+ {{ objective_plot_html | safe }} +
+{% endif %} +{% if candidate_heatmap_html %} +

Parameter Candidate Heatmap

+
+ {{ candidate_heatmap_html | safe }} +
+{% endif %} +{% if candidate_plots_html %} +

Parameter Candidates

+{% for plot_html in candidate_plots_html %} +
+ {{ plot_html | safe }} +
+{% endfor %} +{% endif %} \ No newline at end of file diff --git a/python/mujoco/sysid/report/templates/parameter_confidence.html b/python/mujoco/sysid/report/templates/parameter_confidence.html new file mode 100644 index 00000000..d7b43daf --- /dev/null +++ b/python/mujoco/sysid/report/templates/parameter_confidence.html @@ -0,0 +1,3 @@ +
+ {{ plot_html | safe }} +
diff --git a/python/mujoco/sysid/report/templates/parameters_table.html b/python/mujoco/sysid/report/templates/parameters_table.html new file mode 100644 index 00000000..9be7e4c2 --- /dev/null +++ b/python/mujoco/sysid/report/templates/parameters_table.html @@ -0,0 +1,147 @@ + +
+ + + + {% for header in headers %} + + {% endfor %} + + + + {% for row in table_data %} + + {# Parameter + name #} + {% if row.initial is defined %} + + {% endif %} + {% if row.nominal is defined %} + + {% endif %} + + {% if row.is_frozen %} + + + {% else %} + + + {% endif %} + {% if row.abs_err is defined %} + + {% endif %} + {# Format relative error as percentage #} + {% if row.rel_err is defined %} + + {% endif %} + + {% endfor %} + +
{{ header }}
{{ row.name }}{{ "%.4f" | format(row.initial) }}{{ "%.4f" | format(row.nominal) }}{{ "%.4f" | format(row.pred) }}--{{ "%.4f" | format(row.lower_bound) }}{{ "%.4f" | format(row.upper_bound) }}{{ "%.4f" | format(row.abs_err) }}{{ "%.1f%%" | format(row.rel_err * 100) }}
+ {% if overall_rmse is not none %} +

Overall Parameters RMSE: {{ "%.4f" | format(overall_rmse) }}

+ {% endif %} +

* parameter is frozen

+
\ No newline at end of file diff --git a/python/mujoco/sysid/report/templates/plot_generic.html b/python/mujoco/sysid/report/templates/plot_generic.html new file mode 100644 index 00000000..7619a363 --- /dev/null +++ b/python/mujoco/sysid/report/templates/plot_generic.html @@ -0,0 +1,8 @@ +
+ {{ plot_div | safe }} + {% if caption %} +
+ {{ caption | safe }} +
+ {% endif %} +
\ No newline at end of file diff --git a/python/mujoco/sysid/report/templates/row.html b/python/mujoco/sysid/report/templates/row.html new file mode 100644 index 00000000..6f7f94e6 --- /dev/null +++ b/python/mujoco/sysid/report/templates/row.html @@ -0,0 +1,13 @@ +
+ {% if description %} +

{{ description | safe }}

+ {% endif %} +
+ {% for child in child_sections %} +
+ {{ child.content | safe }} +
+ {% endfor %} +
+
\ No newline at end of file diff --git a/python/mujoco/sysid/report/templates/signals.html b/python/mujoco/sysid/report/templates/signals.html new file mode 100644 index 00000000..d0bab4a2 --- /dev/null +++ b/python/mujoco/sysid/report/templates/signals.html @@ -0,0 +1,3 @@ +
+ {{ plot_div | safe }} +
diff --git a/python/mujoco/sysid/report/templates/video.html b/python/mujoco/sysid/report/templates/video.html new file mode 100644 index 00000000..7523477d --- /dev/null +++ b/python/mujoco/sysid/report/templates/video.html @@ -0,0 +1,11 @@ +
+ + {% if caption %} +

{{ caption | safe }} +

+ {% endif %} +
\ No newline at end of file diff --git a/python/mujoco/sysid/report/utils.py b/python/mujoco/sysid/report/utils.py new file mode 100644 index 00000000..2f774a2f --- /dev/null +++ b/python/mujoco/sysid/report/utils.py @@ -0,0 +1,41 @@ +import math + +from plotly import offline as plt_offline + + +def plotly_script_tag() -> str: + """ + Returns an HTML ' + ) + + +def get_text_color(bg_color_hex: str) -> str: + """ + Determines whether text should be 'black' or 'white' based on the + luminance of the given background hex color. + Useful for heatmaps and colored data tables. + """ + # Convert hex to RGB + hex_color = bg_color_hex.lstrip("#") + if len(hex_color) != 6: + return "black" # Fallback + + rgb = tuple(int(hex_color[i : i + 2], 16) for i in (0, 2, 4)) + r, g, b = [x / 255.0 for x in rgb] + + # Calculate luminance (per WCAG guidelines) + # https://www.w3.org/TR/WCAG20/#relativeluminancedef + def lum_component(c): + return c / 12.92 if c <= 0.03928 else math.pow((c + 0.055) / 1.055, 2.4) + + luminance = ( + 0.2126 * lum_component(r) + 0.7152 * lum_component(g) + 0.0722 * lum_component(b) + ) + + # Return 'black' for light backgrounds, 'white' for dark backgrounds + return "black" if luminance > 0.4 else "white" diff --git a/python/mujoco/sysid/tests/__init__.py b/python/mujoco/sysid/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/python/mujoco/sysid/tests/conftest.py b/python/mujoco/sysid/tests/conftest.py new file mode 100644 index 00000000..843a2830 --- /dev/null +++ b/python/mujoco/sysid/tests/conftest.py @@ -0,0 +1,219 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Shared fixtures for mujoco.sysid tests.""" + +import mujoco +import numpy as np +import pytest + +from mujoco.sysid._src import parameter, timeseries +from mujoco.sysid._src.model_modifier import _infer_inertial + +# --------------------------------------------------------------------------- +# Inline model XML strings — no external file dependencies +# --------------------------------------------------------------------------- + +BOX_XML = """\ + + + + + + + + + + + + + + + + + + + + + + + + +""" + +ARM_XML = """\ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +""" + +OSCILLATOR_XML = """\ + + + + + + + + + + +""" + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def box_spec() -> mujoco.MjSpec: + return mujoco.MjSpec.from_string(BOX_XML) + + +@pytest.fixture +def box_model(box_spec) -> mujoco.MjModel: + return box_spec.compile() + + +@pytest.fixture +def arm_spec() -> mujoco.MjSpec: + """Minimal 5-joint arm with sensors, actuators, textures/materials.""" + return mujoco.MjSpec.from_string(ARM_XML) + + +@pytest.fixture +def arm_model(arm_spec) -> mujoco.MjSpec: + return arm_spec.compile() + + +@pytest.fixture +def oscillator_spec() -> mujoco.MjSpec: + """Single-body oscillator with implicit (geom-based) inertia.""" + return mujoco.MjSpec.from_string(OSCILLATOR_XML) + + +@pytest.fixture +def simple_timeseries() -> timeseries.TimeSeries: + """A TimeSeries with 5 data points and 2 columns: y = [x^2, 2*x^2].""" + times = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) + data = np.array( + [ + [0.0, 0.0], + [1.0, 2.0], + [4.0, 8.0], + [9.0, 18.0], + [16.0, 32.0], + ] + ) + return timeseries.TimeSeries(times=times, data=data) + + +@pytest.fixture +def box_params(box_spec) -> parameter.ParameterDict: + """ParameterDict for box model with modifier callbacks.""" + pdict = parameter.ParameterDict() + pdict.add( + parameter.Parameter( + "box_mass", + [5], + min_value=[4.5], + max_value=[5.5], + modifier=lambda s, p: setattr(_infer_inertial(s, "box"), "mass", p.value[0]), + ) + ) + pdict.add( + parameter.Parameter( + "solref1", + [0.01], + min_value=[0.002], + max_value=[0.02], + modifier=lambda s, p: s.pair("box_floor").solref.__setitem__(0, p.value[0]), + ) + ) + pdict.add( + parameter.Parameter( + "friction2", + [0.005], + min_value=[0], + max_value=[0.01], + modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(1, p.value[0]), + ) + ) + return pdict diff --git a/python/mujoco/sysid/tests/test_integration.py b/python/mujoco/sysid/tests/test_integration.py new file mode 100644 index 00000000..33197f20 --- /dev/null +++ b/python/mujoco/sysid/tests/test_integration.py @@ -0,0 +1,237 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""End-to-end integration test using the box model.""" + +import pathlib +import tempfile + +import mujoco +import mujoco.rollout as mj_rollout +import numpy as np + +from mujoco.sysid import ( + ModelSequences, + Parameter, + ParameterDict, + build_residual_fn, + create_initial_state, + optimize, + save_results, +) +from mujoco.sysid._src import signal_modifier, timeseries +from mujoco.sysid._src.model_modifier import _infer_inertial +from mujoco.sysid.tests.conftest import BOX_XML + + +def _generate_box_data( + spec: mujoco.MjSpec, + duration: float = 1.0, +) -> tuple[timeseries.TimeSeries, timeseries.TimeSeries, np.ndarray]: + """Generate synthetic box-pushing data via rollout.""" + model = spec.compile() + data = mujoco.MjData(model) + + n_steps = int(duration / model.opt.timestep) + t = np.arange(n_steps) * model.opt.timestep + force = (np.sin(t) * 3.0).reshape(-1, 1) + control_ts = timeseries.TimeSeries(t, force) + + initial_state = create_initial_state(model, data.qpos, data.qvel, data.act) + + control_applied = force[:-1] + state, _ = mj_rollout.rollout(model, data, initial_state, control_applied) + state = np.squeeze(state, axis=0) + + sensor_ids = [1, 8] + signal_mapping = { + "pos_x": (timeseries.SignalType.MjStateQPos, np.array([0])), + "vel_x": (timeseries.SignalType.MjStateQVel, np.array([1])), + } + state_times = state[:, 0] + sensordata = timeseries.TimeSeries( + state_times, + state[:, sensor_ids], + signal_mapping, + ) + + return control_ts, sensordata, initial_state + + +def _build_box_params() -> ParameterDict: + """Build parameter dict with modifier callbacks matching box config.""" + pdict = ParameterDict() + + pdict.add( + Parameter( + "box_mass", + [5], + min_value=[4.5], + max_value=[5.5], + modifier=lambda s, p: setattr(_infer_inertial(s, "box"), "mass", p.value[0]), + ) + ) + pdict.add( + Parameter( + "solref1", + [0.01], + min_value=[0.002], + max_value=[0.02], + modifier=lambda s, p: s.pair("box_floor").solref.__setitem__(0, p.value[0]), + ) + ) + pdict.add( + Parameter( + "solref2", + [1.0], + min_value=[0.3], + max_value=[1.7], + frozen=True, + modifier=lambda s, p: s.pair("box_floor").solref.__setitem__(1, p.value[0]), + ) + ) + pdict.add( + Parameter( + "friction1", + [1.6], + min_value=[0], + max_value=[3.0], + frozen=True, + modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(0, p.value[0]), + ) + ) + pdict.add( + Parameter( + "friction2", + [0.005], + min_value=[0], + max_value=[0.01], + modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(1, p.value[0]), + ) + ) + pdict.add( + Parameter( + "friction3", + [0.0001], + min_value=[0], + max_value=[0.001], + frozen=True, + modifier=lambda s, p: s.pair("box_floor").friction.__setitem__(2, p.value[0]), + ) + ) + return pdict + + +def test_box_end_to_end(): + """Full 5-stage pipeline: generate data, build residual, optimize 3 iters, save.""" + spec = mujoco.MjSpec.from_string(BOX_XML) + + # 1. Generate synthetic ground-truth data. + control, sensordata, initial_state = _generate_box_data(spec, duration=1.0) + + # 2. Build config with known parameters. + params = _build_box_params() + + # 3. Create ModelSequences. + models_sequences = [ + ModelSequences( + "box", + spec, + "push", + initial_state, + control, + sensordata, + allow_missing_sensors=True, + ) + ] + + # 4. Define modify_residual (box uses state-based residual). + def modify_residual( + params, + sensordata_predicted, + sensordata_measured, + model, + return_pred_all, + state=None, + **kwargs, + ): + assert state is not None + sensor_ids = [1, 8] + statedata_predicted = timeseries.TimeSeries( + state[:, 0], + state[..., sensor_ids], + { + "pos_x": (timeseries.SignalType.MjStateQPos, np.array([0])), + "vel_x": (timeseries.SignalType.MjStateQVel, np.array([1])), + }, + ) + sensordata_measured = signal_modifier.apply_delayed_ts_window( + sensordata_measured, statedata_predicted, 0.0, 0.0 + ) + statedata_predicted = statedata_predicted.resample(sensordata_measured.times) + res = signal_modifier.weighted_diff( + predicted_data=statedata_predicted.data, + measured_data=sensordata_measured.data, + model=model, + ) + res = signal_modifier.normalize_residual(res, sensordata_measured.data) + return res, statedata_predicted, sensordata_measured + + residual_fn = build_residual_fn( + models_sequences=models_sequences, + modify_residual=modify_residual, + ) + + # 5. Perturb params from nominal and optimize (just 3 iters to verify it runs). + rng = np.random.default_rng(42) + params.randomize(rng=rng) + + # Compute initial cost. + initial_residuals, _, _ = residual_fn(params.as_vector(), params) + initial_cost = sum(np.sum(r**2) for r in initial_residuals) + + opt_params, opt_result = optimize( + initial_params=params, + residual_fn=residual_fn, + optimizer="mujoco", + max_iters=3, + verbose=False, + ) + + # 6. Assert basic properties. + assert opt_result.x.shape == params.as_vector().shape + + # Compute final cost. + final_residuals, _, _ = residual_fn(opt_result.x, opt_params) + final_cost = sum(np.sum(r**2) for r in final_residuals) + assert final_cost <= initial_cost, ( + f"Cost should decrease: {final_cost} > {initial_cost}" + ) + + # 7. Save results to a temp dir. + with tempfile.TemporaryDirectory() as tmpdir: + save_results( + experiment_results_folder=tmpdir, + models_sequences=models_sequences, + initial_params=params, + opt_params=opt_params, + opt_result=opt_result, + residual_fn=residual_fn, + ) + result_dir = pathlib.Path(tmpdir) + assert (result_dir / "params_x_0.yaml").exists() + assert (result_dir / "params_x_hat.yaml").exists() + assert (result_dir / "results.pkl").exists() + assert (result_dir / "confidence.pkl").exists() + assert (result_dir / "box.xml").exists() diff --git a/python/mujoco/sysid/tests/test_model_modifier.py b/python/mujoco/sysid/tests/test_model_modifier.py new file mode 100644 index 00000000..a76236e5 --- /dev/null +++ b/python/mujoco/sysid/tests/test_model_modifier.py @@ -0,0 +1,122 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the model_modifier module.""" + +import mujoco +import numpy as np + +from mujoco.sysid._src import model_modifier + + +def test_apply_pgain(arm_spec): + """Setting a P gain correctly configures the underlying actuator parameters.""" + actuator_name = "actuator5" + pgain_value = 74 + + modified_spec = model_modifier.apply_pgain(arm_spec, actuator_name, pgain_value) + model = modified_spec.compile() + + assert model.actuator(actuator_name).gainprm[0] == pgain_value + assert model.actuator(actuator_name).biasprm[1] == -pgain_value + + +def test_apply_dgain(arm_spec): + """Setting a D gain correctly configures the underlying actuator parameters.""" + actuator_name = "actuator5" + dgain_value = 1.2 + + modified_spec = model_modifier.apply_dgain(arm_spec, actuator_name, dgain_value) + model = modified_spec.compile() + + assert model.actuator(actuator_name).biasprm[2] == -dgain_value + + +def test_apply_pdgain(arm_spec): + """Setting P and D gains together from a single array configures both correctly.""" + actuator_name = "actuator5" + pdgain_value = np.array([74, 1.2]) + + modified_spec = model_modifier.apply_pdgain(arm_spec, actuator_name, pdgain_value) + model = modified_spec.compile() + + assert model.actuator(actuator_name).gainprm[0] == pdgain_value[0] + assert model.actuator(actuator_name).biasprm[1] == -pdgain_value[0] + assert model.actuator(actuator_name).biasprm[2] == -pdgain_value[1] + + +def test_apply_body_mass_explicit(arm_spec): + """Bodies with inertia defined in XML: changing mass proportionally scales inertia.""" + body_name = "link1" + model = arm_spec.compile() + original_mass = model.body(body_name).mass[0] + original_inertia = model.body(body_name).inertia + del model + + scale = 3.3 + new_mass = scale * original_mass + + modified_spec = model_modifier.apply_body_mass_ipos( + arm_spec, body_name, mass=new_mass, rot_inertia_scale=True + ) + model = modified_spec.compile() + + assert model.body(body_name).mass == new_mass + np.testing.assert_allclose(model.body(body_name).inertia, original_inertia * scale) + + +def test_apply_body_mass_implicit(oscillator_spec): + """Bodies with inertia inferred from geoms: changing mass proportionally scales inertia.""" + body_name = "mass" + model = oscillator_spec.compile() + original_mass = model.body(body_name).mass[0] + original_inertia = model.body(body_name).inertia + del model + + scale = 0.077 + new_mass = scale * original_mass + + modified_spec = model_modifier.apply_body_mass_ipos( + oscillator_spec, body_name, mass=new_mass, rot_inertia_scale=True + ) + model = modified_spec.compile() + + assert model.body(body_name).mass == new_mass + np.testing.assert_allclose(model.body(body_name).inertia, original_inertia * scale) + + +def test_remove_visuals(arm_spec): + """Stripping visuals removes all textures and materials for faster compilation.""" + cleaned_spec = model_modifier.remove_visuals(arm_spec) + assert len(cleaned_spec.textures) == 0 + assert len(cleaned_spec.materials) == 0 + + +def test_apply_param_modifiers(box_spec, box_params): + """The full modifier pipeline applies parameter callbacks and produces an updated model.""" + spec = box_spec.copy() + original_model = spec.compile() + original_mass = original_model.body("box").mass[0] + + # Change box_mass parameter. + box_params["box_mass"].update_from_vector(np.array([5.3])) + + modified_model = model_modifier.apply_param_modifiers(box_params, spec) + assert modified_model.body("box").mass[0] != original_mass + np.testing.assert_allclose(modified_model.body("box").mass[0], 5.3, atol=1e-6) + + # Also verify apply_param_modifiers_spec returns MjSpec. + spec2 = box_spec.copy() + result = model_modifier.apply_param_modifiers_spec(box_params, spec2) + assert isinstance(result, mujoco.MjSpec) diff --git a/python/mujoco/sysid/tests/test_parameter.py b/python/mujoco/sysid/tests/test_parameter.py new file mode 100644 index 00000000..b73aa720 --- /dev/null +++ b/python/mujoco/sysid/tests/test_parameter.py @@ -0,0 +1,149 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the Parameter and ParameterDict classes.""" + +import numpy as np + +from mujoco.sysid._src import parameter + + +def test_scalar_parameter(): + """A single-valued parameter round-trips through vector conversion, sampling, and reset.""" + param = parameter.Parameter("test", 1.0, 0.5, 2.0) + + assert param.name == "test" + assert param.size == 1 + assert param.shape == (1,) + assert param.nominal == 1.0 + assert param.value == 1.0 + assert param.min_value == 0.5 + assert param.max_value == 2.0 + + np.testing.assert_array_equal(param.as_vector(), [1.0]) + param.update_from_vector(np.array([1.5])) + np.testing.assert_array_equal(param.value, [1.5]) + np.testing.assert_array_equal(param.as_vector(), [1.5]) + + lower, upper = param.get_bounds() + np.testing.assert_array_equal(lower, [0.5]) + np.testing.assert_array_equal(upper, [2.0]) + + param.reset() + np.testing.assert_array_equal(param.value, [1.0]) + + rng = np.random.default_rng(42) + sample = param.sample(rng) + assert 0.5 <= sample[0] <= 2.0 + + +def test_vector_parameter(): + """A multi-valued parameter preserves element-wise bounds and resets correctly.""" + param = parameter.Parameter("test_vector", [1.0, 2.0], [0.5, 1.0], [2.0, 3.0]) + + assert param.name == "test_vector" + assert param.size == 2 + assert param.shape == (2,) + np.testing.assert_array_equal(param.nominal, [1.0, 2.0]) + np.testing.assert_array_equal(param.value, [1.0, 2.0]) + np.testing.assert_array_equal(param.min_value, [0.5, 1.0]) + np.testing.assert_array_equal(param.max_value, [2.0, 3.0]) + + np.testing.assert_array_equal(param.as_vector(), [1.0, 2.0]) + param.update_from_vector(np.array([1.5, 2.5])) + np.testing.assert_array_equal(param.value, [1.5, 2.5]) + + lower, upper = param.get_bounds() + np.testing.assert_array_equal(lower, [0.5, 1.0]) + np.testing.assert_array_equal(upper, [2.0, 3.0]) + + param.reset() + np.testing.assert_array_equal(param.value, [1.0, 2.0]) + + +def test_parameter_dict(): + """A dict of mixed scalar/vector params flattens to one vector and reconstructs.""" + param1 = parameter.Parameter("param1", 1.0, 0.5, 2.0) + param2 = parameter.Parameter("param2", [2.0, 3.0], [1.0, 2.0], [3.0, 4.0]) + params = parameter.ParameterDict({"param1": param1, "param2": param2}) + + assert params.size == 3 # 1 + 2 + assert len(params) == 2 + + assert params["param1"] is param1 + assert params["param2"] is param2 + + np.testing.assert_array_equal(params.as_vector(), [1.0, 2.0, 3.0]) + + params.update_from_vector(np.array([1.5, 2.5, 3.5])) + np.testing.assert_array_equal(params["param1"].value, [1.5]) + np.testing.assert_array_equal(params["param2"].value, [2.5, 3.5]) + + lower, upper = params.get_bounds() + np.testing.assert_array_equal(lower, [0.5, 1.0, 2.0]) + np.testing.assert_array_equal(upper, [2.0, 3.0, 4.0]) + + params.reset() + np.testing.assert_array_equal(params["param1"].value, [1.0]) + np.testing.assert_array_equal(params["param2"].value, [2.0, 3.0]) + + rng = np.random.default_rng(42) + sample = params.sample(rng=rng) + assert len(sample) == 3 + + +def test_save_and_load_round_trip(tmp_path): + """Saving to YAML and loading back recovers modified values, nominals, and bounds.""" + param1 = parameter.Parameter("p1", 1.0, 0.0, 2.0) + param2 = parameter.Parameter("p2", [3.0, 4.0], [1.0, 2.0], [5.0, 6.0]) + params = parameter.ParameterDict({"p1": param1, "p2": param2}) + params.update_from_vector(np.array([0.7, 3.5, 4.5])) + + path = tmp_path / "params.yaml" + params.save_to_disk(path) + + loaded = parameter.ParameterDict.load_from_disk(path) + np.testing.assert_array_equal(loaded.as_vector(), [0.7, 3.5, 4.5]) + np.testing.assert_array_equal(loaded["p1"].nominal, [1.0]) + np.testing.assert_array_equal(loaded["p2"].min_value, [1.0, 2.0]) + + +def test_randomize_stays_in_bounds(): + """Randomized parameter values always stay within their declared bounds.""" + param1 = parameter.Parameter("a", 5.0, 2.0, 8.0) + param2 = parameter.Parameter("b", [1.0, 2.0], [0.0, 0.0], [3.0, 3.0]) + params = parameter.ParameterDict({"a": param1, "b": param2}) + + rng = np.random.default_rng(0) + for _ in range(10): + params.randomize(rng=rng) + lower, upper = params.get_bounds() + vec = params.as_vector() + assert np.all(vec >= lower) + assert np.all(vec <= upper) + + +def test_frozen_param_excluded(): + """Freezing a parameter hides it from the optimizer: excluded from vector ops.""" + p1 = parameter.Parameter("free", 1.0, 0.0, 2.0) + p2 = parameter.Parameter("frozen", 5.0, 3.0, 7.0, frozen=True) + params = parameter.ParameterDict({"free": p1, "frozen": p2}) + + assert params.size == 1 + np.testing.assert_array_equal(params.as_vector(), [1.0]) + + params.update_from_vector(np.array([1.5])) + np.testing.assert_array_equal(params["free"].value, [1.5]) + # Frozen param unchanged. + np.testing.assert_array_equal(params["frozen"].value, [5.0]) diff --git a/python/mujoco/sysid/tests/test_signal.py b/python/mujoco/sysid/tests/test_signal.py new file mode 100644 index 00000000..f15d90a3 --- /dev/null +++ b/python/mujoco/sysid/tests/test_signal.py @@ -0,0 +1,377 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for signal_modifier and SignalTransform.""" + +import numpy as np +import pytest + +from mujoco.sysid._src import parameter, signal_modifier, timeseries +from mujoco.sysid._src.parameter import Parameter, ParameterDict +from mujoco.sysid._src.signal_transform import SignalTransform + + +# =========================================================================== +# Helpers +# =========================================================================== + + +def _make_pdict(*params: Parameter) -> ParameterDict: + pdict = ParameterDict() + for p in params: + pdict.add(p) + return pdict + + +def _make_arm_sensor_ts(arm_model): + """Create a synthetic TimeSeries with signal_mapping matching arm sensors.""" + n_steps = 10 + n_sensors = arm_model.nsensordata + times = np.linspace(0, 1, n_steps) + data = np.random.default_rng(42).standard_normal((n_steps, n_sensors)) + + mapping = {} + for i in range(arm_model.nsensor): + name = arm_model.sensor(i).name + adr = arm_model.sensor_adr[i] + dim = arm_model.sensor_dim[i] + mapping[name] = (timeseries.SignalType.MjSensor, np.arange(adr, adr + dim)) + + return timeseries.TimeSeries(times, data, signal_mapping=mapping) + + +def _make_resample_ts(n_steps, n_cols, seed=0): + """Deterministic TimeSeries for resample tests.""" + rng = np.random.default_rng(seed) + times = np.linspace(0, 1, n_steps) + data = rng.standard_normal((n_steps, n_cols)) + mapping = { + f"s{i}": (timeseries.SignalType.MjSensor, np.array([i])) for i in range(n_cols) + } + return timeseries.TimeSeries(times, data, signal_mapping=mapping) + + +def _make_transform_sensor_ts(n_steps=50, n_sensors=15, seed=0): + """Deterministic TimeSeries with named MjSensor columns for transforms.""" + rng = np.random.default_rng(seed) + times = np.linspace(0, 1, n_steps) + data = rng.standard_normal((n_steps, n_sensors)) + mapping = { + f"joint{i + 1}_pos": (timeseries.SignalType.MjSensor, np.array([i])) + for i in range(5) + } + mapping.update( + { + f"joint{i - 4}_vel": (timeseries.SignalType.MjSensor, np.array([i])) + for i in range(5, 10) + } + ) + mapping.update( + { + f"joint{i - 9}_torque": (timeseries.SignalType.MjSensor, np.array([i])) + for i in range(10, 15) + } + ) + return timeseries.TimeSeries(times, data, signal_mapping=mapping) + + +def _run_both(ts, times, default_delay, sensor_delays, predicted_data): + """Run grouped and column-wise implementations, return both.""" + delays = signal_modifier._build_per_column_delays( + ts, default_delay, sensor_delays, predicted_data + ) + reference = signal_modifier._apply_resample_and_delay_columnwise(ts, times, delays) + result = signal_modifier.apply_resample_and_delay( + ts, + times, + default_delay, + sensor_delays=sensor_delays, + predicted_data=predicted_data, + ) + return result.data, reference + + +def _run_gains_biases_both(transform, ts, target_label, params): + """Run new and reference implementations, return both results.""" + new_result = transform._apply_gains_biases(ts, target_label, params) + ref_result = transform._apply_gains_biases_reference(ts, target_label, params) + return new_result, ref_result + + +# =========================================================================== +# signal_modifier: get_sensor_indices +# =========================================================================== + + +def test_get_sensor_indices(arm_model): + """Sensor name lookup returns the right data column indices for one or many sensors.""" + indices = signal_modifier.get_sensor_indices(arm_model, "joint1_pos") + assert isinstance(indices, list) + assert len(indices) == 1 + + indices = signal_modifier.get_sensor_indices(arm_model, ["joint1_pos", "joint2_pos"]) + assert len(indices) == 2 + + +# =========================================================================== +# signal_modifier: apply_gain / apply_bias +# =========================================================================== + + +def test_apply_gain(arm_model): + """Gain scaling affects only the named sensor's columns, leaving others untouched.""" + ts = _make_arm_sensor_ts(arm_model) + gain = parameter.Parameter("gain", 2.0, 0.5, 3.0) + + result = signal_modifier.apply_gain(ts, "joint1_torque", gain) + + idx = ts.get_indices("joint1_torque")[1] + np.testing.assert_allclose(result.data[:, idx], ts.data[:, idx] * 2.0) + other = [i for i in range(ts.data.shape[1]) if i not in idx] + np.testing.assert_array_equal(result.data[:, other], ts.data[:, other]) + + +def test_apply_bias(arm_model): + """Bias offset affects only the named sensor's columns, leaving others untouched.""" + ts = _make_arm_sensor_ts(arm_model) + bias = parameter.Parameter("bias", 0.5, -1.0, 1.0) + + result = signal_modifier.apply_bias(ts, "joint1_pos", bias) + + idx = ts.get_indices("joint1_pos")[1] + np.testing.assert_allclose(result.data[:, idx], ts.data[:, idx] + 0.5) + other = [i for i in range(ts.data.shape[1]) if i not in idx] + np.testing.assert_array_equal(result.data[:, other], ts.data[:, other]) + + +# =========================================================================== +# signal_modifier: apply_delayed_ts_window +# =========================================================================== + + +def test_apply_delayed_ts_window(arm_model): + """Time-windowing crops timestamps to the overlapping region between two series.""" + ts = _make_arm_sensor_ts(arm_model) + ts_delayed = _make_arm_sensor_ts(arm_model) + + result = signal_modifier.apply_delayed_ts_window( + ts, ts_delayed, min_delay=0.0, max_delay=0.0 + ) + assert result.times[0] >= ts_delayed.times[0] + assert result.times[-1] <= ts_delayed.times[-1] + + +# =========================================================================== +# signal_modifier: weighted_diff / normalize_residual +# =========================================================================== + + +def test_weighted_diff_basic(): + """Without weights, the residual is simply measured minus predicted.""" + predicted = np.array([[1.0, 2.0], [3.0, 4.0]]) + measured = np.array([[1.1, 2.2], [3.3, 4.4]]) + result = signal_modifier.weighted_diff(predicted, measured) + np.testing.assert_allclose(result, measured - predicted) + + +def test_weighted_diff_with_weights(arm_model): + """Sensor weights let you emphasize or de-emphasize specific channels in the residual.""" + n = arm_model.nsensordata + predicted = np.ones((5, n)) + measured = np.ones((5, n)) * 2.0 + weights = {"joint1_pos": 0.5} + result = signal_modifier.weighted_diff(predicted, measured, arm_model, weights) + idx = signal_modifier.get_sensor_indices(arm_model, "joint1_pos") + np.testing.assert_allclose(result[:, idx], 0.5) + other = [i for i in range(n) if i not in idx] + np.testing.assert_allclose(result[:, other], 1.0) + + +def test_normalize_residual(): + """Normalization makes residuals comparable across sensors with different scales.""" + residual = np.array([[2.0, 4.0], [6.0, 8.0]]) + measured = np.array([[1.0, 2.0], [3.0, 4.0]]) + result = signal_modifier.normalize_residual(residual, measured) + norm = np.linalg.norm(measured, axis=0) / np.sqrt(2) + np.testing.assert_allclose(result, residual / norm) + + +# =========================================================================== +# signal_modifier: resample_and_delay grouped vs columnwise equivalence +# =========================================================================== + + +def test_resample_delay_mixed_delays(): + """Optimized grouped resampling gives identical results to naive per-column resampling.""" + ts = _make_resample_ts(200, 8, seed=42) + out_times = np.linspace(0.05, 0.95, 150) + sensor_delays = { + "s0": 0.01, + "s1": 0.01, + "s2": 0.01, + "s3": 0.03, + "s4": 0.03, + } + result, reference = _run_both(ts, out_times, 0.0, sensor_delays, True) + np.testing.assert_array_equal(result, reference) + + +# =========================================================================== +# SignalTransform: pattern matching +# =========================================================================== + + +class TestPatternMatching: + + def test_basic_glob(self): + """Wildcard patterns select the right sensors (e.g. '*_pos' matches positions only).""" + transform = SignalTransform() + delay_param = Parameter("delay", [0.01], [0.0], [0.05]) + transform.delay("*_pos", delay_param) + pdict = _make_pdict(delay_param) + + resolved = transform._resolve_delays( + ["joint1_pos", "joint2_pos", "joint1_vel"], pdict + ) + assert "joint1_pos" in resolved + assert "joint2_pos" in resolved + assert "joint1_vel" not in resolved + assert resolved["joint1_pos"] == pytest.approx(0.01) + + def test_last_match_wins(self): + """When patterns overlap, the last one registered takes priority.""" + transform = SignalTransform() + general_delay = Parameter("delay_general", [0.01], [0.0], [0.05]) + specific_delay = Parameter("delay_specific", [0.05], [0.0], [0.1]) + transform.delay("*_torque", general_delay) + transform.delay("joint5_torque", specific_delay) + pdict = _make_pdict(general_delay, specific_delay) + + resolved = transform._resolve_delays(["joint1_torque", "joint5_torque"], pdict) + assert resolved["joint1_torque"] == pytest.approx(0.01) + assert resolved["joint5_torque"] == pytest.approx(0.05) + + def test_no_match(self): + """Patterns that don't match any sensor names produce no delay entries.""" + transform = SignalTransform() + delay_param = Parameter("delay", [0.01], [0.0], [0.05]) + transform.delay("*_pos", delay_param) + pdict = _make_pdict(delay_param) + + resolved = transform._resolve_delays(["joint1_vel", "joint2_vel"], pdict) + assert len(resolved) == 0 + + +# =========================================================================== +# SignalTransform: delay bounds +# =========================================================================== + + +class TestDelayBounds: + + def test_single_param(self): + """The min/max delay window is derived from a parameter's declared bounds.""" + transform = SignalTransform() + delay_param = Parameter("delay", [0.01], [-0.02], [0.05]) + transform.delay("*_pos", delay_param) + + min_d, max_d = transform._compute_delay_bounds() + assert min_d == pytest.approx(-0.02) + assert max_d == pytest.approx(0.05) + + def test_dedup_by_name(self): + """Reusing one delay param across patterns doesn't double-count its bounds.""" + transform = SignalTransform() + delay_param = Parameter("delay", [0.01], [-0.01], [0.05]) + transform.delay("*_pos", delay_param) + transform.delay("*_vel", delay_param) + + min_d, max_d = transform._compute_delay_bounds() + assert min_d == pytest.approx(-0.01) + assert max_d == pytest.approx(0.05) + + +# =========================================================================== +# SignalTransform: edge cases +# =========================================================================== + + +class TestEdgeCases: + + def test_enable_sensors_stores_copy(self): + """The sensor list is defensively copied so callers can't mutate it after the fact.""" + transform = SignalTransform() + sensors = ["a", "b"] + transform.enable_sensors(sensors) + sensors.append("c") + assert transform._enabled_sensors == ["a", "b"] + + def test_invalid_target(self): + """Typos in the target argument ('predicted'/'measured'/'both') are caught early.""" + transform = SignalTransform() + param = Parameter("gain", [1.0], [0.5], [2.0]) + with pytest.raises(ValueError, match="target must be"): + transform.gain("*", param, target="invalid") + + bias_param = Parameter("bias", [0.0], [-1.0], [1.0]) + with pytest.raises(ValueError, match="target must be"): + transform.bias("*", bias_param, target="invalid") + + +# =========================================================================== +# SignalTransform: _apply_gains_biases equivalence +# =========================================================================== + + +class TestApplyGainsBiasesEquivalence: + + def test_gains_and_biases_mixed(self): + """Applying gains and biases together produces the same result as the reference path.""" + ts = _make_transform_sensor_ts() + gain = Parameter("torque_scale", [1.5], [0.5], [3.0]) + bias = Parameter("torque_bias", [0.3], [-1.0], [1.0]) + pdict = _make_pdict(gain, bias) + + transform = SignalTransform() + transform.gain("*_torque", gain, target="both") + transform.bias("*_torque", bias, target="both") + + new, ref = _run_gains_biases_both(transform, ts, "predicted", pdict) + np.testing.assert_array_equal(new.data, ref.data) + + def test_target_filtering(self): + """A gain meant for measured data doesn't accidentally affect the predicted side.""" + ts = _make_transform_sensor_ts() + gain = Parameter("gain", [2.0], [0.5], [3.0]) + pdict = _make_pdict(gain) + + transform = SignalTransform() + transform.gain("*_torque", gain, target="measured") + + new, ref = _run_gains_biases_both(transform, ts, "predicted", pdict) + np.testing.assert_array_equal(new.data, ref.data) + np.testing.assert_array_equal(new.data, ts.data) + + def test_original_ts_not_mutated(self): + """Signal transforms produce new data without mutating the input TimeSeries.""" + ts = _make_transform_sensor_ts() + original_data = ts.data.copy() + gain = Parameter("gain", [2.0], [0.5], [3.0]) + pdict = _make_pdict(gain) + + transform = SignalTransform() + transform.gain("*_torque", gain, target="predicted") + + transform._apply_gains_biases(ts, "predicted", pdict) + np.testing.assert_array_equal(ts.data, original_data) diff --git a/python/mujoco/sysid/tests/test_timeseries.py b/python/mujoco/sysid/tests/test_timeseries.py new file mode 100644 index 00000000..92c55f09 --- /dev/null +++ b/python/mujoco/sysid/tests/test_timeseries.py @@ -0,0 +1,332 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the TimeSeries class and factory methods.""" + +import mujoco +import numpy as np +import pytest + +from mujoco.sysid import SignalType, TimeSeries +from mujoco.sysid._src import timeseries + + +# --------------------------------------------------------------------------- +# Local fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def scalar_ts(): + """y = x^2""" + times = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) + data = np.array([0.0, 1.0, 4.0, 9.0, 16.0]) + return timeseries.TimeSeries(times=times, data=data) + + +@pytest.fixture +def multi_ts(): + """y = [x^2, 2*x^2]""" + times = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) + data = np.array( + [ + [0.0, 0.0], + [1.0, 2.0], + [4.0, 8.0], + [9.0, 18.0], + [16.0, 32.0], + ] + ) + return timeseries.TimeSeries(times=times, data=data) + + +# --------------------------------------------------------------------------- +# Core TimeSeries tests +# --------------------------------------------------------------------------- + + +def test_basics(scalar_ts, multi_ts): + """Basic properties: length, times array, and data array are all accessible.""" + assert len(scalar_ts) == 5 + assert len(multi_ts) == 5 + np.testing.assert_array_equal(scalar_ts.times, [0, 1, 2, 3, 4]) + np.testing.assert_array_equal(scalar_ts.data, [0, 1, 4, 9, 16]) + + +@pytest.mark.parametrize( + "times, data, match", + [ + (np.array([]), np.array([]), "Empty"), + (np.array([[0.0], [1.0]]), np.array([0.0, 1.0]), "1D"), + (np.array([0.0, 1.0]), np.array([0.0, 1.0, 2.0]), "Length"), + (np.array([0.0, 2.0, 1.0]), np.array([0.0, 1.0, 2.0]), "strictly increasing"), + ], +) +def test_validation(times, data, match): + """Bad inputs (empty, non-1D times, length mismatch, non-monotonic) are rejected.""" + with pytest.raises(ValueError, match=match): + timeseries.TimeSeries(times=times, data=data) + + +def test_zero_column_data(): + """Zero-column data is valid (state-based models with no sensors).""" + times = np.array([0.0, 1.0, 2.0]) + data = np.empty((3, 0)) + ts = timeseries.TimeSeries(times=times, data=data) + assert len(ts) == 3 + assert ts.data.shape == (3, 0) + + +def test_save_and_load(scalar_ts, multi_ts, tmp_path): + """Saving to .npz and loading back recovers identical times and data.""" + path = tmp_path / "test.npz" + scalar_ts.save_to_disk(path) + loaded = timeseries.TimeSeries.load_from_disk(path) + np.testing.assert_array_equal(loaded.times, scalar_ts.times) + np.testing.assert_array_equal(loaded.data, scalar_ts.data) + + path2 = tmp_path / "test_multi.npz" + multi_ts.save_to_disk(path2) + loaded2 = timeseries.TimeSeries.load_from_disk(path2) + np.testing.assert_array_equal(loaded2.times, multi_ts.times) + np.testing.assert_array_equal(loaded2.data, multi_ts.data) + + +def test_save_and_load_with_signal_mapping(tmp_path): + """Save/load also preserves the signal_mapping (sensor name -> column index map).""" + times = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) + data = np.array( + [ + [0.0, 0.0], + [1.0, 2.0], + [4.0, 8.0], + [9.0, 18.0], + [16.0, 32.0], + ] + ) + signal_mapping = { + "signal1": (timeseries.SignalType.MjSensor, np.array([0])), + "signal2": (timeseries.SignalType.MjSensor, np.array([1])), + } + ts = timeseries.TimeSeries(times=times, data=data, signal_mapping=signal_mapping) + + path = tmp_path / "test_signal_mapping.npz" + ts.save_to_disk(path) + loaded = timeseries.TimeSeries.load_from_disk(path) + + np.testing.assert_array_equal(loaded.times, times) + np.testing.assert_array_equal(loaded.data, data) + assert loaded.signal_mapping is not None + assert loaded.signal_mapping.keys() == signal_mapping.keys() + for key in signal_mapping: + val_type, val_indices = signal_mapping[key] + loaded_type, loaded_indices = loaded.signal_mapping[key] + assert loaded_type == val_type + np.testing.assert_array_equal(loaded_indices, val_indices) + + +@pytest.mark.parametrize( + "method, expected", + [ + ("linear", 6.5), + ("cubic", 6.25), + ("quadratic", 6.25), + ("zero_order_hold", 4.0), + ("zoh", 4.0), + ], +) +def test_interpolate_scalar(scalar_ts, method, expected): + """Each interpolation method (linear, cubic, ZOH, etc.) gives the expected midpoint value.""" + result = scalar_ts.interpolate(2.5, method=method) + assert result[0] == pytest.approx(expected, abs=1e-5) + + +def test_interpolate_array(scalar_ts, multi_ts): + """Interpolating at multiple times simultaneously works for scalar and multi-column data.""" + t_values = np.array([0.5, 1.5, 2.5, 3.5]) + expected = np.array([0.5, 2.5, 6.5, 12.5]) + result = scalar_ts.interpolate(t_values, method="linear") + np.testing.assert_allclose(result, expected, rtol=1e-5) + + expected_multi = np.array( + [ + [0.5, 1.0], + [2.5, 5.0], + [6.5, 13.0], + [12.5, 25.0], + ] + ) + result_multi = multi_ts.interpolate(t_values, method="linear") + np.testing.assert_allclose(result_multi, expected_multi, rtol=1e-5) + + +def test_resample_with_new_times(scalar_ts): + """Resampling onto a finer time grid via explicit new_times gives correct values.""" + new_times = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]) + expected = np.array([0.0, 0.5, 1.0, 2.5, 4.0, 6.5, 9.0, 12.5, 16.0]) + resampled = scalar_ts.resample(new_times=new_times, method="linear") + np.testing.assert_array_equal(resampled.times, new_times) + np.testing.assert_allclose(resampled.data, expected, rtol=1e-5) + + with pytest.raises(ValueError): + scalar_ts.resample(new_times=np.array([0.0, 2.0, 1.0])) + + with pytest.raises(ValueError): + scalar_ts.resample(new_times=np.array([[0.0], [1.0]])) + + +def test_resample_with_target_dt(scalar_ts): + """Resampling by specifying a target timestep generates the right uniform grid.""" + resampled = scalar_ts.resample(target_dt=0.5, method="linear") + expected_times = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0]) + expected_data = np.array([0.0, 0.5, 1.0, 2.5, 4.0, 6.5, 9.0, 12.5, 16.0]) + np.testing.assert_allclose(resampled.times, expected_times, rtol=1e-5) + np.testing.assert_allclose(resampled.data, expected_data, rtol=1e-5) + + with pytest.raises(ValueError): + scalar_ts.resample(target_dt=-0.5) + + with pytest.raises(ValueError): + scalar_ts.resample() + + +# --------------------------------------------------------------------------- +# TimeSeries factory method tests +# --------------------------------------------------------------------------- + + +def test_from_model_controls_auto_resolution(): + """Without explicit names, all model actuators are auto-discovered and mapped.""" + xml = """ + + + + + + + + + + + + + """ + model = mujoco.MjModel.from_xml_string(xml) + times = np.linspace(0, 1, 100) + data = np.random.randn(100, 2) + + ts = TimeSeries.from_control_names(times, data, model) + assert ts.signal_mapping is not None + assert "m1_ctrl" in ts.signal_mapping + assert "m2_ctrl" in ts.signal_mapping + assert ts.signal_mapping["m1_ctrl"][0] == SignalType.MjCtrl + assert ts.signal_mapping["m2_ctrl"][0] == SignalType.MjCtrl + + +def test_from_model_controls_explicit_names(): + """Explicit actuator names are resolved; invalid or wrong-type names are rejected.""" + xml = """ + + + + + + + + + + + + """ + model = mujoco.MjModel.from_xml_string(xml) + times = np.linspace(0, 1, 10) + data = np.zeros((10, 1)) + + ts = TimeSeries.from_control_names(times, data, model, names=["m1"]) + assert ts.signal_mapping is not None + assert "m1_ctrl" in ts.signal_mapping + + with pytest.raises(ValueError, match="Could not resolve signal"): + TimeSeries.from_control_names(times, data, model, names=["invalid"]) + + with pytest.raises(ValueError, match="not allowed"): + TimeSeries.from_control_names( + times, data, model, names=[("m1", SignalType.MjSensor)] + ) + + +def test_from_model_auto_resolution_sensors(): + """Without explicit names, all model sensors are auto-discovered with correct dimensions.""" + xml = """ + + + + + + + + + + + + + """ + model = mujoco.MjModel.from_xml_string(xml) + times = np.linspace(0, 1, 10) + data = np.zeros((10, 6)) + + ts = TimeSeries.from_names(times, data, model) + assert ts.signal_mapping is not None + assert "acc1" in ts.signal_mapping + assert "gyro1" in ts.signal_mapping + assert ts.signal_mapping["acc1"][0] == SignalType.MjSensor + assert ts.signal_mapping["gyro1"][0] == SignalType.MjSensor + + +def test_from_model_state_resolution(): + """State signals (qpos, qvel) can be mapped by passing (name, SignalType) tuples.""" + xml = """ + + + + + + + + + + """ + model = mujoco.MjModel.from_xml_string(xml) + times = np.linspace(0, 1, 10) + data = np.zeros((10, 2)) + + names = [("j1", SignalType.MjStateQPos), ("j2", SignalType.MjStateQPos)] + ts = TimeSeries.from_names(times, data, model, names=names) + assert ts.signal_mapping is not None + assert "j1_qpos" in ts.signal_mapping + assert "j2_qpos" in ts.signal_mapping + + +def test_from_custom(): + """Custom signal definitions (name strings and dimension tuples) are mapped correctly.""" + times = np.linspace(0, 1, 10) + data = np.zeros((10, 3)) + signals = ["a", ("b", 2, SignalType.CustomObs)] + + ts = TimeSeries.from_custom_map(times, data, signals) + assert ts.signal_mapping is not None + assert "a" in ts.signal_mapping + assert "b" in ts.signal_mapping + assert ts.signal_mapping["a"][1].size == 1 + assert ts.signal_mapping["b"][1].size == 2 diff --git a/python/mujoco/sysid/tests/test_trajectory.py b/python/mujoco/sysid/tests/test_trajectory.py new file mode 100644 index 00000000..25016f91 --- /dev/null +++ b/python/mujoco/sysid/tests/test_trajectory.py @@ -0,0 +1,146 @@ +# Copyright 2025 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for the SystemTrajectory class.""" + +from unittest import mock + +import mujoco +import numpy as np +import pytest + +from mujoco.sysid._src import timeseries +from mujoco.sysid._src.trajectory import SystemTrajectory, create_initial_state + + +@pytest.fixture +def mock_model(): + model = mock.Mock(spec=mujoco.MjModel) + model.nsensordata = 2 + model.nu = 1 + model.nq = 1 + model.nv = 1 + return model + + +@pytest.fixture +def sample_trajectory(mock_model): + with mock.patch.object(SystemTrajectory, "check_compatible", return_value=None): + times = np.array([0.0, 1.0, 2.0]) + control_mapping = {"ctrl1": (timeseries.SignalType.MjCtrl, np.array([0]))} + sensordata_mapping = { + "sensor1": (timeseries.SignalType.MjSensor, np.array([0])), + "sensor2": (timeseries.SignalType.MjSensor, np.array([1])), + } + state_mapping = {"qpos1": (timeseries.SignalType.MjStateQPos, np.array([0]))} + + control = timeseries.TimeSeries( + times=times, + data=np.array([[1.0], [2.0], [3.0]]), + signal_mapping=control_mapping, + ) + sensordata = timeseries.TimeSeries( + times=times, + data=np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), + signal_mapping=sensordata_mapping, + ) + state = timeseries.TimeSeries( + times=times, + data=np.array([[0.01], [0.02], [0.03]]), + signal_mapping=state_mapping, + ) + + traj = SystemTrajectory( + model=mock_model, + control=control, + sensordata=sensordata, + initial_state=np.array([0.0]), + state=state, + ) + yield traj, control_mapping, sensordata_mapping, state_mapping + + +def test_save_and_load_with_signal_mapping(sample_trajectory, mock_model, tmp_path): + """Saving and loading a trajectory preserves all signal mappings (control, sensor, state).""" + traj, control_mapping, sensordata_mapping, state_mapping = sample_trajectory + + path = tmp_path / "test_traj.npz" + traj.save_to_disk(path) + + with mock.patch.object(SystemTrajectory, "check_compatible", return_value=None): + loaded = SystemTrajectory.load_from_disk(path, mock_model) + + ctrl_map = loaded.control.signal_mapping + assert ctrl_map is not None + assert ctrl_map.keys() == control_mapping.keys() + for key in control_mapping: + assert ctrl_map[key][0] == control_mapping[key][0] + np.testing.assert_array_equal(ctrl_map[key][1], control_mapping[key][1]) + + sensor_map = loaded.sensordata.signal_mapping + assert sensor_map is not None + assert sensor_map.keys() == sensordata_mapping.keys() + for key in sensordata_mapping: + assert sensor_map[key][0] == sensordata_mapping[key][0] + np.testing.assert_array_equal(sensor_map[key][1], sensordata_mapping[key][1]) + + assert loaded.state is not None + state_map = loaded.state.signal_mapping + assert state_map is not None + assert state_map.keys() == state_mapping.keys() + for key in state_mapping: + assert state_map[key][0] == state_mapping[key][0] + np.testing.assert_array_equal(state_map[key][1], state_mapping[key][1]) + + +def test_create_initial_state(box_model): + """The initial MuJoCo state (qpos, qvel, act) is packed into a flat vector for rollout.""" + qpos = np.zeros(box_model.nq) + qvel = np.zeros(box_model.nv) + state = create_initial_state(box_model, qpos, qvel) + expected_size = mujoco.mj_stateSize(box_model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + assert state.shape == (expected_size,) + + +def test_create_initial_state_wrong_qpos(box_model): + """Wrong-sized qpos is caught early rather than causing a silent rollout bug.""" + with pytest.raises(ValueError, match="qpos"): + create_initial_state(box_model, np.zeros(999)) + + +def test_split(sample_trajectory): + """A long trajectory can be split into smaller chunks for batched optimization.""" + traj, *_ = sample_trajectory + chunks = traj.split(chunk_size=1) + assert len(chunks) == 3 + assert len(chunks[0].sensordata) == 1 + + +def test_check_compatible_sensor_mismatch(box_model): + """Mismatched sensor dimensions between data and model are caught before rollout.""" + times = np.array([0.0, 0.01, 0.02]) + sensordata = timeseries.TimeSeries(times, np.ones((3, 5))) + control = timeseries.TimeSeries(times, np.ones((3, box_model.nu))) + initial_state = np.zeros( + mujoco.mj_stateSize(box_model, mujoco.mjtState.mjSTATE_FULLPHYSICS) + ) + traj = SystemTrajectory( + model=box_model, + control=control, + sensordata=sensordata, + initial_state=initial_state, + state=None, + ) + with pytest.raises(ValueError, match="Sensor data dimension"): + traj.check_compatible() diff --git a/python/pyproject.toml b/python/pyproject.toml index f8ca2c2c..ebcfd100 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -62,6 +62,18 @@ mujoco = [ ] [project.optional-dependencies] +sysid = [ + "absl-py", + "colorama", + "imageio[ffmpeg]", + "jinja2", + "matplotlib", + "plotly", + "pyyaml", + "scipy", + "tabulate", + "typing_extensions", +] usd = [ "usd-core", "pillow" From a1f2b875483b45369f4aa45c5dad06f5e3cbd257 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Mon, 9 Feb 2026 09:58:34 -0800 Subject: [PATCH 2/2] Ignore sysid tests. --- .github/workflows/build_steps.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build_steps.sh b/.github/workflows/build_steps.sh index 6a6889bc..97c33628 100755 --- a/.github/workflows/build_steps.sh +++ b/.github/workflows/build_steps.sh @@ -215,7 +215,9 @@ install_python_bindings() { test_python_bindings() { echo "Testing Python bindings..." source ${TMPDIR}/venv/bin/activate && - pytest -v --pyargs mujoco + # TODO(kevinzakka): Add sysid tests to CI once sysid dependencies are installed in the CI environment. + MUJOCO_PATH=$(python -c "import mujoco; print(mujoco.__path__[0])") && + pytest -v --pyargs mujoco --ignore="${MUJOCO_PATH}/sysid/tests" }