jaxsim 0.2.dev191__tar.gz → 0.2.dev366__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (160) hide show
  1. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.github/workflows/ci_cd.yml +13 -2
  2. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/PKG-INFO +4 -6
  3. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/environment.yml +1 -2
  4. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/pyproject.toml +1 -1
  5. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/setup.cfg +4 -5
  6. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/__init__.py +3 -4
  7. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/_version.py +2 -2
  8. jaxsim-0.2.dev366/src/jaxsim/api/__init__.py +3 -0
  9. jaxsim-0.2.dev366/src/jaxsim/api/com.py +240 -0
  10. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/api/common.py +13 -2
  11. jaxsim-0.2.dev366/src/jaxsim/api/contact.py +271 -0
  12. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/api/data.py +112 -71
  13. jaxsim-0.2.dev366/src/jaxsim/api/joint.py +189 -0
  14. jaxsim-0.2.dev366/src/jaxsim/api/kin_dyn_parameters.py +777 -0
  15. jaxsim-0.2.dev366/src/jaxsim/api/link.py +337 -0
  16. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/api/model.py +542 -269
  17. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/api/ode.py +86 -74
  18. jaxsim-0.2.dev366/src/jaxsim/api/ode_data.py +694 -0
  19. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/api/references.py +12 -11
  20. jaxsim-0.2.dev366/src/jaxsim/integrators/__init__.py +2 -0
  21. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/integrators/common.py +110 -24
  22. jaxsim-0.2.dev366/src/jaxsim/integrators/fixed_step.py +102 -0
  23. jaxsim-0.2.dev366/src/jaxsim/integrators/variable_step.py +610 -0
  24. jaxsim-0.2.dev366/src/jaxsim/math/__init__.py +11 -0
  25. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/math/adjoint.py +24 -2
  26. jaxsim-0.2.dev366/src/jaxsim/math/joint_model.py +335 -0
  27. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/math/quaternion.py +44 -3
  28. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/math/rotation.py +4 -4
  29. jaxsim-0.2.dev366/src/jaxsim/math/transform.py +93 -0
  30. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/descriptions/link.py +2 -2
  31. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/rod/utils.py +7 -8
  32. jaxsim-0.2.dev366/src/jaxsim/rbda/__init__.py +7 -0
  33. jaxsim-0.2.dev366/src/jaxsim/rbda/aba.py +295 -0
  34. jaxsim-0.2.dev366/src/jaxsim/rbda/collidable_points.py +142 -0
  35. {jaxsim-0.2.dev191/src/jaxsim/physics/algos → jaxsim-0.2.dev366/src/jaxsim/rbda}/crba.py +43 -42
  36. jaxsim-0.2.dev366/src/jaxsim/rbda/forward_kinematics.py +113 -0
  37. jaxsim-0.2.dev366/src/jaxsim/rbda/jacobian.py +201 -0
  38. jaxsim-0.2.dev366/src/jaxsim/rbda/rnea.py +237 -0
  39. jaxsim-0.2.dev366/src/jaxsim/rbda/soft_contacts.py +296 -0
  40. jaxsim-0.2.dev366/src/jaxsim/rbda/utils.py +152 -0
  41. jaxsim-0.2.dev366/src/jaxsim/terrain/__init__.py +2 -0
  42. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/utils/__init__.py +1 -4
  43. jaxsim-0.2.dev366/src/jaxsim/utils/hashless.py +18 -0
  44. jaxsim-0.2.dev366/src/jaxsim/utils/jaxsim_dataclass.py +360 -0
  45. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim.egg-info/PKG-INFO +4 -6
  46. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim.egg-info/SOURCES.txt +28 -42
  47. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim.egg-info/requires.txt +3 -5
  48. jaxsim-0.2.dev366/tests/conftest.py +364 -0
  49. jaxsim-0.2.dev366/tests/test_api_com.py +63 -0
  50. jaxsim-0.2.dev366/tests/test_api_data.py +145 -0
  51. jaxsim-0.2.dev366/tests/test_api_joint.py +34 -0
  52. jaxsim-0.2.dev366/tests/test_api_link.py +159 -0
  53. jaxsim-0.2.dev366/tests/test_api_model.py +319 -0
  54. jaxsim-0.2.dev366/tests/test_automatic_differentiation.py +430 -0
  55. jaxsim-0.2.dev366/tests/test_pytree.py +60 -0
  56. jaxsim-0.2.dev366/tests/test_simulations.py +92 -0
  57. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/tests/utils_idyntree.py +177 -15
  58. jaxsim-0.2.dev191/src/jaxsim/api/__init__.py +0 -1
  59. jaxsim-0.2.dev191/src/jaxsim/api/contact.py +0 -194
  60. jaxsim-0.2.dev191/src/jaxsim/api/joint.py +0 -148
  61. jaxsim-0.2.dev191/src/jaxsim/api/link.py +0 -262
  62. jaxsim-0.2.dev191/src/jaxsim/high_level/__init__.py +0 -2
  63. jaxsim-0.2.dev191/src/jaxsim/high_level/common.py +0 -11
  64. jaxsim-0.2.dev191/src/jaxsim/high_level/joint.py +0 -148
  65. jaxsim-0.2.dev191/src/jaxsim/high_level/link.py +0 -259
  66. jaxsim-0.2.dev191/src/jaxsim/high_level/model.py +0 -1686
  67. jaxsim-0.2.dev191/src/jaxsim/integrators/__init__.py +0 -2
  68. jaxsim-0.2.dev191/src/jaxsim/integrators/fixed_step.py +0 -158
  69. jaxsim-0.2.dev191/src/jaxsim/math/conv.py +0 -114
  70. jaxsim-0.2.dev191/src/jaxsim/math/joint.py +0 -102
  71. jaxsim-0.2.dev191/src/jaxsim/math/plucker.py +0 -100
  72. jaxsim-0.2.dev191/src/jaxsim/physics/__init__.py +0 -12
  73. jaxsim-0.2.dev191/src/jaxsim/physics/algos/__init__.py +0 -0
  74. jaxsim-0.2.dev191/src/jaxsim/physics/algos/aba.py +0 -254
  75. jaxsim-0.2.dev191/src/jaxsim/physics/algos/aba_motors.py +0 -284
  76. jaxsim-0.2.dev191/src/jaxsim/physics/algos/forward_kinematics.py +0 -79
  77. jaxsim-0.2.dev191/src/jaxsim/physics/algos/jacobian.py +0 -98
  78. jaxsim-0.2.dev191/src/jaxsim/physics/algos/rnea.py +0 -180
  79. jaxsim-0.2.dev191/src/jaxsim/physics/algos/rnea_motors.py +0 -196
  80. jaxsim-0.2.dev191/src/jaxsim/physics/algos/soft_contacts.py +0 -523
  81. jaxsim-0.2.dev191/src/jaxsim/physics/algos/utils.py +0 -69
  82. jaxsim-0.2.dev191/src/jaxsim/physics/model/__init__.py +0 -0
  83. jaxsim-0.2.dev191/src/jaxsim/physics/model/ground_contact.py +0 -53
  84. jaxsim-0.2.dev191/src/jaxsim/physics/model/physics_model.py +0 -388
  85. jaxsim-0.2.dev191/src/jaxsim/physics/model/physics_model_state.py +0 -283
  86. jaxsim-0.2.dev191/src/jaxsim/simulation/__init__.py +0 -4
  87. jaxsim-0.2.dev191/src/jaxsim/simulation/integrators.py +0 -393
  88. jaxsim-0.2.dev191/src/jaxsim/simulation/ode.py +0 -290
  89. jaxsim-0.2.dev191/src/jaxsim/simulation/ode_data.py +0 -96
  90. jaxsim-0.2.dev191/src/jaxsim/simulation/ode_integration.py +0 -62
  91. jaxsim-0.2.dev191/src/jaxsim/simulation/simulator.py +0 -543
  92. jaxsim-0.2.dev191/src/jaxsim/simulation/simulator_callbacks.py +0 -79
  93. jaxsim-0.2.dev191/src/jaxsim/simulation/utils.py +0 -15
  94. jaxsim-0.2.dev191/src/jaxsim/sixd/__init__.py +0 -2
  95. jaxsim-0.2.dev191/src/jaxsim/utils/jaxsim_dataclass.py +0 -109
  96. jaxsim-0.2.dev191/src/jaxsim/utils/oop.py +0 -536
  97. jaxsim-0.2.dev191/src/jaxsim/utils/vmappable.py +0 -117
  98. jaxsim-0.2.dev191/tests/__init__.py +0 -0
  99. jaxsim-0.2.dev191/tests/test_ad_physics.py +0 -190
  100. jaxsim-0.2.dev191/tests/test_eom.py +0 -130
  101. jaxsim-0.2.dev191/tests/test_forward_dynamics.py +0 -71
  102. jaxsim-0.2.dev191/tests/test_jax_oop.py +0 -422
  103. jaxsim-0.2.dev191/tests/utils_models.py +0 -56
  104. jaxsim-0.2.dev191/tests/utils_rng.py +0 -96
  105. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.devcontainer/Dockerfile +0 -0
  106. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.devcontainer/devcontainer.json +0 -0
  107. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.github/CODEOWNERS +0 -0
  108. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.github/workflows/read_the_docs.yml +0 -0
  109. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.github/workflows/style.yml +0 -0
  110. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.gitignore +0 -0
  111. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.pre-commit-config.yaml +0 -0
  112. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/.readthedocs.yaml +0 -0
  113. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/CONTRIBUTING.md +0 -0
  114. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/LICENSE +0 -0
  115. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/README.md +0 -0
  116. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/Makefile +0 -0
  117. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/conf.py +0 -0
  118. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/guide/install.rst +0 -0
  119. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/index.rst +0 -0
  120. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/make.bat +0 -0
  121. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/high_level.rst +0 -0
  122. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/math.rst +0 -0
  123. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/parsers.rst +0 -0
  124. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/physics.rst +0 -0
  125. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/simulation.rst +0 -0
  126. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/typing.rst +0 -0
  127. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/docs/modules/utils.rst +0 -0
  128. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/.gitattributes +0 -0
  129. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/.gitignore +0 -0
  130. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/PD_controller.ipynb +0 -0
  131. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/Parallel_computing.ipynb +0 -0
  132. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/README.md +0 -0
  133. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/assets/cartpole.urdf +0 -0
  134. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/pixi.lock +0 -0
  135. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/examples/pixi.toml +0 -0
  136. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/setup.py +0 -0
  137. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/logging.py +0 -0
  138. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/math/cross.py +0 -0
  139. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/math/inertia.py +0 -0
  140. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/math/skew.py +0 -0
  141. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/mujoco/__init__.py +0 -0
  142. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/mujoco/__main__.py +0 -0
  143. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/mujoco/loaders.py +0 -0
  144. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/mujoco/model.py +0 -0
  145. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/mujoco/visualizer.py +0 -0
  146. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/__init__.py +0 -0
  147. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/descriptions/__init__.py +0 -0
  148. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/descriptions/collision.py +0 -0
  149. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/descriptions/joint.py +0 -0
  150. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/descriptions/model.py +0 -0
  151. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/kinematic_graph.py +0 -0
  152. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/rod/__init__.py +0 -0
  153. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/parsers/rod/parser.py +0 -0
  154. {jaxsim-0.2.dev191/src/jaxsim/physics/algos → jaxsim-0.2.dev366/src/jaxsim/terrain}/terrain.py +0 -0
  155. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/typing.py +0 -0
  156. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim/utils/tracing.py +0 -0
  157. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim.egg-info/dependency_links.txt +0 -0
  158. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim.egg-info/not-zip-safe +0 -0
  159. {jaxsim-0.2.dev191 → jaxsim-0.2.dev366}/src/jaxsim.egg-info/top_level.txt +0 -0
  160. {jaxsim-0.2.dev191/src/jaxsim/math → jaxsim-0.2.dev366/tests}/__init__.py +0 -0
@@ -100,11 +100,22 @@ jobs:
100
100
  with:
101
101
  fetch-depth: 0
102
102
 
103
- - name: Install Gazebo Classic
103
+ # - name: Install Gazebo Classic
104
+ # if: contains(matrix.os, 'ubuntu')
105
+ # run: |
106
+ # sudo apt-get update
107
+ # sudo apt-get install gazebo
108
+
109
+ # https://gazebosim.org/docs/harmonic/install_ubuntu
110
+ - name: Install Gazebo Sim
104
111
  if: contains(matrix.os, 'ubuntu')
105
112
  run: |
106
113
  sudo apt-get update
107
- sudo apt-get install gazebo
114
+ sudo apt-get install lsb-release wget gnupg
115
+ sudo wget https://packages.osrfoundation.org/gazebo.gpg -O /usr/share/keyrings/pkgs-osrf-archive-keyring.gpg
116
+ echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/pkgs-osrf-archive-keyring.gpg] http://packages.osrfoundation.org/gazebo/ubuntu-stable $(lsb_release -cs) main" | sudo tee /etc/apt/sources.list.d/gazebo-stable.list > /dev/null
117
+ sudo apt-get update
118
+ sudo apt-get install gz-harmonic
108
119
 
109
120
  - name: Run the Python tests
110
121
  if: contains(matrix.os, 'ubuntu')
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: jaxsim
3
- Version: 0.2.dev191
3
+ Version: 0.2.dev366
4
4
  Summary: A physics engine in reduced coordinates implemented with JAX.
5
5
  Home-page: https://github.com/ami-iit/jaxsim
6
6
  Author: Diego Ferigo
@@ -31,12 +31,12 @@ Requires-Python: >=3.11
31
31
  Description-Content-Type: text/markdown
32
32
  License-File: LICENSE
33
33
  Requires-Dist: coloredlogs
34
- Requires-Dist: jax<0.4.25,>=0.4.13
35
- Requires-Dist: jaxlib<0.4.25,>=0.4.13
34
+ Requires-Dist: jax>=0.4.13
35
+ Requires-Dist: jaxlib>=0.4.13
36
36
  Requires-Dist: jaxlie>=1.3.0
37
37
  Requires-Dist: jax_dataclasses>=1.4.0
38
38
  Requires-Dist: pptree
39
- Requires-Dist: rod
39
+ Requires-Dist: rod>=0.2.0
40
40
  Requires-Dist: typing_extensions; python_version < "3.12"
41
41
  Provides-Extra: style
42
42
  Requires-Dist: black[jupyter]; extra == "style"
@@ -45,7 +45,6 @@ Requires-Dist: pre-commit; extra == "style"
45
45
  Provides-Extra: testing
46
46
  Requires-Dist: idyntree; extra == "testing"
47
47
  Requires-Dist: pytest>=6.0; extra == "testing"
48
- Requires-Dist: pytest-forked; extra == "testing"
49
48
  Requires-Dist: pytest-icdiff; extra == "testing"
50
49
  Requires-Dist: robot-descriptions; extra == "testing"
51
50
  Provides-Extra: viz
@@ -58,7 +57,6 @@ Requires-Dist: isort; extra == "all"
58
57
  Requires-Dist: pre-commit; extra == "all"
59
58
  Requires-Dist: idyntree; extra == "all"
60
59
  Requires-Dist: pytest>=6.0; extra == "all"
61
- Requires-Dist: pytest-forked; extra == "all"
62
60
  Requires-Dist: pytest-icdiff; extra == "all"
63
61
  Requires-Dist: robot-descriptions; extra == "all"
64
62
  Requires-Dist: lxml; extra == "all"
@@ -10,7 +10,7 @@ dependencies:
10
10
  - jaxlie >= 1.3.0
11
11
  - jax-dataclasses >= 1.4.0
12
12
  - pptree
13
- - rod
13
+ - rod >= 0.2.0
14
14
  - typing_extensions # python<3.12
15
15
  # Optional dependencies from setup.cfg
16
16
  # [style]
@@ -19,7 +19,6 @@ dependencies:
19
19
  # [testing]
20
20
  - idyntree
21
21
  - pytest
22
- - pytest-forked
23
22
  - pytest-icdiff
24
23
  - robot_descriptions
25
24
  # [viz]
@@ -19,7 +19,7 @@ multi_line_output = 3
19
19
 
20
20
  [tool.pytest.ini_options]
21
21
  minversion = "6.0"
22
- addopts = "-rsxX -v --strict-markers --forked"
22
+ addopts = "-rsxX -v --strict-markers"
23
23
  testpaths = [
24
24
  "tests",
25
25
  ]
@@ -50,12 +50,12 @@ package_dir =
50
50
  python_requires = >=3.11
51
51
  install_requires =
52
52
  coloredlogs
53
- jax >= 0.4.13,< 0.4.25
54
- jaxlib >= 0.4.13,< 0.4.25
53
+ jax >= 0.4.13
54
+ jaxlib >= 0.4.13
55
55
  jaxlie >= 1.3.0
56
56
  jax_dataclasses >= 1.4.0
57
57
  pptree
58
- rod
58
+ rod >= 0.2.0
59
59
  typing_extensions ; python_version < '3.12'
60
60
 
61
61
  [options.packages.find]
@@ -68,8 +68,7 @@ style =
68
68
  pre-commit
69
69
  testing =
70
70
  idyntree
71
- pytest >= 6.0
72
- pytest-forked
71
+ pytest >=6.0
73
72
  pytest-icdiff
74
73
  robot-descriptions
75
74
  viz =
@@ -61,7 +61,6 @@ del _jnp_options
61
61
  del _np_options
62
62
  del _is_editable
63
63
 
64
- from . import high_level, logging, math, simulation, sixd
65
- from .high_level.common import VelRepr
66
- from .simulation.ode_integration import IntegratorType
67
- from .simulation.simulator import JaxSim
64
+ from . import terrain # isort:skip
65
+ from . import api, integrators, logging, math, rbda
66
+ from .api.common import VelRepr
@@ -12,5 +12,5 @@ __version__: str
12
12
  __version_tuple__: VERSION_TUPLE
13
13
  version_tuple: VERSION_TUPLE
14
14
 
15
- __version__ = version = '0.2.dev191'
16
- __version_tuple__ = version_tuple = (0, 2, 'dev191')
15
+ __version__ = version = '0.2.dev366'
16
+ __version_tuple__ = version_tuple = (0, 2, 'dev366')
@@ -0,0 +1,3 @@
1
+ from . import common # isort:skip
2
+ from . import model, data # isort:skip
3
+ from . import com, contact, joint, kin_dyn_parameters, link, ode, ode_data, references
@@ -0,0 +1,240 @@
1
+ import jax
2
+ import jax.numpy as jnp
3
+ import jaxlie
4
+
5
+ import jaxsim.api as js
6
+ import jaxsim.math
7
+ import jaxsim.typing as jtp
8
+
9
+ from .common import VelRepr
10
+
11
+
12
+ @jax.jit
13
+ def com_position(
14
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
15
+ ) -> jtp.Vector:
16
+ """
17
+ Compute the position of the center of mass of the model.
18
+
19
+ Args:
20
+ model: The model to consider.
21
+ data: The data of the considered model.
22
+
23
+ Returns:
24
+ The position of the center of mass of the model w.r.t. the world frame.
25
+ """
26
+
27
+ m = js.model.total_mass(model=model)
28
+
29
+ W_H_L = js.model.forward_kinematics(model=model, data=data)
30
+ W_H_B = data.base_transform()
31
+ B_H_W = jaxlie.SE3.from_matrix(W_H_B).inverse().as_matrix()
32
+
33
+ def B_p̃_LCoM(i) -> jtp.Vector:
34
+ m = js.link.mass(model=model, link_index=i)
35
+ L_p_LCoM = js.link.com_position(
36
+ model=model, data=data, link_index=i, in_link_frame=True
37
+ )
38
+ return m * B_H_W @ W_H_L[i] @ jnp.hstack([L_p_LCoM, 1])
39
+
40
+ com_links = jax.vmap(B_p̃_LCoM)(jnp.arange(model.number_of_links()))
41
+
42
+ B_p̃_CoM = (1 / m) * com_links.sum(axis=0)
43
+ B_p̃_CoM = B_p̃_CoM.at[3].set(1)
44
+
45
+ return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
46
+
47
+
48
+ @jax.jit
49
+ def com_linear_velocity(
50
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
51
+ ) -> jtp.Vector:
52
+ r"""
53
+ Compute the linear velocity of the center of mass of the model.
54
+
55
+ Args:
56
+ model: The model to consider.
57
+ data: The data of the considered model.
58
+
59
+ Returns:
60
+ The linear velocity of the center of mass of the model in the
61
+ active representation.
62
+
63
+ Note:
64
+ The linear velocity of the center of mass is expressed in the mixed frame
65
+ :math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
66
+ active velocity representation is either inertial-fixed or mixed,
67
+ and :math:`[C] = [B]` if the active velocity representation is body-fixed.
68
+ """
69
+
70
+ # Extract the linear component of the 6D average centroidal velocity.
71
+ # This is expressed in G[B] in body-fixed representation, and in G[W] in
72
+ # inertial-fixed or mixed representation.
73
+ G_vl_WG = average_centroidal_velocity(model=model, data=data)[0:3]
74
+
75
+ return G_vl_WG
76
+
77
+
78
+ @jax.jit
79
+ def centroidal_momentum(
80
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
81
+ ) -> jtp.Vector:
82
+ r"""
83
+ Compute the centroidal momentum of the model.
84
+
85
+ Args:
86
+ model: The model to consider.
87
+ data: The data of the considered model.
88
+
89
+ Returns:
90
+ The centroidal momentum of the model.
91
+
92
+ Note:
93
+ The centroidal momentum is expressed in the mixed frame
94
+ :math:`({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`C = W` if the
95
+ active velocity representation is either inertial-fixed or mixed,
96
+ and :math:`C = B` if the active velocity representation is body-fixed.
97
+ """
98
+
99
+ ν = data.generalized_velocity()
100
+ G_J = centroidal_momentum_jacobian(model=model, data=data)
101
+
102
+ return G_J @ ν
103
+
104
+
105
+ @jax.jit
106
+ def centroidal_momentum_jacobian(
107
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
108
+ ) -> jtp.Matrix:
109
+ r"""
110
+ Compute the Jacobian of the centroidal momentum of the model.
111
+
112
+ Args:
113
+ model: The model to consider.
114
+ data: The data of the considered model.
115
+
116
+ Returns:
117
+ The Jacobian of the centroidal momentum of the model.
118
+
119
+ Note:
120
+ The frame corresponding to the output representation of this Jacobian is either
121
+ :math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
122
+ or :math:`G[B]`, if the active velocity representation is body-fixed.
123
+
124
+ Note:
125
+ This Jacobian is also known in the literature as Centroidal Momentum Matrix.
126
+ """
127
+
128
+ # Compute the Jacobian of the total momentum with body-fixed output representation.
129
+ # We convert the output representation either to G[W] or G[B] below.
130
+ B_Jh = js.model.total_momentum_jacobian(
131
+ model=model, data=data, output_vel_repr=VelRepr.Body
132
+ )
133
+
134
+ W_H_B = data.base_transform()
135
+ B_H_W = jaxsim.math.Transform.inverse(W_H_B)
136
+
137
+ W_p_CoM = com_position(model=model, data=data)
138
+
139
+ match data.velocity_representation:
140
+ case VelRepr.Inertial | VelRepr.Mixed:
141
+ W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
142
+ case VelRepr.Body:
143
+ W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
144
+ case _:
145
+ raise ValueError(data.velocity_representation)
146
+
147
+ # Compute the transform for 6D forces.
148
+ G_Xf_B = jaxsim.math.Adjoint.from_transform(transform=B_H_W @ W_H_G).T
149
+
150
+ return G_Xf_B @ B_Jh
151
+
152
+
153
+ @jax.jit
154
+ def locked_centroidal_spatial_inertia(
155
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
156
+ ):
157
+ """
158
+ Compute the locked centroidal spatial inertia of the model.
159
+
160
+ Args:
161
+ model: The model to consider.
162
+ data: The data of the considered model.
163
+
164
+ Returns:
165
+ The locked centroidal spatial inertia of the model.
166
+ """
167
+
168
+ with data.switch_velocity_representation(VelRepr.Body):
169
+ B_Mbb_B = js.model.locked_spatial_inertia(model=model, data=data)
170
+
171
+ W_H_B = data.base_transform()
172
+ W_p_CoM = com_position(model=model, data=data)
173
+
174
+ match data.velocity_representation:
175
+ case VelRepr.Inertial | VelRepr.Mixed:
176
+ W_H_G = W_H_GW = jnp.eye(4).at[0:3, 3].set(W_p_CoM)
177
+ case VelRepr.Body:
178
+ W_H_G = W_H_GB = W_H_B.at[0:3, 3].set(W_p_CoM)
179
+ case _:
180
+ raise ValueError(data.velocity_representation)
181
+
182
+ B_H_G = jaxlie.SE3.from_matrix(jaxsim.math.Transform.inverse(W_H_B) @ W_H_G)
183
+
184
+ B_Xv_G = B_H_G.adjoint()
185
+ G_Xf_B = B_Xv_G.transpose()
186
+
187
+ return G_Xf_B @ B_Mbb_B @ B_Xv_G
188
+
189
+
190
+ @jax.jit
191
+ def average_centroidal_velocity(
192
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
193
+ ) -> jtp.Vector:
194
+ r"""
195
+ Compute the average centroidal velocity of the model.
196
+
197
+ Args:
198
+ model: The model to consider.
199
+ data: The data of the considered model.
200
+
201
+ Returns:
202
+ The average centroidal velocity of the model.
203
+
204
+ Note:
205
+ The average velocity is expressed in the mixed frame
206
+ :math:`G = ({}^W \mathbf{p}_{\text{CoM}}, [C])`, where :math:`[C] = [W]` if the
207
+ active velocity representation is either inertial-fixed or mixed,
208
+ and :math:`[C] = [B]` if the active velocity representation is body-fixed.
209
+ """
210
+
211
+ ν = data.generalized_velocity()
212
+ G_J = average_centroidal_velocity_jacobian(model=model, data=data)
213
+
214
+ return G_J @ ν
215
+
216
+
217
+ @jax.jit
218
+ def average_centroidal_velocity_jacobian(
219
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
220
+ ) -> jtp.Matrix:
221
+ r"""
222
+ Compute the Jacobian of the average centroidal velocity of the model.
223
+
224
+ Args:
225
+ model: The model to consider.
226
+ data: The data of the considered model.
227
+
228
+ Returns:
229
+ The Jacobian of the average centroidal velocity of the model.
230
+
231
+ Note:
232
+ The frame corresponding to the output representation of this Jacobian is either
233
+ :math:`G[W]`, if the active velocity representation is inertial-fixed or mixed,
234
+ or :math:`G[B]`, if the active velocity representation is body-fixed.
235
+ """
236
+
237
+ G_J = centroidal_momentum_jacobian(model=model, data=data)
238
+ G_Mbb = locked_centroidal_spatial_inertia(model=model, data=data)
239
+
240
+ return jnp.linalg.inv(G_Mbb) @ G_J
@@ -1,6 +1,7 @@
1
1
  import abc
2
2
  import contextlib
3
3
  import dataclasses
4
+ import enum
4
5
  import functools
5
6
  from typing import ContextManager
6
7
 
@@ -11,7 +12,6 @@ import jaxlie
11
12
  from jax_dataclasses import Static
12
13
 
13
14
  import jaxsim.typing as jtp
14
- from jaxsim.high_level.common import VelRepr
15
15
  from jaxsim.utils import JaxsimDataclass, Mutability
16
16
 
17
17
  try:
@@ -20,6 +20,17 @@ except ImportError:
20
20
  from typing_extensions import Self
21
21
 
22
22
 
23
+ @enum.unique
24
+ class VelRepr(enum.IntEnum):
25
+ """
26
+ Enumeration of all supported 6D velocity representations.
27
+ """
28
+
29
+ Body = enum.auto()
30
+ Mixed = enum.auto()
31
+ Inertial = enum.auto()
32
+
33
+
23
34
  @jax_dataclasses.pytree_dataclass
24
35
  class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
25
36
  """
@@ -59,7 +70,7 @@ class ModelDataWithVelocityRepresentation(JaxsimDataclass, abc.ABC):
59
70
  # We run this in a mutable context with restoration so that any exception
60
71
  # occurring, we restore the original object in case it was modified.
61
72
  with self.mutable_context(
62
- mutability=self._mutability(), restore_after_exception=True
73
+ mutability=self.mutability(), restore_after_exception=True
63
74
  ):
64
75
  yield self
65
76
 
@@ -0,0 +1,271 @@
1
+ import functools
2
+
3
+ import jax
4
+ import jax.numpy as jnp
5
+
6
+ import jaxsim.api as js
7
+ import jaxsim.rbda
8
+ import jaxsim.typing as jtp
9
+
10
+ from .common import VelRepr
11
+
12
+
13
+ @jax.jit
14
+ def collidable_point_kinematics(
15
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
16
+ ) -> tuple[jtp.Matrix, jtp.Matrix]:
17
+ """
18
+ Compute the position and 3D velocity of the collidable points in the world frame.
19
+
20
+ Args:
21
+ model: The model to consider.
22
+ data: The data of the considered model.
23
+
24
+ Returns:
25
+ The position and velocity of the collidable points in the world frame.
26
+
27
+ Note:
28
+ The collidable point velocity is the plain coordinate derivative of the position.
29
+ If we attach a frame C = (p_C, [C]) to the collidable point, it corresponds to
30
+ the linear component of the mixed 6D frame velocity.
31
+ """
32
+
33
+ from jaxsim.rbda import collidable_points
34
+
35
+ with data.switch_velocity_representation(VelRepr.Inertial):
36
+ W_p_Ci, W_ṗ_Ci = collidable_points.collidable_points_pos_vel(
37
+ model=model,
38
+ base_position=data.base_position(),
39
+ base_quaternion=data.base_orientation(dcm=False),
40
+ joint_positions=data.joint_positions(model=model),
41
+ base_linear_velocity=data.base_velocity()[0:3],
42
+ base_angular_velocity=data.base_velocity()[3:6],
43
+ joint_velocities=data.joint_velocities(model=model),
44
+ )
45
+
46
+ return W_p_Ci, W_ṗ_Ci
47
+
48
+
49
+ @jax.jit
50
+ def collidable_point_positions(
51
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
52
+ ) -> jtp.Matrix:
53
+ """
54
+ Compute the position of the collidable points in the world frame.
55
+
56
+ Args:
57
+ model: The model to consider.
58
+ data: The data of the considered model.
59
+
60
+ Returns:
61
+ The position of the collidable points in the world frame.
62
+ """
63
+
64
+ return collidable_point_kinematics(model=model, data=data)[0]
65
+
66
+
67
+ @jax.jit
68
+ def collidable_point_velocities(
69
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
70
+ ) -> jtp.Matrix:
71
+ """
72
+ Compute the 3D velocity of the collidable points in the world frame.
73
+
74
+ Args:
75
+ model: The model to consider.
76
+ data: The data of the considered model.
77
+
78
+ Returns:
79
+ The 3D velocity of the collidable points.
80
+ """
81
+
82
+ return collidable_point_kinematics(model=model, data=data)[1]
83
+
84
+
85
+ @jax.jit
86
+ def collidable_point_forces(
87
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
88
+ ) -> jtp.Matrix:
89
+ """
90
+ Compute the 6D forces applied to each collidable point.
91
+
92
+ Args:
93
+ model: The model to consider.
94
+ data: The data of the considered model.
95
+
96
+ Returns:
97
+ The 6D forces applied to each collidable point expressed in the frame
98
+ corresponding to the active representation.
99
+ """
100
+
101
+ f_Ci, _ = collidable_point_dynamics(model=model, data=data)
102
+
103
+ return f_Ci
104
+
105
+
106
+ @jax.jit
107
+ def collidable_point_dynamics(
108
+ model: js.model.JaxSimModel, data: js.data.JaxSimModelData
109
+ ) -> tuple[jtp.Matrix, jtp.Matrix]:
110
+ r"""
111
+ Compute the 6D force applied to each collidable point and the corresponding
112
+ material deformation rate.
113
+
114
+ Args:
115
+ model: The model to consider.
116
+ data: The data of the considered model.
117
+
118
+ Returns:
119
+ The 6D force applied to each collidable point and the corresponding
120
+ material deformation rate.
121
+
122
+ Note:
123
+ The material deformation rate is always returned in the mixed frame
124
+ `C[W] = ({}^W \mathbf{p}_C, [W])`. This is convenient for integration purpose.
125
+ Instead, the 6D forces are returned in the active representation.
126
+ """
127
+
128
+ # Compute the position and linear velocities (mixed representation) of
129
+ # all collidable points belonging to the robot.
130
+ W_p_Ci, W_ṗ_Ci = js.contact.collidable_point_kinematics(model=model, data=data)
131
+
132
+ # Build the soft contact model.
133
+ soft_contacts = jaxsim.rbda.SoftContacts(
134
+ parameters=data.soft_contacts_params, terrain=model.terrain
135
+ )
136
+
137
+ # Compute the 6D force expressed in the inertial frame and applied to each
138
+ # collidable point, and the corresponding material deformation rate.
139
+ # Note that the material deformation rate is always returned in the mixed frame
140
+ # C[W] = (W_p_C, [W]). This is convenient for integration purpose.
141
+ W_f_Ci, CW_ṁ = jax.vmap(soft_contacts.contact_model)(
142
+ W_p_Ci, W_ṗ_Ci, data.state.soft_contacts.tangential_deformation
143
+ )
144
+
145
+ # Convert the 6D forces to the active representation.
146
+ f_Ci = jax.vmap(
147
+ lambda W_f_C: data.inertial_to_other_representation(
148
+ array=W_f_C,
149
+ other_representation=data.velocity_representation,
150
+ transform=data.base_transform(),
151
+ is_force=True,
152
+ )
153
+ )(W_f_Ci)
154
+
155
+ return f_Ci, CW_ṁ
156
+
157
+
158
+ @functools.partial(jax.jit, static_argnames=["link_names"])
159
+ def in_contact(
160
+ model: js.model.JaxSimModel,
161
+ data: js.data.JaxSimModelData,
162
+ *,
163
+ link_names: tuple[str, ...] | None = None,
164
+ ) -> jtp.Vector:
165
+ """
166
+ Return whether the links are in contact with the terrain.
167
+
168
+ Args:
169
+ model: The model to consider.
170
+ data: The data of the considered model.
171
+ link_names:
172
+ The names of the links to consider. If None, all links are considered.
173
+
174
+ Returns:
175
+ A boolean vector indicating whether the links are in contact with the terrain.
176
+ """
177
+
178
+ link_names = link_names if link_names is not None else model.link_names()
179
+
180
+ if set(link_names).difference(model.link_names()):
181
+ raise ValueError("One or more link names are not part of the model")
182
+
183
+ W_p_Ci = collidable_point_positions(model=model, data=data)
184
+
185
+ terrain_height = jax.vmap(lambda x, y: model.terrain.height(x=x, y=y))(
186
+ W_p_Ci[:, 0], W_p_Ci[:, 1]
187
+ )
188
+
189
+ below_terrain = W_p_Ci[:, 2] <= terrain_height
190
+
191
+ links_in_contact = jax.vmap(
192
+ lambda link_index: jnp.where(
193
+ jnp.array(model.kin_dyn_parameters.contact_parameters.body) == link_index,
194
+ below_terrain,
195
+ jnp.zeros_like(below_terrain, dtype=bool),
196
+ ).any()
197
+ )(js.link.names_to_idxs(link_names=link_names, model=model))
198
+
199
+ return links_in_contact
200
+
201
+
202
+ @jax.jit
203
+ def estimate_good_soft_contacts_parameters(
204
+ model: js.model.JaxSimModel,
205
+ *,
206
+ standard_gravity: jtp.FloatLike = jaxsim.math.StandardGravity,
207
+ static_friction_coefficient: jtp.FloatLike = 0.5,
208
+ number_of_active_collidable_points_steady_state: jtp.IntLike = 1,
209
+ damping_ratio: jtp.FloatLike = 1.0,
210
+ max_penetration: jtp.FloatLike | None = None,
211
+ ) -> jaxsim.rbda.soft_contacts.SoftContactsParams:
212
+ """
213
+ Estimate good soft contacts parameters for the given model.
214
+
215
+ Args:
216
+ model: The model to consider.
217
+ standard_gravity: The standard gravity constant.
218
+ static_friction_coefficient: The static friction coefficient.
219
+ number_of_active_collidable_points_steady_state:
220
+ The number of active collidable points in steady state supporting
221
+ the weight of the robot.
222
+ damping_ratio: The damping ratio.
223
+ max_penetration:
224
+ The maximum penetration allowed in steady state when the robot is
225
+ supported by the configured number of active collidable points.
226
+
227
+ Returns:
228
+ The estimated good soft contacts parameters.
229
+
230
+ Note:
231
+ This method provides a good starting point for the soft contacts parameters.
232
+ The user is encouraged to fine-tune the parameters based on the
233
+ specific application.
234
+ """
235
+
236
+ def estimate_model_height(model: js.model.JaxSimModel) -> jtp.Float:
237
+ """"""
238
+
239
+ zero_data = js.data.JaxSimModelData.build(
240
+ model=model,
241
+ soft_contacts_params=jaxsim.rbda.soft_contacts.SoftContactsParams(),
242
+ )
243
+
244
+ W_pz_CoM = js.com.com_position(model=model, data=zero_data)[2]
245
+
246
+ if model.floating_base():
247
+ W_pz_C = collidable_point_positions(model=model, data=zero_data)[:, -1]
248
+ return 2 * (W_pz_CoM - W_pz_C.min())
249
+
250
+ return 2 * W_pz_CoM
251
+
252
+ max_δ = (
253
+ max_penetration
254
+ if max_penetration is not None
255
+ else 0.005 * estimate_model_height(model=model)
256
+ )
257
+
258
+ nc = number_of_active_collidable_points_steady_state
259
+
260
+ sc_parameters = (
261
+ jaxsim.rbda.soft_contacts.SoftContactsParams.build_default_from_jaxsim_model(
262
+ model=model,
263
+ standard_gravity=standard_gravity,
264
+ static_friction_coefficient=static_friction_coefficient,
265
+ max_penetration=max_δ,
266
+ number_of_active_collidable_points_steady_state=nc,
267
+ damping_ratio=damping_ratio,
268
+ )
269
+ )
270
+
271
+ return sc_parameters