Theseus installation issues (python 3.8, pytorch 1.12.1) ImportError: cannot import name 'tree_map_only' from 'torch.utils._pytree'

[DanielTakeshi]

I am trying to build Theseus with an earlier pytorch version (1.12.1) and it seems to be running into some problems. There are several ways I am trying to do this following current documentation and prior discussions. I'm using Ubuntu 18.04.

Attempt 1 (pip install theseus-ai)

Steps to reproduce, run these commands:

conda create -n thtest python=3.8 -y && conda activate thtest
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge -y
pip install theseus-ai

Then the error that I get (which I can also reproduce on a second machine that I have):

(thtest) seita@starship:~/tmp $ python
Python 3.8.17 (default, Jul  5 2023, 21:04:15) 
[GCC 11.2.0] :: Anaconda, Inc. on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import theseus
/home/seita/miniconda3/envs/thtest/lib/python3.8/site-packages/theseus/_version.py:30: FutureWarning: Using torch < 2.0 for theseus is deprecated and compatibility will be discontinued in future releases.
  warnings.warn(
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/seita/miniconda3/envs/thtest/lib/python3.8/site-packages/theseus/__init__.py", line 8, in <module>
    from .global_params import set_global_params  # usort: skip
  File "/home/seita/miniconda3/envs/thtest/lib/python3.8/site-packages/theseus/global_params.py", line 10, in <module>
    from torchlie.global_params import _TORCHLIE_GLOBAL_PARAMS as _LIE_GP
  File "/home/seita/miniconda3/envs/thtest/lib/python3.8/site-packages/torchlie/__init__.py", line 8, in <module>
    from .lie_tensor import (  # usort: skip
  File "/home/seita/miniconda3/envs/thtest/lib/python3.8/site-packages/torchlie/lie_tensor.py", line 12, in <module>
    from torch.utils._pytree import tree_flatten, tree_map_only
ImportError: cannot import name 'tree_map_only' from 'torch.utils._pytree' (/home/seita/miniconda3/envs/thtest/lib/python3.8/site-packages/torch/utils/_pytree.py)
>>> exit()
(thtest) seita@starship:~/tmp $ 

Attempt 2 (pip install with a custom-made wheel)

Following prior discussions #442 and my earlier one #530 let's try and make a custom wheel file. I'm testing on Ubuntu 18.04 machine with nvcc --version returning CUDA 11.6. The only thing I changed in the build script was the python version (making it 3.8) as well as making the torch version 1.12.1, see the git diff:

(thtest) seita@starship:~/theseus (main) $ git diff
diff --git a/build_scripts/build_wheel.sh b/build_scripts/build_wheel.sh
index c969b31..a868565 100755
--- a/build_scripts/build_wheel.sh
+++ b/build_scripts/build_wheel.sh
@@ -37,7 +37,7 @@ CUDA_VERSION=$3
 TH_VERSION=${4:-${COMMIT}}
 INCLUDE_LABS=${5:-0}
 NIGHTLY=${6:-0}
-PY_VER=${7:-3.10}
+PY_VER=${7:-3.8}
 
 if [[ ${INCLUDE_LABS} != 0 ]]
 then
@@ -66,7 +66,7 @@ CUDA_VERSION_IS_SUPPORTED=$(echo "cpu ${SUPPORTED_CUDA_VERSIONS}" | grep -w ${CU
 
 CUDA_SUFFIX=$(echo ${CUDA_VERSION} | sed 's/[.]//g')
 
-TORCH_VERSION='"torch>=1.13"'
+TORCH_VERSION='"torch==1.12.1"'
 if [[ ${CUDA_VERSION} == "cpu" ]]
 then 
     DEVICE_TAG=cpu
(thtest) seita@starship:~/theseus (main) $ 

The command that I run is:

./build_scripts/build_wheel.sh ROOT_DIR 0.2.1 11.6

This produces a new .whl file. It took a while but it ran to completion and made the file. The 0.2.1 is because I use the latest version of theseus (as of today, on PyPI) and the 11.6 is the CUDA version. Then I start a new conda environment and try to pip install the above in the 3rd command:

conda create -n thtest python=3.8 -y && conda activate thtest
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge -y
pip install ROOT_DIR/theseus_docker_3.8/theseus_ai-0.2.1+cu116-cp38-cp38-manylinux_2_17_x86_64.whl

However it still appears to run into the same error when I import theseus, on both Ubutnu 18.04 machines I'm testing this on (same error as above). Not sure why as #442 reported it was successful but there might be subtle changes.

There's one point about functorch in that discussion, though. Interestingly doing pip install functorch then appears to override my PyTorch version and creates torch 2.0.1 in my conda environment.

I want to stick to the earlier PyTorch version so I deleted thtest env, and then did this:

conda create -n thtest python=3.8 -y && conda activate thtest
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.6 -c pytorch -c conda-forge -y
pip install functorch==0.2.1
pip install ROOT_DIR/theseus_docker_3.8/theseus_ai-0.2.1+cu116-cp38-cp38-manylinux_2_17_x86_64.whl

It runs into the same error.

Just checking and building upon our earlier discussion should we expect to need to use pytorch 2.0 or higher for theseus to properly work? (And if this is the case maybe we can update the installation instructions to make this more explicit?) Thanks!

[DanielTakeshi]

Maybe there's something that happened between pytorch 1.12.1 and 1.13 since I could get 1.13 with a similar setup. I updated my earlier discussion #530

Or maybe I didn't use TORCH_VERSION correctly here?

[luisenp]

IIRC, tree_map_only was introduced in 1.13, which means that torchlie is not compatible with older versions of PyTorch. Thanks for pointing this out, we should update the README.

[DanielTakeshi]

Thanks, alternatively is there a way to remove torchlie or does it come built-into theseus? I don't know if I would need the torchlie stuff. Or do you think it's too embedded within theseus to be worth the effort to disentangle it?

[luisenp]

As of version 0.2.0, our Lie groups optimization variables are based on torchlie, so there is no way to remove this dependency anymore.

[DanielTakeshi]

Thanks @luisenp . Also I did some further tests, and I'm not sure if 1.13 should be considered a compatible version of PyTorch. To test I created two conda environments:

conda create -n torch113 python=3.8 -y && conda activate torch113
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia -y
pip install theseus-ai==0.2.1

and

conda create -n torch200 python=3.8 -y && conda activate torch200
conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.7 -c pytorch -c nvidia -y
pip install theseus-ai==0.2.1

I then did a fresh git clone of the current theseus repository as of today (main branch). Then I went into the examples directory and tried python inverse_kinematics.py. The conda environment with pytorch==2.0.0 works in that the code runs without errors. However the conda environment with the earlier version of pytorch results in:

(torch113) seita@starship:~/theseus/examples (main) $ python inverse_kinematics.py 
/home/seita/miniconda3/envs/torch113/lib/python3.8/site-packages/theseus/_version.py:30: FutureWarning: Using torch < 2.0 for theseus is deprecated and compatibility will be discontinued in future releases.
  warnings.warn(
---------------------------------------------------
Use Body Jacobian
---------------------------------------------------
Traceback (most recent call last):
  File "inverse_kinematics.py", line 96, in <module>
    run_ik_using_body_or_spatial_jacobian(jfk_b, 100, 0.5, use_body_jacobian=True)
  File "inverse_kinematics.py", line 81, in run_ik_using_body_or_spatial_jacobian
    target_pose: torch.Tensor = fk(target_theta)[0]
  File "/home/seita/miniconda3/envs/torch113/lib/python3.8/site-packages/torchkin/forward_kinematics.py", line 144, in forward_kinematics
    output = ForwardKinematics.apply(angles)
TypeError: forward() takes 1 positional argument but 2 were given
(torch113) seita@starship:~/theseus/examples (main) $ 

So I would suggest having torch 2.0 as the minimum compatible version.

[luisenp]

Thanks for reporting, I was not aware of this restriction. Unless I'm forgetting about something, this should be only an issue for torchkin, but not for torchlie (we wrote this with <2.0 compatibility in mind). I'll try to fix backward compatibility if I have time over the next week, otherwise we'll report 2.0 as the minimum compatible version. Thanks.

cc @fantaosha