test with more menagerie models

Author: Peter

.github/workflows/python-package.yml

@@ -30,6 +30,10 @@ jobs:
           python -m pip install .[test]
           python -m pip install flake8 pytest
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+      - name: Clone mujoco_menagerie repository into the tests/ folder
+        run: |
+          git clone https://github.com/google-deepmind/mujoco_menagerie
+        working-directory: ${{ runner.workspace }}/pytorch_kinematics/tests
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
@@ -38,35 +42,4 @@ jobs:
           flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
       - name: Test with pytest
         run: |
-          pytest
-
-  test-shadow-hand:
-
-    runs-on: ubuntu-latest
-
-    steps:
-      - uses: actions/checkout@v3
-      - name: Install dependencies
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install .[test]
-          python -m pip install pytest
-      - name: Clone mujoco_menagerie repository
-        run: |
-          git clone https://github.com/google-deepmind/mujoco_menagerie
-        working-directory: ${{ runner.workspace }}
-
-      - name: Debug
-        run: |
-          tree
-        working-directory: ${{ runner.workspace }}
-
-      - name: Copy test_shadow_hand.py to mujoco_menagerie/shadow_hand
-        run: |
-          cp pytorch_kinematics/tests/menagerie_test_shadow_hand.py mujoco_menagerie/shadow_hand/test_shadow_hand.py
-        working-directory: ${{ runner.workspace }}
-
-      - name: Run test_shadow_hand.py
-        run: |
-          pytest test_shadow_hand.py
-        working-directory: ${{ runner.workspace }}/mujoco_menagerie/shadow_hand
+          pytest

.gitignore

@@ -3,4 +3,8 @@
 __pycache__
 temp*
 build
-dist
+dist
+# These are cloned/generated when testing with mujoco
+tests/MUJOCO_LOG.TXT
+tests/mujoco_menagerie/
+

tests/menagerie_test_shadow_hand.py

@@ -0,0 +1,51 @@
+import os
+import pathlib
+
+import numpy as np
+
+import pytorch_kinematics as pk
+
+# Find all files named "scene*.xml" in the "mujoco_menagerie" directory
+_MENAGERIE_ROOT = pathlib.Path(__file__).parent / 'mujoco_menagerie'
+_XMLS_AND_BODIES = {
+    # 'agility_cassie/scene.xml':            'cassie-pelvis',  # not supported because it has a ball joint
+    'anybotics_anymal_b/scene.xml':        'base',
+    'anybotics_anymal_c/scene.xml':        'base',
+    'franka_emika_panda/scene.xml':        'link0',
+    'google_barkour_v0/scene.xml':         'chassis',
+    'google_barkour_v0/scene_barkour.xml': 'chassis',
+    # 'hello_robot_stretch/scene.xml':       'base_link', # not supported because it has composite joints
+    'kuka_iiwa_14/scene.xml':              'base',
+    'rethink_robotics_sawyer/scene.xml':   'base',
+    'robotiq_2f85/scene.xml':              'base_mount',
+    'robotis_op3/scene.xml':               'body_link',
+    'shadow_hand/scene_left.xml':          'lh_forearm',
+    'shadow_hand/scene_right.xml':         'rh_forearm',
+    'ufactory_xarm7/scene.xml':            'link_base',
+    'unitree_a1/scene.xml':                'trunk',
+    'unitree_go1/scene.xml':               'trunk',
+    'universal_robots_ur5e/scene.xml':     'base',
+    'wonik_allegro/scene_left.xml':        'palm',
+    'wonik_allegro/scene_right.xml':       'palm',
+}
+
+
+def test_menagerie():
+    for xml_filename, body in _XMLS_AND_BODIES.items():
+        xml_filename = _MENAGERIE_ROOT / xml_filename
+        xml_dir = xml_filename.parent
+        # Menagerie files assume the current working directory is the directory of the scene.xml
+        os.chdir(xml_dir)
+        with xml_filename.open('r') as f:
+            xml = f.read()
+        chain = pk.build_chain_from_mjcf(xml, body)
+        print(xml_filename)
+        print("=" * 32)
+        print(f"\t {chain.get_frame_names()}")
+        print(f"\t {chain.get_joint_parameter_names()}")
+        th = np.zeros(len(chain.get_joint_parameter_names()))
+        fk_dict = chain.forward_kinematics(th, end_only=True)
+
+
+if __name__ == '__main__':
+    test_menagerie()