Supports warp backend for camera unprojection operations (#2)

* implements unprojection operations in torch
* adds GPU backend support in play_camera.py
* adds play-camera to docs
* adds utilities methods to convert any array type to different backends
* adds utilities to convert arrays/tensors in a nested dictionary to a specified backend
* updates changelog and extension.toml
Co-authored-by: Mayank Mittal <mittalma@leggedrobotics.com>

docs/conf.py

@@ -96,6 +96,7 @@ autodoc_mock_imports = [
     "numpy",
     "scipy",
     "carb",
+    "warp",
     "pxr",
     "omni.kit",
     "omni.usd",

docs/source/setup/sample.rst

@@ -31,6 +31,15 @@ A few quick demo scripts to run and checkout:
 
       ./orbit.sh -p source/standalone/demo/play_ik_control.py --robot franka_panda --num_envs 128
 
+-  Spawn a camera and visualize the obtained pointcloud:
+
+   .. code:: bash
+
+      # CPU
+      ./orbit.sh -p source/standalone/demo/play_camera.py
+      # GPU
+      ./orbit.sh -p source/standalone/demo/play_camera.py --gpu
+
 Environments
 ------------
 

source/extensions/omni.isaac.orbit/config/extension.toml

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.1.1"
+version = "0.2.0"
 
 # Description
 title = "ORBIT framework for Robot Learning"

source/extensions/omni.isaac.orbit/docs/CHANGELOG.rst

 Changelog
 ---------
 
+0.2.0 (2023-01-25)
+~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added support for warp backend in camera utilities.
+* Extended the ``play_camera.py`` with ``--gpu`` flag to use GPU replicator backend.
+
 0.1.1 (2023-01-24)
 ~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.orbit/omni/isaac/orbit/sensors/height_scanner/utils.py

@@ -12,6 +12,8 @@ from matplotlib.axes import Axes
 from matplotlib.image import AxesImage
 from typing import Tuple
 
+__all__ = ["create_points_from_grid", "plot_height_grid"]
+
 
 def create_points_from_grid(size: Tuple[float, float], resolution: float) -> np.ndarray:
     """Creates a list of points from 2D meshgrid.

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/__init__.py

@@ -14,18 +14,24 @@ Sub-module containing utilities for the Orbit framework.
 * `timer`: Provides a timer class (uses contextlib) for benchmarking.
 """
 
+from .array import TENSOR_TYPE_CONVERSIONS, TENSOR_TYPES, convert_to_torch
 from .configclass import configclass
-from .dict import class_to_dict, print_dict, update_class_from_dict, update_dict
+from .dict import class_to_dict, convert_dict_to_backend, print_dict, update_class_from_dict, update_dict
 from .string import to_camel_case, to_snake_case
 from .timer import Timer
 
 __all__ = [
+    # arrays
+    "TENSOR_TYPES",
+    "TENSOR_TYPE_CONVERSIONS",
+    "convert_to_torch",
     # config wrapper
     "configclass",
     # dictionary utilities
     "class_to_dict",
-    "update_class_from_dict",
+    "convert_dict_to_backend",
     "print_dict",
+    "update_class_from_dict",
     "update_dict",
     # string utilities
     "to_camel_case",

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/array.py

+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES, ETH Zurich, and University of Toronto
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Utilities for working with different array backends."""
+
+import numpy as np
+import torch
+from typing import Optional, Sequence, Union
+
+import warp as wp
+
+__all__ = ["TENSOR_TYPES", "TENSOR_TYPE_CONVERSIONS", "convert_to_torch"]
+
+TENSOR_TYPES = {
+    "numpy": np.ndarray,
+    "torch": torch.Tensor,
+    "warp": wp.array,
+}
+"""A dictionary containing the types for each backend.
+
+The keys are the name of the backend ("numpy", "torch", "warp") and the values are the corresponding type
+(``np.ndarray``, ``torch.Tensor``, ``wp.array``).
+"""
+
+TENSOR_TYPE_CONVERSIONS = {
+    "numpy": {wp.array: lambda x: x.numpy(), torch.Tensor: lambda x: x.detach().cpu().numpy()},
+    "torch": {wp.array: lambda x: wp.torch.to_torch(x), np.ndarray: lambda x: torch.from_numpy(x)},
+    "warp": {np.array: lambda x: wp.array(x), torch.Tensor: lambda x: wp.torch.from_torch(x)},
+}
+"""A nested dictionary containing the conversion functions for each backend.
+
+The keys of the outer dictionary are the name of target backend ("numpy", "torch", "warp"). The keys of the
+inner dictionary are the source backend (``np.ndarray``, ``torch.Tensor``, ``wp.array``).
+"""
+
+
+def convert_to_torch(
+    array: Sequence[float],
+    dtype: torch.dtype = None,
+    device: Optional[Union[torch.device, str]] = None,
+) -> torch.Tensor:
+    """Converts a given array into a torch tensor.
+
+    The function tries to convert the array to a torch tensor. If the array is a numpy/warp arrays, or python
+    list/tuples, it is converted to a torch tensor. If the array is already a torch tensor, it is returned
+    directly.
+
+    If ``device`` is :obj:`None`, then the function deduces the current device of the data. For numpy arrays,
+    this defaults to "cpu", for torch tensors it is "cpu" or "cuda", and for warp arrays it is "cuda".
+
+    Args:
+        array (Sequence[float]): The input array. It can be a numpy array, warp array, python list/tuple, or torch tensor.
+        dtype (torch.dtype, optional): Target data-type for the tensor.
+        device (Optional[Union[torch.device, str]], optional): The target device for the tensor. Defaults to None.
+
+    Returns:
+        torch.Tensor: The converted array as torch tensor.
+    """
+    # Convert array to tensor
+    if isinstance(array, torch.Tensor):
+        tensor = array
+    elif isinstance(array, np.ndarray):
+        tensor = torch.from_numpy(array)
+    elif isinstance(array, wp.array):
+        tensor = wp.to_torch(array)
+    else:
+        tensor = torch.Tensor(array)
+    # Convert tensor to the right device
+    if device is not None and str(tensor.device) != str(device):
+        tensor = tensor.to(device)
+    # Convert dtype of tensor if requested
+    if dtype is not None and tensor.dtype != dtype:
+        tensor = tensor.type(dtype)
+
+    return tensor

source/extensions/omni.isaac.orbit/omni/isaac/orbit/utils/dict.py

@@ -10,7 +10,9 @@ import collections.abc
 import importlib
 from typing import Any, Callable, Dict, Iterable, Mapping
 
-__all__ = ["class_to_dict", "update_class_from_dict", "update_dict", "print_dict"]
+from .array import TENSOR_TYPE_CONVERSIONS, TENSOR_TYPES
+
+__all__ = ["class_to_dict", "update_class_from_dict", "convert_dict_to_backend", "update_dict", "print_dict"]
 
 """
 Dictionary <-> Class operations.
@@ -74,7 +76,9 @@ def update_class_from_dict(obj, data: Dict[str, Any], _ns: str = "") -> None:
         KeyError: When dictionary has a key that does not exist in the default config type.
     """
     for key, value in data.items():
+        # key_ns is the full namespace of the key
         key_ns = _ns + "/" + key
+        # check if key is present in the object
         if hasattr(obj, key):
             obj_mem = getattr(obj, key)
             if isinstance(obj_mem, Mapping):
@@ -115,6 +119,75 @@ Dictionary operations.
 """
 
 
+def convert_dict_to_backend(
+    data: dict, backend: str = "numpy", array_types: Iterable[str] = ("numpy", "torch", "warp")
+) -> dict:
+    """Convert all arrays or tensors in a dictionary to a given backend.
+
+    This function iterates over the dictionary, converts all arrays or tensors with the given types to
+    the desired backend, and stores them in a new dictionary. It also works with nested dictionaries.
+
+    Currently supported backends are "numpy", "torch", and "warp".
+
+    Note:
+        This function only converts arrays or tensors. Other types of data are left unchanged. Mutable types
+        (e.g. lists) are referenced by the new dictionary, so they are not copied.
+
+    Args:
+        data (dict): An input dict containing array or tensor data as values.
+        backend(str): The backend ("numpy", "torch", "warp") to which arrays in this dict should be converted.
+            Defaults to "numpy".
+        array_types(Iterable[str]): A list containing the types of arrays that should be converted to
+            the desired backend. Defaults to ("numpy", "torch", "warp").
+
+    Raises:
+        ValueError: If the specified ``backend`` or ``array_types`` are unknown, i.e. not in the list of supported
+            backends ("numpy", "torch", "warp").
+
+    Returns:
+        dict: The updated dict with the data converted to the desired backend.
+    """
+    # THINK: Should we also support converting to a specific device, e.g. "cuda:0"?
+    # Define the conversion functions for each backend.
+    if backend not in TENSOR_TYPE_CONVERSIONS:
+        raise ValueError(f"Unknown backend '{backend}'. Supported backends are 'numpy', 'torch', and 'warp'.")
+    else:
+        tensor_type_conversions = TENSOR_TYPE_CONVERSIONS[backend]
+
+    # Parse the array types and convert them to the corresponding types: "numpy" -> np.ndarray, etc.
+    parsed_types = list()
+    for t in array_types:
+        # Check type is valid.
+        if t not in TENSOR_TYPES:
+            raise ValueError(f"Unknown array type: '{t}'. Supported array types are 'numpy', 'torch', and 'warp'.")
+        # Exclude types that match the backend, since we do not need to convert these.
+        if t == backend:
+            continue
+        # Convert the string types to the corresponding types.
+        parsed_types.append(TENSOR_TYPES[t])
+
+    # Convert the data to the desired backend.
+    output_dict = dict()
+    for key, value in data.items():
+        # Obtain the data type of the current value.
+        data_type = type(value)
+        # -- arrays
+        if data_type in parsed_types:
+            # check if we have a known conversion.
+            if data_type not in tensor_type_conversions:
+                raise ValueError(f"No registered conversion for data type: {data_type} to {backend}!")
+            else:
+                output_dict[key] = tensor_type_conversions[data_type](value)
+        # -- nested dictionaries
+        elif isinstance(data[key], dict):
+            output_dict[key] = convert_dict_to_backend(value)
+        # -- everything else
+        else:
+            output_dict[key] = value
+
+    return output_dict
+
+
 def update_dict(orig_dict: dict, new_dict: collections.abc.Mapping) -> dict:
     """Updates existing dictionary with values from a new dictionary.
 

source/standalone/demo/play_camera.py

@@ -2,7 +2,6 @@
 # All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
-
 """
 This script shows how to use the camera sensor from the Orbit framework.
 
@@ -12,14 +11,15 @@ the simulator or OpenGL convention for the camera, we use the robotics or ROS co
 
 """Launch Isaac Sim Simulator first."""
 
-
 import argparse
 
+# omni-isaac-orbit
 from omni.isaac.kit import SimulationApp
 
 # add argparse arguments
 parser = argparse.ArgumentParser("Welcome to Orbit: Omniverse Robotics Environments!")
 parser.add_argument("--headless", action="store_true", default=False, help="Force display off at all times.")
+parser.add_argument("--gpu", action="store_true", default=False, help="Use gpu for pointcloud unprojection.")
 args_cli = parser.parse_args()
 
 # launch omniverse app
@@ -45,6 +45,7 @@ from pxr import Gf, UsdGeom
 import omni.isaac.orbit.utils.kit as kit_utils
 from omni.isaac.orbit.sensors.camera import Camera, PinholeCameraCfg
 from omni.isaac.orbit.sensors.camera.utils import create_pointcloud_from_rgbd
+from omni.isaac.orbit.utils import convert_dict_to_backend
 
 """
 Helpers
@@ -69,6 +70,8 @@ def design_scene():
         translation=(-4.5, 3.5, 10.0),
         attributes={"radius": 2.5, "intensity": 600.0, "color": (1.0, 1.0, 1.0)},
     )
+    # Xform to hold objects
+    prim_utils.create_prim("/World/Objects", "Xform")
     # Random objects
     for i in range(8):
         # sample random position
@@ -101,6 +104,7 @@ Main
 def main():
     """Runs a camera sensor from orbit."""
 
+    device = "cuda" if args_cli.gpu else "cpu"
     # Load kit helper
     sim = SimulationContext(stage_units_in_meters=1.0, physics_dt=0.005, rendering_dt=0.005, backend="torch")
     # Set main camera
@@ -118,8 +122,12 @@ def main():
         data_types=["rgb", "distance_to_image_plane", "normals", "motion_vectors"],
         usd_params=PinholeCameraCfg.UsdCameraCfg(clipping_range=(0.1, 1.0e5)),
     )
-    camera = Camera(cfg=camera_cfg, device="cpu")
+    camera = Camera(cfg=camera_cfg, device=device)
+
+    # Spawn camera
     camera.spawn("/World/CameraSensor")
+    # Initialize camera
+    # note: For rendering based sensors, it is not necessary to initialize before playing the simulation.
     camera.initialize()
 
     # Create replicator writer
@@ -144,13 +152,16 @@ def main():
         sim.step()
         # Update camera data
         camera.update(dt=0.0)
+
+        # Print camera info
         print(camera)
         print("Received shape of rgb   image: ", camera.data.output["rgb"].shape)
         print("Received shape of depth image: ", camera.data.output["distance_to_image_plane"].shape)
         print("-------------------------------")
 
         # Save images
-        rep_writer.write(camera.data.output)
+        # note: BasicWriter only supports saving data in numpy format
+        rep_writer.write(convert_dict_to_backend(camera.data.output, backend="numpy"))
 
         # Pointcloud in world frame
         pointcloud_w, pointcloud_rgb = create_pointcloud_from_rgbd(
@@ -162,7 +173,13 @@ def main():
             normalize_rgb=True,
             num_channels=4,
         )
-        # visualize the points
+
+        # Convert to numpy for visualization
+        if not isinstance(pointcloud_w, np.ndarray):
+            pointcloud_w = pointcloud_w.cpu().numpy()
+        if not isinstance(pointcloud_rgb, np.ndarray):
+            pointcloud_rgb = pointcloud_rgb.cpu().numpy()
+        # Visualize the points
         num_points = pointcloud_w.shape[0]
         points_size = [1.25] * num_points
         points_color = pointcloud_rgb