Graduate FK code out of labs (#550)

* Replaced differentiable robot model with labs FK implementation. Jacobians tests not passing yet. * Revert back to dict based API. * Moved FK out of labs. * Add urdf_parser_py to requirements. * Remove duplicate pose computation when computing jacobians. * Moved FK code to new torchkin project. * Update precommit config with torchkin info. * Update theseus with torchkin import. * Remove unnecessary function. * Fixed torchkin build errors. * Move FK test data to a different folder. * test_jacobian works * no need to compute dlocal since we assume these joint states yield the poses * Update precommit config with torchkin info. * add warning message * update readme * Update warning for body jacobians. * Fix CI errors. --------- Co-authored-by: Taosha Fan <[email protected]> Co-authored-by: Mustafa Mukadam <[email protected]>
facebookresearch · Jun 27, 2023 · a12c152 · a12c152
1 parent 2281b53
commit a12c152
Show file tree

Hide file tree

Showing 22 changed files with 560 additions and 130 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -25,13 +25,6 @@ update_and_install_python: &update_and_install_python
         sudo apt-get update
         sudo apt-get install -y python3.9 python3.9-dev
 
-update_and_install_conda: &update_and_install_conda
-  - run:
-      name: "Preparing environment: python"
-      command: |
-        wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
-        /bin/bash ~/miniconda.sh -b -p ~/conda
-
 install_suitesparse: &install_suitesparse
   - run:
       name: "Preparing environment: suitesparse"
@@ -83,9 +76,16 @@ setup_environment: &setup_environment
         pip install --progress-bar off --upgrade pip
         pip install --progress-bar off --upgrade setuptools
 
+# The following two are only used by the build wheel test
+update_and_install_conda: &update_and_install_conda
+  - run:
+      name: "Preparing environment: python"
+      command: |
+        wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
+        /bin/bash ~/miniconda.sh -b -p ~/conda
 setup_environment_310: &setup_environment_310
   - run:
-      name: Setup virtualenv and tools
+      name: Setup virtualenv and tools for python 3.10
       working_directory: ~/project
       command: |
         export PATH=~/conda/bin:$PATH
@@ -99,6 +99,9 @@ install_torch_cuda11: &install_torch_cuda11
       name: Install Torch for cuda11
       working_directory: ~/project
       command: |
+        export PATH=~/conda/bin:$PATH
+        source activate theseus
+        which python && which pip
         pip install --progress-bar off torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
         python -c 'import torch; print("Torch version:", torch.__version__); assert torch.cuda.is_available()'
 
@@ -155,6 +158,16 @@ setup_project: &setup_project
         pip install -r requirements/dev.txt
         BASPACHO_ROOT_DIR=~/baspacho pip install --progress-bar off -e .
 
+install_torchlie_torchkin: &install_torchlie_torchkin
+  - run:
+      name: Installing torchlie and torchkin
+      working_directory: ~/project
+      command: |
+        cd torchlie
+        pip install -e .
+        cd ../torchkin
+        pip install -e .
+
 run_pre_commit: &run_pre_commit
   - run:
       name: Running precommit checks
@@ -205,10 +218,12 @@ build_cuda11_wheel: &build_cuda11_wheel
       name: Building wheel for CUDA 11
       working_directory: /home/circleci
       command: |
+        export PATH=~/conda/bin:$PATH
+        source activate theseus
+        which python && which pip
         THESEUS_GIT_COMMIT=$(git --git-dir project/.git log --format="%H" -n 1)
         THESEUS_VERSION=$(grep -Eo  "[0-9].[0-9].[0-9][.0-9a-z]*" project/theseus/_version.py | tail -n 1)
-        ./project/build_scripts/build_wheel.sh . ${THESEUS_GIT_COMMIT} 11.6 ${THESEUS_VERSION}
-        source ~/conda/bin/activate theseus
+        ./project/build_scripts/build_wheel.sh . ${THESEUS_GIT_COMMIT} 11.7 ${THESEUS_VERSION}
         pip install $(ls */*.whl)
         pip install -r project/requirements/dev.txt
 
@@ -218,7 +233,9 @@ run_tests_from_wheel: &run_tests_from_wheel
       working_directory: ~/project
       command: |
           mv theseus theseus_tmp
-          source ~/conda/bin/activate theseus
+          export PATH=~/conda/bin:$PATH
+          source activate theseus
+          which python && which pip
           python -m pytest tests/theseus_tests/test_theseus_layer.py
           pytest -s tests/theseus_tests/ -m "cudaext"
 
@@ -239,6 +256,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *install_torch
       - <<: *build_baspacho_cpu
+      - <<: *install_torchlie_torchkin
       - <<: *run_pre_commit
 
   py39_precommit:
@@ -249,6 +267,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *install_torch
       - <<: *build_baspacho_cpu
+      - <<: *install_torchlie_torchkin
       - <<: *run_pre_commit
 
   py310_precommit:
@@ -259,6 +278,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *install_torch
       - <<: *build_baspacho_cpu
+      - <<: *install_torchlie_torchkin
       - <<: *run_pre_commit
 
   # End-to-end tests
@@ -271,6 +291,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *install_torch
       - <<: *build_baspacho_cpu
+      - <<: *install_torchlie_torchkin
       - <<: *run_optimizer_tests
 
   py39_optimizer_tests:
@@ -282,6 +303,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *install_torch
       - <<: *build_baspacho_cpu
+      - <<: *install_torchlie_torchkin
       - <<: *run_optimizer_tests
 
   py310_optimizer_tests:
@@ -293,6 +315,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *install_torch
       - <<: *build_baspacho_cpu
+      - <<: *install_torchlie_torchkin
       - <<: *run_optimizer_tests
 
   # Lie group tests
@@ -304,6 +327,7 @@ jobs:
       - checkout
       - <<: *install_suitesparse
       - <<: *install_torch
+      - <<: *install_torchlie_torchkin
       - <<: *run_lie_group_tests
 
   py39_lie_tests:
@@ -314,6 +338,7 @@ jobs:
       - checkout
       - <<: *install_suitesparse
       - <<: *install_torch
+      - <<: *install_torchlie_torchkin
       - <<: *run_lie_group_tests
 
   py310_lie_tests:
@@ -324,6 +349,7 @@ jobs:
       - checkout
       - <<: *install_suitesparse
       - <<: *install_torch
+      - <<: *install_torchlie_torchkin
       - <<: *run_lie_group_tests
 
   # Other unit tests
@@ -334,6 +360,7 @@ jobs:
       - checkout
       - <<: *install_suitesparse
       - <<: *install_torch
+      - <<: *install_torchlie_torchkin
       - <<: *run_other_unit_tests
 
   py39_other_unit_tests:
@@ -343,6 +370,7 @@ jobs:
       - checkout
       - <<: *install_suitesparse
       - <<: *install_torch
+      - <<: *install_torchlie_torchkin
       - <<: *run_other_unit_tests
 
   py310_other_unit_tests:
@@ -352,6 +380,7 @@ jobs:
       - checkout
       - <<: *install_suitesparse
       - <<: *install_torch
+      - <<: *install_torchlie_torchkin
       - <<: *run_other_unit_tests
 
 # ------- GPU only jobs -------
@@ -366,6 +395,7 @@ jobs:
       - <<: *install_torch_cuda11
       - <<: *install_recent_cmake_2004
       - <<: *build_baspacho_cuda
+      - <<: *install_torchlie_torchkin
       - <<: *setup_project
       - <<: *run_gpu_tests
 
@@ -378,6 +408,7 @@ jobs:
       - <<: *install_suitesparse
       - <<: *setup_cuda11_libs
       - <<: *setup_environment_310
+      - <<: *install_torch_cuda11
       - <<: *build_cuda11_wheel
       - <<: *run_tests_from_wheel
 

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,28 +1,28 @@
-exclude: "setup.py|torchlie/setup.py"
+exclude: "setup.py|torchlie/setup.py|torchkin/setup.py"
 
 repos:
   - repo: https://github.com/psf/black
     rev: 23.3.0
     hooks:
       - id: black
-        files: "^theseus|^examples|^tests|^torchlie"
+        files: "^theseus|^examples|^tests|^torchlie|^torchkin"
 
   - repo: https://github.com/pycqa/flake8
     rev: 6.0.0
     hooks:
       - id: flake8
-        files: "^theseus|^torchlie"
+        files: "^theseus|^torchlie|^torchkin"
 
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.4.0
     hooks:
       - id: mypy
         additional_dependencies: [torch, tokenize-rt==3.2.0, types-PyYAML, types-mock]
         args: [--no-strict-optional, --ignore-missing-imports, --implicit-reexport, --explicit-package-bases]
-        files: "^theseus|^torchlie"
+        files: "^theseus|^torchlie|^torchkin"
 
   - repo: https://github.com/pycqa/isort
     rev: 5.12.0
     hooks:
       - id: isort
-        files: "^theseus|^examples|^torchlie"
+        files: "^theseus|^examples|^torchlie|^torchkin"
diff --git a/requirements/dev.txt b/requirements/dev.txt
@@ -3,9 +3,9 @@ flake8>=3.8.4
 mypy>=0.981
 pre-commit>=2.9.2
 isort>=5.6.4
-differentiable-robot-model>=0.2.3
 types-PyYAML==5.4.3
 mock>=4.0.3
 types-mock>=4.0.8
 Sphinx==5.0.2
-sphinx-rtd-theme==1.0.0
+sphinx-rtd-theme==1.0.0
+hydra-core
diff --git a/requirements/main.txt b/requirements/main.txt
@@ -4,4 +4,5 @@ scikit-sparse>=0.4.5
 pytest>=6.2.1
 pybind11>=2.7.1
 semantic-version==2.10.0
-torchlie
+torchlie
+torchkin
diff --git a/tests/theseus_tests/embodied/kinematics/test_inverse_kinematics.py b/tests/theseus_tests/embodied/kinematics/test_inverse_kinematics.py
@@ -14,12 +14,7 @@
 EE_NAME = "panda_virtual_ee_link"
 HOME_POSE = torch.Tensor([-0.1394, -0.0205, -0.0520, -2.0691, 0.0506, 2.0029, -0.9168])
 ERR_SCALING = torch.Tensor([1.0, 1.0, 1.0, 5.0, 5.0, 5.0])
-
-
-@pytest.fixture
-def robot_model():
-    urdf_path = os.path.join(os.path.dirname(__file__), "data/panda_no_gripper.urdf")
-    return th.eb.UrdfRobotModel(urdf_path)
+URDF_PATH = os.path.join(os.path.dirname(__file__), "data/panda_no_gripper.urdf")
 
 
 @pytest.fixture(params=[1, 3])
@@ -54,9 +49,11 @@ def ee_pose_target(batch_size):
 
 
 @pytest.mark.parametrize("is_grad_enabled", [True, False])
-def test_ik_optimization(robot_model, batch_size, ee_pose_target, is_grad_enabled):
+def test_ik_optimization(batch_size, ee_pose_target, is_grad_enabled):
     """Sets up inverse kinematics as an optimization problem that uses forward kinematics"""
 
+    robot_model = th.eb.UrdfRobotModel(URDF_PATH, link_names=[EE_NAME])
+
     # Define cost (distance between desired and current ee pose)
     def ee_pose_err_fn(optim_vars, aux_vars):
         (theta,) = optim_vars
@@ -83,7 +80,7 @@ def ee_pose_err_fn(optim_vars, aux_vars):
     objective.add(cost_function)
     optimizer = th.LevenbergMarquardt(
         objective,
-        max_iterations=15,
+        max_iterations=25,
         step_size=0.5,
     )
     theseus_optim = th.TheseusLayer(optimizer)
@@ -106,4 +103,6 @@ def ee_pose_err_fn(optim_vars, aux_vars):
         )
 
     # Check result
-    assert torch.allclose(info.best_err, torch.zeros(batch_size), atol=1e-2)
+    torch.testing.assert_close(
+        torch.zeros(batch_size), info.best_err, atol=1e-3, rtol=1e-3
+    )