Merge branch 'og-develop' into feat/curobo-primitives

.github/workflows/build-push-containers.yml

@@ -22,7 +22,7 @@ jobs:
           sudo rm -rf \
             /usr/share/dotnet /usr/local/lib/android /opt/ghc \
             /usr/local/share/powershell /usr/share/swift /usr/local/.ghcup \
-            /usr/lib/jvm || true
+            /usr/lib/jvm /opt/hostedtoolcache/CodeQL || true
           echo "some directories deleted"
           sudo apt install aptitude -y >/dev/null 2>&1
           sudo aptitude purge aria2 ansible azure-cli shellcheck rpm xorriso zsync \
@@ -97,6 +97,18 @@ jobs:
           tags: |
             type=ref,event=branch
             type=semver,pattern={{version}}
+      -
+        name: Metadata for actions Image
+        id: meta-actions
+        uses: docker/metadata-action@v5
+        # The actions image should only be built if the push is to og-develop
+        if: github.ref == 'refs/heads/og-develop'
+        with:
+          images: |
+            stanfordvl/omnigibson-gha
+          tags: |
+            # We only push to the latest tag for the actions image
+            type=raw,value=latest
       -
         name: Build and push prod image
         id: build-prod
@@ -107,8 +119,8 @@ jobs:
           tags: ${{ steps.meta-prod.outputs.tags }}
           labels: ${{ steps.meta-prod.outputs.labels }}
           file: docker/prod.Dockerfile
-          cache-from: type=registry,ref=stanfordvl/omnigibson:og-develop
-          cache-to: type=inline
+          cache-from: type=registry,ref=stanfordvl/omnigibson:build-cache
+          cache-to: type=registry,ref=stanfordvl/omnigibson:build-cache,mode=max
 
       -
         name: Build and push dev image
@@ -121,8 +133,8 @@ jobs:
           tags: ${{ steps.meta-dev.outputs.tags }}
           labels: ${{ steps.meta-dev.outputs.labels }}
           file: docker/prod.Dockerfile
-          cache-from: type=registry,ref=stanfordvl/omnigibson:og-develop  # OK to share cache here.
-          cache-to: type=inline
+          cache-from: type=registry,ref=stanfordvl/omnigibson:build-cache  # OK to share cache here.
+          cache-to: type=registry,ref=stanfordvl/omnigibson:build-cache,mode=max
 
       - name: Update vscode image Dockerfile with prod image tag
         run: |
@@ -137,5 +149,25 @@ jobs:
           tags: ${{ steps.meta-vscode.outputs.tags }}
           labels: ${{ steps.meta-vscode.outputs.labels }}
           file: docker/vscode.Dockerfile
-          cache-from: type=registry,ref=stanfordvl/omnigibson:og-develop  # OK to share cache here.
-          cache-to: type=inline
+          cache-from: type=registry,ref=stanfordvl/omnigibson:build-cache  # OK to share cache here.
+          cache-to: type=registry,ref=stanfordvl/omnigibson:build-cache,mode=max
+
+      - name: Update actions image Dockerfile with dev image tag
+        # The actions image should only be built if the push is to og-develop
+        if: github.ref == 'refs/heads/og-develop'
+        run: |
+          sed -i "s/omnigibson-dev:og-develop/omnigibson-dev@${{ steps.build-dev.outputs.digest }}/g" docker/gh-actions/Dockerfile && cat docker/gh-actions/Dockerfile
+      -
+        name: Build and push actions image
+        id: build-actions
+        uses: docker/build-push-action@v5
+        # The actions image should only be built if the push is to og-develop
+        if: github.ref == 'refs/heads/og-develop'
+        with:
+          context: docker/gh-actions
+          push: true
+          tags: ${{ steps.meta-actions.outputs.tags }}
+          labels: ${{ steps.meta-actions.outputs.labels }}
+          file: docker/gh-actions/Dockerfile
+          cache-from: type=registry,ref=stanfordvl/omnigibson:build-cache  # OK to share cache here.
+          cache-to: type=registry,ref=stanfordvl/omnigibson:build-cache,mode=max

.github/workflows/examples-as-test.yml

@@ -28,7 +28,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Generate example tests
         working-directory: omnigibson-src
@@ -76,7 +76,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Run tests
         working-directory: omnigibson-src

.github/workflows/profiling.yml

@@ -34,7 +34,7 @@ jobs:
 
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev]
+        run: pip install -e .[dev,primitives]
 
       - name: Run performance benchmark
         run: bash scripts/profiling.sh

.github/workflows/tests.yml

@@ -50,12 +50,9 @@ jobs:
           submodules: true
           path: omnigibson-src
 
-      - name: Install CUDA and PyTorch
-        run: micromamba install cuda=11.8 -c conda-forge
-
       - name: Install
         working-directory: omnigibson-src
-        run: pip install -e .[dev] --no-build-isolation
+        run: pip install -e .[dev,primitives] --no-build-isolation
 
       - name: Print env
         run: printenv

docker/prod.Dockerfile

@@ -7,6 +7,8 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \
   && rm -rf /var/lib/apt/lists/*
 
 RUN rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/gym*
+RUN rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/torch*
+RUN rm -rf /isaac-sim/exts/omni.isaac.ml_archive/pip_prebundle/functorch*
 RUN rm -rf /isaac-sim/kit/extscore/omni.kit.pip_archive/pip_prebundle/numpy*
 RUN /isaac-sim/python.sh -m pip install click~=8.1.3
 
@@ -23,6 +25,29 @@ ENV MAMBA_ROOT_PREFIX /micromamba
 RUN micromamba create -n omnigibson -c conda-forge python=3.10
 RUN micromamba shell init --shell=bash
 
+# Install torch
+RUN micromamba run -n omnigibson micromamba install \
+  pytorch torchvision pytorch-cuda=11.8 \
+  -c pytorch -c nvidia -c conda-forge
+
+# Install curobo. This can normally be installed when OmniGibson is pip
+# installed, but we need to install it beforehand here so that it doesn't
+# have to happen on every time a CI action is run (otherwise it's just
+# very slow).
+# This also allows us to uninstall the cuda toolkit after curobo is built
+# to save space (meaning curobo will not be able to be rebuilt at runtime).
+# Here we also compile this such that it is compatible with GPU architectures
+# Turing, Ampere, and Ada; which correspond to 20, 30, and 40 series GPUs.
+# We also suppress the output of the installation to avoid the log limit.
+RUN wget --no-verbose -O /cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \ 
+  dpkg -i /cuda-keyring.deb && rm /cuda-keyring.deb && apt-get update && \
+  DEBIAN_FRONTEND=noninteractive apt-get install -y cuda-toolkit-11-8 && \
+  TORCH_CUDA_ARCH_LIST='7.5;8.0;8.6+PTX' PATH=/usr/local/cuda-11.8/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH \
+    micromamba run -n omnigibson pip install \
+    git+https://github.com/StanfordVL/curobo@6a4eb2ca8677829b0f57451ad107e0a3186525e9#egg=nvidia_curobo \
+    --no-build-isolation > /dev/null && \
+  apt-get remove -y cuda-toolkit-11-8 && apt-get autoremove -y && apt-get autoclean -y && rm -rf /var/lib/apt/lists/*
+
 # Make sure isaac gets properly sourced every time omnigibson gets called
 ARG CONDA_ACT_FILE="/micromamba/envs/omnigibson/etc/conda/activate.d/env_vars.sh"
 RUN mkdir -p "/micromamba/envs/omnigibson/etc/conda/activate.d"
@@ -33,31 +58,6 @@ RUN echo "source /isaac-sim/setup_conda_env.sh" >> $CONDA_ACT_FILE
 
 RUN echo "micromamba activate omnigibson" >> /root/.bashrc
 
-# Prepare to build OMPL
-ENV CXX="g++"
-ENV MAKEFLAGS="-j `nproc`"
-RUN micromamba run -n omnigibson micromamba install -c conda-forge boost && \
-    micromamba run -n omnigibson pip install pyplusplus && \
-    git clone https://github.com/ompl/ompl.git /ompl && \
-    mkdir -p /ompl/build/Release && \
-    sed -i "s/find_program(PYPY/# find_program(PYPY/g" /ompl/CMakeModules/Findpypy.cmake
-
-# Build and install OMPL 
-RUN micromamba run -n omnigibson /bin/bash --login -c 'source /isaac-sim/setup_conda_env.sh && (which python > /root/PYTHON_EXEC) && (echo $PYTHONPATH > /root/PYTHONPATH)' && \
-    cd /ompl/build/Release && \
-    micromamba run -n omnigibson cmake ../.. \
-      -DCMAKE_INSTALL_PREFIX="$CONDA_PREFIX" \
-      -DBOOST_ROOT="$CONDA_PREFIX" \
-      -DPYTHON_EXEC=$(cat /root/PYTHON_EXEC) \
-      -DPYTHONPATH=$(cat /root/PYTHONPATH) && \
-    micromamba run -n omnigibson make -j 4 update_bindings && \
-    micromamba run -n omnigibson make -j 4 && \
-    cd py-bindings && \
-    micromamba run -n omnigibson make install
-
-# Test OMPL
-RUN micromamba run -n omnigibson python -c "from ompl import base"
-
 # Copy over omnigibson source
 ADD . /omnigibson-src
 WORKDIR /omnigibson-src
@@ -72,7 +72,7 @@ ENV DEV_MODE=${DEV_MODE}
 ARG WORKDIR_PATH=/omnigibson-src
 RUN if [ "$DEV_MODE" != "1" ]; then \
       echo "OMNIGIBSON_NO_OMNIVERSE=1 python omnigibson/download_datasets.py" >> /root/.bashrc; \
-      micromamba run -n omnigibson pip install -e .[dev]; \
+      micromamba run -n omnigibson pip install -e .[dev,primitives]; \
     else \
       WORKDIR_PATH=/; \
       cd / && rm -rf /omnigibson-src; \

omnigibson/configs/tiago_primitives.yaml

@@ -51,13 +51,13 @@ robots:
         name: JointController
       arm_left:
         name: JointController
+        subsume_controllers: [trunk]
         motor_type: position
         command_input_limits: null
         command_output_limits: null
         use_delta_commands: false
       arm_right:
         name: JointController
-        subsume_controllers: [trunk]
         motor_type: position
         command_input_limits: null
         command_output_limits: null

omnigibson/controllers/ik_controller.py

@@ -175,8 +175,9 @@ def limiter(target_pos: Array[float], target_quat: Array[float], control_dict: D
             control_freq=control_freq,
             control_limits=control_limits,
             dof_idx=dof_idx,
-            kp=kp,
-            damping_ratio=damping_ratio,
+            pos_kp=pos_kp,
+            pos_damping_ratio=pos_damping_ratio,
+            vel_kp=vel_kp,
             motor_type="position",
             use_delta_commands=False,
             use_impedances=use_impedances,

omnigibson/controllers/null_joint_controller.py

@@ -68,8 +68,9 @@ def __init__(
             dof_idx=dof_idx,
             command_input_limits=command_input_limits,
             command_output_limits=command_output_limits,
-            kp=kp,
-            damping_ratio=damping_ratio,
+            pos_kp=pos_kp,
+            pos_damping_ratio=pos_damping_ratio,
+            vel_kp=vel_kp,
             use_impedances=use_impedances,
             use_delta_commands=False,
         )

omnigibson/objects/controllable_object.py

@@ -226,6 +226,9 @@ def _load_controllers(self):
         # Generate the controller config
         self._controller_config = self._generate_controller_config(custom_config=self._controller_config)
 
+        # We copy the controller config here because we add/remove some keys in-place that shouldn't persist
+        _controller_config = deepcopy(self._controller_config)
+
         # Store dof idx mapping to dof name
         self.dof_names_ordered = list(self._joints.keys())
 
@@ -237,8 +240,8 @@ def _load_controllers(self):
         subsume_names = set()
         for name in self._raw_controller_order:
             # Make sure we have the valid controller name specified
-            assert_valid_key(key=name, valid_keys=self._controller_config, name="controller name")
-            cfg = self._controller_config[name]
+            assert_valid_key(key=name, valid_keys=_controller_config, name="controller name")
+            cfg = _controller_config[name]
             subsume_controllers = cfg.pop("subsume_controllers", [])
             # If this controller subsumes other controllers, it cannot be subsumed by another controller
             # (i.e.: we don't allow nested / cyclical subsuming)
@@ -262,11 +265,11 @@ def _load_controllers(self):
             # If this controller is subsumed by another controller, simply skip it
             if name in subsume_names:
                 continue
-            cfg = self._controller_config[name]
+            cfg = _controller_config[name]
             # If we subsume other controllers, prepend the subsumed' dof idxs to this controller's idxs
             if name in controller_subsumes:
                 for subsumed_name in controller_subsumes[name]:
-                    subsumed_cfg = self._controller_config[subsumed_name]
+                    subsumed_cfg = _controller_config[subsumed_name]
                     cfg["dof_idx"] = th.concatenate([subsumed_cfg["dof_idx"], cfg["dof_idx"]])
 
             # If we're using normalized action space, override the inputs for all controllers

omnigibson/utils/python_utils.py

@@ -151,9 +151,9 @@ def create_object_from_init_info(init_info):
 
 def safe_equal(a, b):
     if isinstance(a, th.Tensor) and isinstance(b, th.Tensor):
-        return (a == b).all().item()
+        return a.shape == b.shape and (a == b).all().item()
     elif isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)):
-        return all(safe_equal(a_item, b_item) for a_item, b_item in zip(a, b))
+        return len(a) == len(b) and all(safe_equal(a_item, b_item) for a_item, b_item in zip(a, b))
     else:
         return a == b
 

omnigibson/utils/transform_utils.py

@@ -273,8 +273,8 @@ def quat_slerp(quat0, quat1, frac, shortestpath=True, eps=1.0e-15):
     # type: (Tensor, Tensor, Tensor, bool, float) -> Tensor
     # reshape quaternion
     quat_shape = quat0.shape
-    quat0 = unit_vector(quat0.reshape(-1, 4), dim=-1)
-    quat1 = unit_vector(quat1.reshape(-1, 4), dim=-1)
+    quat0 = unit_vector(quat0.reshape(-1, 4), dim=-1, out=None)
+    quat1 = unit_vector(quat1.reshape(-1, 4), dim=-1, out=None)
 
     # Check for endpoint cases
     where_start = frac <= 0.0
@@ -481,8 +481,8 @@ def vec2quat(vec: torch.Tensor, up: torch.Tensor = torch.tensor([0.0, 0.0, 1.0])
     if up.dim() == 1:
         up = up.unsqueeze(0)
 
-    vec_n = torch.nn.functional.normalize(vec, dim=-1)
-    up_n = torch.nn.functional.normalize(up, dim=-1)
+    vec_n = normalize(vec, dim=-1, eps=1e-10)
+    up_n = normalize(up, dim=-1, eps=1e-10)
 
     s_n = torch.cross(up_n, vec_n, dim=-1)
     u_n = torch.cross(vec_n, s_n, dim=-1)
@@ -1141,8 +1141,8 @@ def vecs2axisangle(vec0, vec1):
         vec1 (torch.tensor): (..., 3) (x,y,z) 3D vector, possibly unnormalized
     """
     # Normalize vectors
-    vec0 = normalize(vec0, dim=-1)
-    vec1 = normalize(vec1, dim=-1)
+    vec0 = normalize(vec0, dim=-1, eps=1e-10)
+    vec1 = normalize(vec1, dim=-1, eps=1e-10)
 
     # Get cross product for direction of angle, and multiply by arcos of the dot product which is the angle
     return torch.linalg.cross(vec0, vec1) * torch.arccos((vec0 * vec1).sum(-1, keepdim=True))
@@ -1162,8 +1162,8 @@ def vecs2quat(vec0: torch.Tensor, vec1: torch.Tensor, normalized: bool = False)
     """
     # Normalize vectors if requested
     if not normalized:
-        vec0 = normalize(vec0, dim=-1)
-        vec1 = normalize(vec1, dim=-1)
+        vec0 = normalize(vec0, dim=-1, eps=1e-10)
+        vec1 = normalize(vec1, dim=-1, eps=1e-10)
 
     # Half-way Quaternion Solution -- see https://stackoverflow.com/a/11741520
     cos_theta = torch.sum(vec0 * vec1, dim=-1, keepdim=True)

setup.py

@@ -47,7 +47,6 @@
         "rtree>=1.2.0",
         "graphviz>=0.20",
         "matplotlib>=3.0.0",
-        "nvidia-curobo @ git+https://github.com/StanfordVL/curobo@main",
     ],
     extras_require={
         "dev": [
@@ -63,7 +62,11 @@
             "mkdocs-section-index",
             "mkdocs-literate-nav",
             "telemoma~=0.1.2",
-        ]
+        ],
+        "primitives": [
+            "nvidia-curobo @ git+https://github.com/StanfordVL/curobo@6a4eb2ca8677829b0f57451ad107e0a3186525e9",
+            "ompl @ https://storage.googleapis.com/gibson_scenes/ompl-1.6.0-cp310-cp310-manylinux_2_28_x86_64.whl",
+        ],
     },
     tests_require=[],
     python_requires=">=3",

tests/test_curobo.py

@@ -228,8 +228,8 @@ def test_curobo():
         env.scene.reset()
 
         # Create CuRobo instance
-        batch_size = 10
-        n_samples = 30
+        batch_size = 2
+        n_samples = 20
 
         cmg = CuRoboMotionGenerator(
             robot=robot,
@@ -386,7 +386,8 @@ def test_curobo():
             [0.01 if joint.joint_type == "PrismaticJoint" else 3.0 / 180.0 * math.pi for joint in robot.joints.values()]
         )
 
-        for bypass_physics in [True, False]:
+        # for bypass_physics in [True, False]:
+        for bypass_physics in [True]:
             for traj_idx, (success, traj_path) in enumerate(zip(successes, traj_paths)):
                 if not success:
                     continue

tests/test_robot_states_flatcache.py

@@ -61,7 +61,7 @@ def camera_pose_test(flatcache):
         relative_pose_transform(sensor_world_pos, sensor_world_ori, robot_world_pos, robot_world_ori)
     )
 
-    sensor_world_pos_gt = th.tensor([150.1620, 149.9999, 101.2193])
+    sensor_world_pos_gt = th.tensor([150.1628, 149.9993, 101.3773])
     sensor_world_ori_gt = th.tensor([-0.2952, 0.2959, 0.6427, -0.6421])
 
     assert th.allclose(sensor_world_pos, sensor_world_pos_gt, atol=1e-3)