Add TAXPoseD

.dockerignore

@@ -0,0 +1,3 @@
+third_party/ndf_robot/src/ndf_robot/data/
+third_party/ndf_robot/src/ndf_robot/eval_data/
+third_party/ndf_robot/src/ndf_robot/descriptions/

.gitignore

@@ -377,8 +377,11 @@ results/
 rollouts/
 wandb/
 trained_models/ndf/
+artifacts/
+wandb_artifacts/
 recent_training_paths.txt
 
 # Ignore all checkpoint configs except for release files.
 configs/checkpoints/**/*.yaml
 !configs/checkpoints/**/release.yaml
+!configs/checkpoints/**/pretraining.yaml

.vscode/settings.json

@@ -6,11 +6,14 @@
         "torch.*"
     ],
     "editor.codeActionsOnSave": {
-        "source.organizeImports": true
+        "source.organizeImports": "explicit"
     },
     "files.exclude": {
         "data/*": true,
     },
+    "python.analysis.exclude": [
+        "logs/**",
+    ],
     "python.linting.enabled": true,
     "files.watcherExclude": {
         "**/.git/objects/**": true,

Dockerfile

@@ -0,0 +1,88 @@
+# Use the official Ubuntu 20.04 image as the base
+FROM ubuntu:20.04
+
+# Set environment variables to avoid interactive prompts during installation
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install necessary dependencies
+RUN apt-get update && \
+    apt-get install -y curl git build-essential libssl-dev zlib1g-dev libbz2-dev \
+    git \
+    libreadline-dev libsqlite3-dev wget llvm libncurses5-dev libncursesw5-dev \
+    xz-utils tk-dev libffi-dev liblzma-dev python-openssl && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install pyenv
+ENV CODING_ROOT="/opt/baeisner"
+
+WORKDIR $CODING_ROOT
+RUN git clone --depth=1 https://github.com/pyenv/pyenv.git .pyenv
+
+ENV PYENV_ROOT="$CODING_ROOT/.pyenv"
+ENV PATH="$PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH"
+
+# Install Python 3.10 using pyenv
+RUN pyenv install 3.9.12
+RUN pyenv global 3.9.12
+
+# Install PyTorch with CUDA support (make sure to adjust this depending on your CUDA version)
+RUN pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+
+# Install pytorch geometric.
+RUN pip install torch-scatter==2.0.9 torch-sparse==0.6.15 torch-cluster==1.6.0 torch-spline-conv==1.2.1 pyg_lib==0.1.0 -f https://data.pyg.org/whl/torch-1.13.0+cu116.html
+
+# Install pytorch3d
+RUN pip install fvcore iopath && \
+    pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu116_pyt1130/download.html
+
+# Download CoppeliaSim
+RUN mkdir $CODING_ROOT/.coppelia
+WORKDIR $CODING_ROOT/.coppelia
+RUN curl https://www.coppeliarobotics.com/files/V4_1_0/CoppeliaSim_Edu_V4_1_0_Ubuntu20_04.tar.xz -o CoppeliaSim_Edu_V4_1_0_Ubuntu20_04.tar.xz && \
+    tar -xf CoppeliaSim_Edu_V4_1_0_Ubuntu20_04.tar.xz && \
+    rm CoppeliaSim_Edu_V4_1_0_Ubuntu20_04.tar.xz
+
+# modify environment variables
+ENV COPPELIASIM_ROOT="$CODING_ROOT/.coppelia/CoppeliaSim_Edu_V4_1_0_Ubuntu20_04"
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:$COPPELIASIM_ROOT"
+ENV QT_QPA_PLATFORM_PLUGIN_PATH="$COPPELIASIM_ROOT"
+
+# Install CFFI
+RUN pip install cffi==1.14.2 wheel
+
+# Install PyRep
+RUN pip install --no-build-isolation "pyrep @ git+https://github.com/stepjam/PyRep.git"
+
+# Make the working directory the home directory
+RUN mkdir $CODING_ROOT/code
+WORKDIR $CODING_ROOT/code
+
+# Only copy in the source code that is necessary for the dependencies to install
+COPY ./taxpose $CODING_ROOT/code/taxpose
+COPY ./third_party $CODING_ROOT/code/third_party
+COPY ./setup.py $CODING_ROOT/code/setup.py
+COPY ./pyproject.toml $CODING_ROOT/code/pyproject.toml
+RUN pip install -e ".[rlbench]"
+RUN pip install -e third_party/ndf_robot
+
+# Changes to the configs and scripts will not require a rebuild
+COPY ./configs $CODING_ROOT/code/configs
+COPY ./scripts $CODING_ROOT/code/scripts
+
+RUN git config --global --add safe.directory /root/code
+
+# Make a data directory.
+RUN mkdir $CODING_ROOT/data
+
+# Make a logs directory.
+RUN mkdir $CODING_ROOT/logs
+
+# Install gdown.
+RUN pip install gdown
+
+# Copy the download script.
+COPY ./download_data.sh $CODING_ROOT/code/download_data.sh
+
+# Set up the entry point
+CMD ["python", "-c", "import torch; print(torch.cuda.is_available())"]

README.md

@@ -57,7 +57,9 @@ For our experiments, we installed torch 1.11 with cuda 11.3:
 
 ```
 # pip install torch==1.11.0+cu113 torchvision==0.12.0+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
-pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+# pip install torch==1.13.0+cu116 torchvision==0.14.0+cu116 --extra-index-url https://download.pytorch.org/whl/cu116
+pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
+
 ```
 
 ### `pytorch-geometric`
@@ -68,7 +70,8 @@ For our experiments, we installed with the following command:
 
 ```
 # pip install torch-scatter==2.0.9 torch-sparse==0.6.15 torch-cluster==1.6.0 torch-spline-conv==1.2.1 pyg_lib==0.1.0 -f https://data.pyg.org/whl/torch-1.11.0+cu113.html
-pip install torch-scatter==2.0.9 torch-sparse==0.6.15 torch-cluster==1.6.0 torch-spline-conv==1.2.1 pyg_lib==0.1.0 -f https://data.pyg.org/whl/torch-1.13.0+cu116.html
+# pip install torch-scatter==2.0.9 torch-sparse==0.6.15 torch-cluster==1.6.0 torch-spline-conv==1.2.1 pyg_lib==0.1.0 -f https://data.pyg.org/whl/torch-1.13.0+cu116.html
+pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.1.0+cu121.html
 ```
 
 ### `pytorch3d`
@@ -80,7 +83,9 @@ We ran the following:
 ```
 pip install fvcore iopath
 # pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu113_pyt1110/download.html
-pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu116_pyt1130/download.html
+# pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu116_pyt1130/download.html
+# pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu118_pyt201/download.html
+
 ```
 
 ### `dgl`
@@ -648,3 +653,35 @@ These are real-world experiments we ran, and so can't be reproduced purely from
 ## Supplement Tables 11-14: Expanded results
 
 These are granular results of the experiments in Table 1.
+
+
+## Docker
+
+### Build a Docker container.
+
+```
+docker build -t beisner/taxpose .
+```
+
+### Run training.
+
+```
+WANDB_API_KEY=<API_KEY>
+USERNAME=<USERNAME>
+# Optional: mount current directory to run / test new code.
+# Mount data directory to access data.
+docker run \
+    --shm-size=256m\
+    -v /data/ndf:/opt/baeisner/data \
+    -v $(pwd)/trained_models:/opt/baeisner/code/trained_models \
+    -v $(pwd)/logs:/opt/baeisner/logs \
+    --gpus all \
+    -e WANDB_API_KEY=$WANDB_API_KEY \
+    -e WANDB_DOCKER_IMAGE=beisner/taxpose \
+    beisner/taxpose python scripts/train_residual_flow.py \
+        task=mug_grasp \
+        model=taxpose \
+        +mode=train \
+        benchmark.dataset_root=/opt/baeisner/data \
+        log_dir=/opt/baeisner/logs
+```

autobot.md

@@ -0,0 +1,48 @@
+
+Downloading the data.
+
+```bash
+singularity exec \
+--nv \
+--pwd /opt/$(whoami)/code \
+-B /scratch/$(whoami)/data:/opt/data \
+docker://beisner/taxpose \
+/opt/code/scripts/download_data.sh \
+    /opt/data/ndf
+```
+
+
+Pretraining.
+
+```bash
+SINGULARITYENV_CUDA_VISIBLE_DEVICES=0 \
+SINGULARITYENV_WANDB_DOCKER_IMAGE=taxpose \
+singularity exec \
+--nv \
+-B /scratch/$(whoami)/data:/data \
+-B /scratch/$(whoami)/logs:/opt/logs \
+docker://beisner/taxpose \
+python scripts/train_residual_flow.py \
+    log_dir=/opt/logs
+```
+
+
+
+Training the model.
+
+```bash
+SINGULARITYENV_CUDA_VISIBLE_DEVICES=0 \
+SINGULARITYENV_WANDB_DOCKER_IMAGE=taxpose \
+singularity exec \
+--nv \
+--pwd /opt/$(whoami)/code \
+-B /scratch/$(whoami)/data:/opt/data \
+-B /scratch/$(whoami)/logs:/opt/logs \
+docker://beisner/taxpose \
+python scripts/train_residual_flow.py \
+    task=mug_grasp \
+    model=taxpose \
+    +mode=train \
+    benchmark.dataset_root=/opt/data \
+    log_dir=/opt/logs
+```

configs/_logging.yaml

@@ -0,0 +1,32 @@
+# Where logs go, i.e. the top folder.
+log_dir: ${hydra:runtime.cwd}/logs
+
+output_dir: ${hydra:runtime.output_dir}
+
+# This has to come from above.
+job_type: ???
+
+hydra:
+  run:
+    dir: ${log_dir}/${hydra.job.name}/${now:%Y-%m-%d}/${now:%H-%M-%S}
+  sweep:
+    dir: ${log_dir}/${hydra.job.name}/sweep/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  job:
+    chdir: True
+    name: ${job_type}
+
+lightning:
+  checkpoint_dir: ${output_dir}/checkpoints
+
+wandb:
+  entity: r-pad
+  project: taxpose
+
+  # Group is for grouping runs together (i.e. a train run and an eval run).
+  group: ???
+
+  # Where to dump wandb logs, etc.
+  save_dir: ${output_dir}
+  # Put artifacts at the toplevel so that we don't have to re-download each time...
+  artifact_dir: ${hydra:runtime.cwd}/wandb_artifacts

configs/benchmark/ndf.yaml

@@ -0,0 +1,2 @@
+name: ndf
+dataset_root: ${..data_root}/ndf

configs/benchmark/real_world_mug.yaml

@@ -0,0 +1,2 @@
+name: real_world_mug
+dataset_root: /home/beisner/code/multi_project/cam_ready_trainingdata

configs/benchmark/rlbench.yaml

@@ -0,0 +1,2 @@
+name: rlbench
+dataset_root: ${..data_root}/rlbench

configs/checkpoints/ndf/mug/release.yaml

@@ -2,8 +2,8 @@ defaults:
   - taxpose_base
 
 arbitrary:
-  grasp: ${hydra:runtime.cwd}/trained_models/ndf/${object_class.name}/arbitrary/grasp.ckpt
-  place: ${hydra:runtime.cwd}/trained_models/ndf/${object_class.name}/arbitrary/place.ckpt
+  grasp: ${hydra:runtime.cwd}/trained_models/ndf/${task.name}/arbitrary/grasp.ckpt
+  place: ${hydra:runtime.cwd}/trained_models/ndf/${task.name}/arbitrary/place.ckpt
 upright:
-  grasp: ${hydra:runtime.cwd}/trained_models/ndf/${object_class.name}/upright/grasp.ckpt
-  place: ${hydra:runtime.cwd}/trained_models/ndf/${object_class.name}/upright/place.ckpt
+  grasp: ${hydra:runtime.cwd}/trained_models/ndf/${task.name}/upright/grasp.ckpt
+  place: ${hydra:runtime.cwd}/trained_models/ndf/${task.name}/upright/place.ckpt

configs/checkpoints/pretraining.yaml

@@ -0,0 +1,68 @@
+dgcnn:
+  nce_loss:
+    # These got retrained!
+    # gripper: ${hydra:runtime.cwd}/trained_models/pretraining_gripper_embnn_weights.ckpt
+    gripper: r-pad/taxpose/model-wajnzcwn:v0
+    # mug: ${hydra:runtime.cwd}/trained_models/pretraining_mug_embnn_weights.ckpt
+    mug: r-pad/taxpose/model-f81p2p5d:v0
+    # rack: ${hydra:runtime.cwd}/trained_models/pretraining_rack_embnn_weights.ckpt
+    rack: r-pad/taxpose/model-1on3wem2:v0
+    # bottle: ${hydra:runtime.cwd}/trained_models/pretraining_bottle_embnn_weights.ckpt
+    bottle: r-pad/taxpose/model-nogq98i9:v0
+    # bowl: ${hydra:runtime.cwd}/trained_models/pretraining_bowl_embnn_weights.ckpt
+    bowl: r-pad/taxpose/model-i7e35k1f:v0
+    # slab: ${hydra:runtime.cwd}/trained_models/pretraining_slab_embnn_weights.ckpt
+    slab: r-pad/taxpose/model-ysgwf55q:v0
+  none:
+    gripper: null
+    mug: null
+    rack: null
+    wine: null
+    wine_rack: null
+    bottle: null
+    bowl: null
+    slab: null
+    square: null
+    base: null
+    phone: null
+    phone_base: null
+    toilet_roll: null
+    toilet_roll_stand: null
+    puzzle_piece: null
+    puzzle: null
+    hanger: null
+    hanger_rack: null
+vn_dgcnn:
+  nce_loss:
+    # 0
+    # mug: /home/beisner/code/multi_project/taxpose/outputs/2023-05-22/17-09-00/embedding_pretraining_mug/2eakon6k/checkpoints/embnn_callback_epoch_2_global_step_33000.ckpt
+    # # 1
+    # rack: /home/beisner/code/multi_project/taxpose/outputs/2023-05-23/14-31-55/embedding_pretraining_mug/wcdwkh8r/checkpoints/embnn_callback_epoch_999_global_step_13000.ckpt
+    # # 2
+    # gripper: /home/beisner/code/multi_project/taxpose/outputs/2023-05-23/14-31-58/embedding_pretraining_mug/1m959icc/checkpoints/embnn_callback_epoch_999_global_step_13000.ckpt
+    gripper: r-pad/taxpose/model-3bvf82th:v0
+    mug: r-pad/taxpose/model-p3fluqtg:v0
+    rack: r-pad/taxpose/model-kxc51apt:v0
+    bottle: r-pad/taxpose/model-x3c82iam:v0
+    bowl: r-pad/taxpose/model-qbtkzw5y:v0
+    slab: r-pad/taxpose/model-tmymvff2:v0
+
+  none:
+    gripper: null
+    mug: null
+    rack: null
+    wine: null
+    wine_rack: null
+    bottle: null
+    bowl: null
+    slab: null
+    square: null
+    base: null
+    phone: null
+    phone_base: null
+    toilet_roll: null
+    toilet_roll_stand: null
+    puzzle_piece: null
+    puzzle: null
+    hanger: null
+    hanger_rack: null

configs/commands/ablations/se3_augmentation/_base.yaml

@@ -0,0 +1,14 @@
+# @package _global_
+
+defaults:
+  - /commands/ndf/_train@_here_
+  - _self_
+
+dm:
+  train_dset:
+    rotation_variance: 1e-4
+    translation_variance: 1e-4
+
+  val_dset:
+    rotation_variance: 3.1415
+    translation_variance: 0.5

configs/commands/ablations/se3_augmentation/eval_mlat_mug_arbitrary.yaml

@@ -0,0 +1,11 @@
+# @package _global_
+
+defaults:
+  - /commands/ndf/_eval@_here_
+  - override /model: mlat_s256_vnn
+  - override /task: mug
+  - override /pose_dist: arbitrary
+  - _self_
+
+checkpoint_file_grasp: r-pad/taxpose/model-sysvj0fh:v0
+checkpoint_file_place: r-pad/taxpose/model-wbjsiwz1:v0