Create Manager Based Cartpole Vision Example Environments

CONTRIBUTORS.md

@@ -40,6 +40,7 @@ Guidelines for modifications:
 * Calvin Yu
 * Chenyu Yang
 * David Yang
+* Gary Lvov
 * HoJin Jeon
 * Jean Tampon
 * Jia Lin Yuan

docs/source/overview/environments.rst

@@ -53,8 +53,10 @@ Classic environments that are based on IsaacGymEnvs implementation of MuJoCo-sty
     |                  |                             |                                                                         |
     |                  | |cartpole-direct-link|      |                                                                         |
     +------------------+-----------------------------+-------------------------------------------------------------------------+
-    | |cartpole|       | |cartpole-camera-rgb-link|  | Move the cart to keep the pole upwards in the classic cartpole control  |
+    | |cartpole|       | |manager-camera-rgb-link|   | Move the cart to keep the pole upwards in the classic cartpole control  |
     |                  |                             | and perceptive inputs                                                   |
+    |                  | |manager-camera-dpt-link|   |                                                                         |
+    |                  | |cartpole-camera-rgb-link|  |                                                                         |
     |                  | |cartpole-camera-dpt-link|  |                                                                         |
     +------------------+-----------------------------+-------------------------------------------------------------------------+
 
@@ -68,11 +70,14 @@ Classic environments that are based on IsaacGymEnvs implementation of MuJoCo-sty
 
 .. |humanoid-direct-link| replace:: `Isaac-Humanoid-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/humanoid/humanoid_env.py>`__
 .. |ant-direct-link| replace:: `Isaac-Ant-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/ant/ant_env.py>`__
+.. |manager-camera-rgb-link| replace:: `Isaac-Cartpole-RGB-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py>`__
+.. |manager-camera-dpt-link| replace:: `Isaac-Cartpole-Depth-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py>`__
 .. |cartpole-direct-link| replace:: `Isaac-Cartpole-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_env.py>`__
+.. |manager-camera-rgb-link| replace:: `Isaac-Cartpole-RGB-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py>`__
+.. |manager-camera-dpt-link| replace:: `Isaac-Cartpole-Depth-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_camera_env_cfg.py>`__
 .. |cartpole-camera-rgb-link| replace:: `Isaac-Cartpole-RGB-Camera-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py>`__
 .. |cartpole-camera-dpt-link| replace:: `Isaac-Cartpole-Depth-Camera-Direct-v0 <https://github.com/isaac-sim/IsaacLab/blob/main/source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/direct/cartpole/cartpole_camera_env.py>`__
 
-
 Manipulation
 ~~~~~~~~~~~~
 

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

@@ -1,6 +1,25 @@
 Changelog
 ---------
 
+0.22.15 (2024-09-20)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added :meth:`grab_images` to be able to use images for an observation term in manager based environments
+
+0.24.14 (2024-09-20)
+~~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added :meth:`convert_perspective_depth_to_orthogonal_depth`. :meth:`unproject_depth` assumes
+that the input depth image is orthogonal. The new :meth:`convert_perspective_depth_to_orthogonal_depth`
+can be used to convert a perspective depth image into an orthogonal depth image, so that the point cloud
+can be unprojected correctly with :meth:`unproject_depth`.
+
 
 0.24.13 (2024-09-08)
 ~~~~~~~~~~~~~~~~~~~~

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py

@@ -17,7 +17,7 @@
 import omni.isaac.lab.utils.math as math_utils
 from omni.isaac.lab.assets import Articulation, RigidObject
 from omni.isaac.lab.managers import SceneEntityCfg
-from omni.isaac.lab.sensors import RayCaster
+from omni.isaac.lab.sensors import Camera, RayCaster, RayCasterCamera, TiledCamera
 
 if TYPE_CHECKING:
     from omni.isaac.lab.envs import ManagerBasedEnv, ManagerBasedRLEnv
@@ -182,6 +182,41 @@ def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> tor
     return link_incoming_forces.view(env.num_envs, -1)
 
 
+def grab_images(
+    env: ManagerBasedEnv,
+    sensor_cfg: SceneEntityCfg = SceneEntityCfg("tiled_camera"),
+    data_type: str = "rgb",
+    convert_perspective_to_orthogonal: bool = True,
+    normalize: bool = True,
+) -> torch.Tensor:
+    """Grab all of the images of a specific datatype produced by a specific camera at the last timestep.
+
+    Args:
+        env: The environment the cameras are placed within.
+        sensor_cfg: The desired sensor to take . Defaults to SceneEntityCfg("tiled_camera").
+        data_type: The data type to pull from the desired camera. Defaults to "rgb".
+        convert_perspective_to_orthogonal: Whether to convert perspective
+            depth images to orthogonal depth images. Defaults to True.
+        normalize: Set to true normalize images. Highly recommended for training performance.
+
+    Returns:
+        The images produced at the last timestep
+    """
+    sensor: TiledCamera | Camera | RayCasterCamera = env.scene.sensors[sensor_cfg.name]
+    images = sensor.data.output[data_type]
+    if (data_type == "distance_to_camera") and convert_perspective_to_orthogonal:
+        images = math_utils.convert_perspective_depth_to_orthogonal_depth(images, sensor.data.intrinsic_matrices)
+
+    if normalize:
+        if data_type == "rgb":
+            images = images / 255
+            mean_tensor = torch.mean(images, dim=(1, 2), keepdim=True)
+            images -= mean_tensor
+        elif "distance_to" in data_type or "depth" in data_type:
+            images[images == float("inf")] = 0
+    return images.clone()
+
+
 """
 Actions.
 """

source/extensions/omni.isaac.lab/omni/isaac/lab/utils/math.py

@@ -988,7 +988,12 @@ def transform_points(
 
 @torch.jit.script
 def unproject_depth(depth: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
-    r"""Unproject depth image into a pointcloud.
+    r"""Unproject depth image into a pointcloud. This method assumes that depth
+    is provided orthagonally relative to the image plane, as opposed to absolutely relative to the camera's
+    principal point (perspective depth). To unproject a perspective depth image, use
+    :meth:`convert_perspective_depth_to_orthogonal_depth` to convert
+    to an orthogonal depth image prior to calling this method as otherwise the
+    created point cloud will be distorted, especially around the edges.
 
     This function converts depth images into points given the calibration matrix of the camera.
 
@@ -1059,6 +1064,105 @@ def unproject_depth(depth: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tens
     return points_xyz
 
 
+@torch.jit.script
+def convert_perspective_depth_to_orthogonal_depth(
+    perspective_depth: torch.Tensor, intrinsics: torch.Tensor
+) -> torch.Tensor:
+    r"""Provided depth image(s) where depth is provided as the distance to the principal
+    point of the camera (perspective depth), this function converts it so that depth
+    is provided as the distance to the camera's image plane (orthogonal depth).
+
+    This is helpful because `unproject_depth` assumes that depth is expressed in
+    the orthogonal depth format.
+
+    If `perspective_depth` is a batch of depth images and `intrinsics` is a single intrinsic matrix,
+    the same calibration matrix is applied to all depth images in the batch.
+
+    The function assumes that the width and height are both greater than 1.
+
+    Args:
+        perspective_depth: The depth measurement obtained with the distance_to_camera replicator.
+            Shape is (H, W) or or (H, W, 1) or (N, H, W) or (N, H, W, 1).
+        intrinsics: A tensor providing camera's calibration matrix. Shape is (3, 3) or (N, 3, 3).
+
+    Returns:
+        The depth image as if obtained by the distance_to_image_plane replicator. Shape
+            matches the input shape of depth
+
+    Raises:
+        ValueError: When depth is not of shape (H, W) or (H, W, 1) or (N, H, W) or (N, H, W, 1).
+        ValueError: When intrinsics is not of shape (3, 3) or (N, 3, 3).
+    """
+
+    # Clone inputs to avoid in-place modifications
+    perspective_depth_batch = perspective_depth.clone()
+    intrinsics_batch = intrinsics.clone()
+
+    # Check if inputs are batched
+    is_batched = perspective_depth_batch.dim() == 4 or (
+        perspective_depth_batch.dim() == 3 and perspective_depth_batch.shape[-1] != 1
+    )
+
+    # Track whether the last dimension was singleton
+    add_last_dim = False
+    if perspective_depth_batch.dim() == 4 and perspective_depth_batch.shape[-1] == 1:
+        add_last_dim = True
+        perspective_depth_batch = perspective_depth_batch.squeeze(dim=3)  # (N, H, W, 1) -> (N, H, W)
+    if perspective_depth_batch.dim() == 3 and perspective_depth_batch.shape[-1] == 1:
+        add_last_dim = True
+        perspective_depth_batch = perspective_depth_batch.squeeze(dim=2)  # (H, W, 1) -> (H, W)
+
+    if perspective_depth_batch.dim() == 2:
+        perspective_depth_batch = perspective_depth_batch[None]  # (H, W) -> (1, H, W)
+
+    if intrinsics_batch.dim() == 2:
+        intrinsics_batch = intrinsics_batch[None]  # (3, 3) -> (1, 3, 3)
+
+    if is_batched and intrinsics_batch.shape[0] == 1:
+        intrinsics_batch = intrinsics_batch.expand(perspective_depth_batch.shape[0], -1, -1)  # (1, 3, 3) -> (N, 3, 3)
+
+    # Validate input shapes
+    if perspective_depth_batch.dim() != 3:
+        raise ValueError(f"Expected perspective_depth to have 2, 3, or 4 dimensions; got {perspective_depth.shape}.")
+    if intrinsics_batch.dim() != 3:
+        raise ValueError(f"Expected intrinsics to have shape (3, 3) or (N, 3, 3); got {intrinsics.shape}.")
+
+    # Image dimensions
+    im_height, im_width = perspective_depth_batch.shape[1:]
+
+    # Get the intrinsics parameters
+    fx = intrinsics_batch[:, 0, 0].view(-1, 1, 1)
+    fy = intrinsics_batch[:, 1, 1].view(-1, 1, 1)
+    cx = intrinsics_batch[:, 0, 2].view(-1, 1, 1)
+    cy = intrinsics_batch[:, 1, 2].view(-1, 1, 1)
+
+    # Create meshgrid of pixel coordinates
+    u_grid = torch.arange(im_width, device=perspective_depth.device, dtype=perspective_depth.dtype)
+    v_grid = torch.arange(im_height, device=perspective_depth.device, dtype=perspective_depth.dtype)
+    u_grid, v_grid = torch.meshgrid(u_grid, v_grid, indexing="xy")
+
+    # Expand the grids for batch processing
+    u_grid = u_grid.unsqueeze(0).expand(perspective_depth_batch.shape[0], -1, -1)
+    v_grid = v_grid.unsqueeze(0).expand(perspective_depth_batch.shape[0], -1, -1)
+
+    # Compute the squared terms for efficiency
+    x_term = ((u_grid - cx) / fx) ** 2
+    y_term = ((v_grid - cy) / fy) ** 2
+
+    # Calculate the orthogonal (normal) depth
+    normal_depth = perspective_depth_batch / torch.sqrt(1 + x_term + y_term)
+
+    # Restore the last dimension if it was present in the input
+    if add_last_dim:
+        normal_depth = normal_depth.unsqueeze(-1)
+
+    # Return to original shape if input was not batched
+    if not is_batched:
+        normal_depth = normal_depth.squeeze(0)
+
+    return normal_depth
+
+
 @torch.jit.script
 def project_points(points: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
     r"""Projects 3D points into 2D image plane.

source/extensions/omni.isaac.lab/test/utils/test_math.py

@@ -376,6 +376,24 @@ def iter_old_quat_rotate_inverse(q: torch.Tensor, v: torch.Tensor) -> torch.Tens
                 iter_old_quat_rotate_inverse(q_rand, v_rand),
             )
 
+    def test_depth_perspective_conversion(self):
+        # Create a sample perspective depth image (N, H, W)
+        perspective_depth = torch.tensor([[[10.0, 0.0, 100.0], [0.0, 3000.0, 0.0], [100.0, 0.0, 100.0]]])
+
+        # Create sample intrinsic matrix (3, 3)
+        intrinsics = torch.tensor([[500.0, 0.0, 5.0], [0.0, 500.0, 5.0], [0.0, 0.0, 1.0]])
+
+        # Convert perspective depth to orthogonal depth
+        orthogonal_depth = math_utils.convert_perspective_depth_to_orthogonal_depth(perspective_depth, intrinsics)
+
+        # Manually compute expected orthogonal depth based on the formula for comparison
+        expected_orthogonal_depth = torch.tensor(
+            [[[9.9990, 0.0000, 99.9932], [0.0000, 2999.8079, 0.0000], [99.9932, 0.0000, 99.9964]]]
+        )
+
+        # Assert that the output is close to the expected result
+        torch.testing.assert_close(orthogonal_depth, expected_orthogonal_depth)
+
 
 if __name__ == "__main__":
     run_tests()

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

@@ -1,6 +1,15 @@
 Changelog
 ---------
 
+0.10.6 (2024-09-20)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+* Added ``Isaac-Cartpole-RGB-Camera-v0`` and ``Isaac-Cartpole-Depth-Camera-v0``
+  manager based camera cartpole environments.
+
+
 0.10.5 (2024-09-11)
 ~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/__init__.py

@@ -10,6 +10,7 @@
 import gymnasium as gym
 
 from . import agents
+from .cartpole_camera_env_cfg import CartpoleDepthCameraEnvCfg, CartpoleRGBCameraEnvCfg
 from .cartpole_env_cfg import CartpoleEnvCfg
 
 ##
@@ -28,3 +29,23 @@
         "sb3_cfg_entry_point": f"{agents.__name__}:sb3_ppo_cfg.yaml",
     },
 )
+
+gym.register(
+    id="Isaac-Cartpole-RGB-Camera-v0",
+    entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": CartpoleRGBCameraEnvCfg,
+        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml",
+    },
+)
+
+gym.register(
+    id="Isaac-Cartpole-Depth-Camera-v0",
+    entry_point="omni.isaac.lab.envs:ManagerBasedRLEnv",
+    disable_env_checker=True,
+    kwargs={
+        "env_cfg_entry_point": CartpoleDepthCameraEnvCfg,
+        "rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_camera_ppo_cfg.yaml",
+    },
+)