Replaces TensorDict with native dictionary (#1348)

# Description

Remove TensorDict usage from Isaac Lab

## Type of change

- Breaking change... for those ones using TensorDict specific API such
as `camera.data.output.to_dict()` or `camera.data.output.sorted_keys`

## Screenshots

Before

![Screenshot from 2024-10-30
12-45-06](https://github.com/user-attachments/assets/2dc0e827-3e12-4ae9-849e-e9f75c718157)

After

![Screenshot from 2024-10-30
16-04-35](https://github.com/user-attachments/assets/715b6cb0-9f87-4938-8dbd-5c56203cb90e)

## Checklist

- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Co-authored-by: Kelly Guo <kellyguo123@hotmail.com>

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

 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.27.11"
+version = "0.27.12"
 
 # Description
 title = "Isaac Lab framework for Robot Learning"

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

 Changelog
 ---------
 
+0.27.12 (2024-01-04)
+~~~~~~~~~~~~~~~~~~~
+
+Removed
+^^^^^^^
+
+* Removed TensorDict usage in favor of Python dictionary in sensors
+
+
 0.27.11 (2024-10-31)
 ~~~~~~~~~~~~~~~~~~~~
 

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera.py

@@ -9,7 +9,6 @@ import numpy as np
 import re
 import torch
 from collections.abc import Sequence
-from tensordict import TensorDict
 from typing import TYPE_CHECKING, Any, Literal
 
 import carb
@@ -156,7 +155,7 @@ class Camera(SensorBase):
         # message for class
         return (
             f"Camera @ '{self.cfg.prim_path}': \n"
-            f"\tdata types   : {self.data.output.sorted_keys} \n"
+            f"\tdata types   : {list(self.data.output.keys())} \n"
             f"\tsemantic filter : {self.cfg.semantic_filter}\n"
             f"\tcolorize semantic segm.   : {self.cfg.colorize_semantic_segmentation}\n"
             f"\tcolorize instance segm.   : {self.cfg.colorize_instance_segmentation}\n"
@@ -497,7 +496,7 @@ class Camera(SensorBase):
         self._update_poses(env_ids)
         # -- read the data from annotator registry
         # check if buffer is called for the first time. If so then, allocate the memory
-        if len(self._data.output.sorted_keys) == 0:
+        if len(self._data.output) == 0:
             # this is the first time buffer is called
             # it allocates memory for all the sensors
             self._create_annotator_data()
@@ -552,7 +551,7 @@ class Camera(SensorBase):
         # lazy allocation of data dictionary
         # since the size of the output data is not known in advance, we leave it as None
         # the memory will be allocated when the buffer() function is called for the first time.
-        self._data.output = TensorDict({}, batch_size=self._view.count, device=self.device)
+        self._data.output = {}
         self._data.info = [{name: None for name in self.cfg.data_types} for _ in range(self._view.count)]
 
     def _update_intrinsic_matrices(self, env_ids: Sequence[int]):

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/camera_data.py

@@ -5,7 +5,6 @@
 
 import torch
 from dataclasses import dataclass
-from tensordict import TensorDict
 from typing import Any
 
 from omni.isaac.lab.utils.math import convert_camera_frame_orientation_convention
@@ -47,7 +46,7 @@ class CameraData:
     Shape is (N, 3, 3) where N is the number of sensors.
     """
 
-    output: TensorDict = None
+    output: dict[str, torch.Tensor] = None
     """The retrieved sensor data with sensor types as key.
 
     The format of the data is available in the `Replicator Documentation`_. For semantic-based data,

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/camera/tiled_camera.py

@@ -9,7 +9,6 @@ import math
 import numpy as np
 import torch
 from collections.abc import Sequence
-from tensordict import TensorDict
 from typing import TYPE_CHECKING, Any
 
 import carb
@@ -106,7 +105,7 @@ class TiledCamera(Camera):
         # message for class
         return (
             f"Tiled Camera @ '{self.cfg.prim_path}': \n"
-            f"\tdata types   : {self.data.output.sorted_keys} \n"
+            f"\tdata types   : {list(self.data.output.keys())} \n"
             f"\tsemantic filter : {self.cfg.semantic_filter}\n"
             f"\tcolorize semantic segm.   : {self.cfg.colorize_semantic_segmentation}\n"
             f"\tcolorize instance segm.   : {self.cfg.colorize_instance_segmentation}\n"
@@ -372,7 +371,7 @@ class TiledCamera(Camera):
                     (self._view.count, self.cfg.height, self.cfg.width, 1), device=self.device, dtype=torch.int32
                 ).contiguous()
 
-        self._data.output = TensorDict(data_dict, batch_size=self._view.count, device=self.device)
+        self._data.output = data_dict
         self._data.info = dict()
 
     def _tiled_image_shape(self) -> tuple[int, int]:

source/extensions/omni.isaac.lab/omni/isaac/lab/sensors/ray_caster/ray_caster_camera.py

@@ -7,7 +7,6 @@ from __future__ import annotations
 
 import torch
 from collections.abc import Sequence
-from tensordict import TensorDict
 from typing import TYPE_CHECKING, ClassVar, Literal
 
 import omni.isaac.core.utils.stage as stage_utils
@@ -347,7 +346,7 @@ class RayCasterCamera(RayCaster):
         self._data.image_shape = self.image_shape
         # -- output data
         # create the buffers to store the annotator data.
-        self._data.output = TensorDict({}, batch_size=self._view.count, device=self.device)
+        self._data.output = {}
         self._data.info = [{name: None for name in self.cfg.data_types}] * self._view.count
         for name in self.cfg.data_types:
             if name in ["distance_to_image_plane", "distance_to_camera"]:

source/extensions/omni.isaac.lab/test/sensors/test_camera.py

@@ -121,7 +121,7 @@ class TestCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for im_data in camera.data.output.to_dict().values():
+            for im_data in camera.data.output.values():
                 self.assertEqual(im_data.shape, (1, self.camera_cfg.height, self.camera_cfg.width, 1))
 
     def test_camera_init_offset(self):
@@ -228,7 +228,7 @@ class TestCamera(unittest.TestCase):
             cam_2.update(self.dt)
             # check image data
             for cam in [cam_1, cam_2]:
-                for im_data in cam.data.output.to_dict().values():
+                for im_data in cam.data.output.values():
                     self.assertEqual(im_data.shape, (1, self.camera_cfg.height, self.camera_cfg.width, 1))
 
     def test_multi_camera_with_different_resolution(self):
@@ -705,7 +705,7 @@ class TestCamera(unittest.TestCase):
             with Timer(f"Time taken for writing data with shape {camera.image_shape}   "):
                 # Pack data back into replicator format to save them using its writer
                 rep_output = {"annotators": {}}
-                camera_data = convert_dict_to_backend(camera.data.output[0].to_dict(), backend="numpy")
+                camera_data = convert_dict_to_backend({k: v[0] for k, v in camera.data.output.items()}, backend="numpy")
                 for key, data, info in zip(camera_data.keys(), camera_data.values(), camera.data.info[0].values()):
                     if info is not None:
                         rep_output["annotators"][key] = {"render_product": {"data": data, **info}}

source/extensions/omni.isaac.lab/test/sensors/test_ray_caster_camera.py

@@ -129,7 +129,7 @@ class TestWarpCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for im_data in camera.data.output.to_dict().values():
+            for im_data in camera.data.output.values():
                 self.assertEqual(
                     im_data.shape, (1, self.camera_cfg.pattern_cfg.height, self.camera_cfg.pattern_cfg.width, 1)
                 )
@@ -147,7 +147,7 @@ class TestWarpCamera(unittest.TestCase):
             self.sim.step()
         camera.update(self.dt)
         # access image data and compare shapes
-        for im_data in camera.data.output.to_dict().values():
+        for im_data in camera.data.output.values():
             self.assertTrue(
                 im_data.shape == (1, self.camera_cfg.pattern_cfg.height, self.camera_cfg.pattern_cfg.width, 1)
             )
@@ -289,7 +289,7 @@ class TestWarpCamera(unittest.TestCase):
             cam_2.update(self.dt)
             # check image data
             for cam in [cam_1, cam_2]:
-                for im_data in cam.data.output.to_dict().values():
+                for im_data in cam.data.output.values():
                     self.assertEqual(
                         im_data.shape, (1, self.camera_cfg.pattern_cfg.height, self.camera_cfg.pattern_cfg.width, 1)
                     )
@@ -392,7 +392,7 @@ class TestWarpCamera(unittest.TestCase):
             with Timer(f"Time taken for writing data with shape {camera.image_shape}   "):
                 # Pack data back into replicator format to save them using its writer
                 rep_output = {"annotators": {}}
-                camera_data = convert_dict_to_backend(camera.data.output[0].to_dict(), backend="numpy")
+                camera_data = convert_dict_to_backend({k: v[0] for k, v in camera.data.output.items()}, backend="numpy")
                 for key, data, info in zip(camera_data.keys(), camera_data.values(), camera.data.info[0].values()):
                     if info is not None:
                         rep_output["annotators"][key] = {"render_product": {"data": data, **info}}

source/extensions/omni.isaac.lab/test/sensors/test_tiled_camera.py

@@ -111,7 +111,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for im_type, im_data in camera.data.output.to_dict().items():
+            for im_type, im_data in camera.data.output.items():
                 if im_type == "rgb":
                     self.assertEqual(im_data.shape, (1, self.camera_cfg.height, self.camera_cfg.width, 3))
                     self.assertGreater((im_data / 255.0).mean().item(), 0.0)
@@ -162,7 +162,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for im_type, im_data in camera.data.output.to_dict().items():
+            for im_type, im_data in camera.data.output.items():
                 if im_type == "rgb":
                     self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3))
                     for i in range(4):
@@ -347,7 +347,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4))
                 for i in range(4):
                     self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0)
@@ -399,7 +399,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1))
                 for i in range(4):
                     self.assertGreater((im_data[i]).mean().item(), 0.0)
@@ -451,7 +451,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1))
                 for i in range(4):
                     self.assertGreater((im_data[i]).mean().item(), 0.0)
@@ -503,7 +503,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3))
                 for i in range(4):
                     self.assertGreater((im_data[i]).mean().item(), 0.0)
@@ -555,7 +555,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 2))
                 for i in range(4):
                     self.assertGreater((im_data[i]).mean().item(), 0.0)
@@ -607,7 +607,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4))
                 for i in range(4):
                     self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0)
@@ -660,7 +660,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4))
                 for i in range(num_cameras):
                     self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0)
@@ -713,7 +713,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 4))
                 for i in range(num_cameras):
                     self.assertGreater((im_data[i] / 255.0).mean().item(), 0.0)
@@ -767,7 +767,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1))
                 for i in range(num_cameras):
                     self.assertGreater(im_data[i].to(dtype=float).mean().item(), 0.0)
@@ -822,7 +822,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1))
                 for i in range(num_cameras):
                     self.assertGreater(im_data[i].to(dtype=float).mean().item(), 0.0)
@@ -876,7 +876,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for _, im_data in camera.data.output.to_dict().items():
+            for _, im_data in camera.data.output.items():
                 self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 1))
                 for i in range(num_cameras):
                     self.assertGreater(im_data[i].to(dtype=float).mean().item(), 0.0)
@@ -941,7 +941,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for data_type, im_data in camera.data.output.to_dict().items():
+            for data_type, im_data in camera.data.output.items():
                 if data_type in ["rgb", "normals"]:
                     self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3))
                 elif data_type in [
@@ -1039,7 +1039,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for data_type, im_data in camera.data.output.to_dict().items():
+            for data_type, im_data in camera.data.output.items():
                 if data_type in ["rgb", "normals"]:
                     self.assertEqual(im_data.shape, (num_cameras, camera_cfg.height, camera_cfg.width, 3))
                 elif data_type in [
@@ -1135,7 +1135,7 @@ class TestTiledCamera(unittest.TestCase):
             # update camera
             camera.update(self.dt)
             # check image data
-            for data_type, im_data in camera.data.output.to_dict().items():
+            for data_type, im_data in camera.data.output.items():
                 if data_type in ["rgb", "normals"]:
                     self.assertEqual(im_data.shape, (num_cameras, self.camera_cfg.height, self.camera_cfg.width, 3))
                 elif data_type in [
@@ -1201,7 +1201,7 @@ class TestTiledCamera(unittest.TestCase):
             with Timer(f"Time taken for updating camera with shape {camera.image_shape}"):
                 camera.update(self.dt)
             # Check image data
-            for im_type, im_data in camera.data.output.to_dict().items():
+            for im_type, im_data in camera.data.output.items():
                 if im_type == "rgb":
                     self.assertEqual(im_data.shape, (1, camera_cfg.height, camera_cfg.width, 3))
                     self.assertGreater((im_data / 255.0).mean().item(), 0.0)

source/standalone/tutorials/04_sensors/run_ray_caster_camera.py

@@ -129,14 +129,9 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict):
             # Extract camera data
             camera_index = 0
             # note: BasicWriter only supports saving data in numpy format, so we need to convert the data to numpy.
-            if sim.backend == "torch":
-                # tensordict allows easy indexing of tensors in the dictionary
-                single_cam_data = convert_dict_to_backend(camera.data.output[camera_index], backend="numpy")
-            else:
-                # for numpy, we need to manually index the data
-                single_cam_data = dict()
-                for key, value in camera.data.output.items():
-                    single_cam_data[key] = value[camera_index]
+            single_cam_data = convert_dict_to_backend(
+                {k: v[camera_index] for k, v in camera.data.output.items()}, backend="numpy"
+            )
             # Extract the other information
             single_cam_info = camera.data.info[camera_index]
 

source/standalone/tutorials/04_sensors/run_usd_camera.py

@@ -228,8 +228,9 @@ def run_simulator(sim: sim_utils.SimulationContext, scene_entities: dict):
         if args_cli.save:
             # Save images from camera at camera_index
             # note: BasicWriter only supports saving data in numpy format, so we need to convert the data to numpy.
-            # tensordict allows easy indexing of tensors in the dictionary
-            single_cam_data = convert_dict_to_backend(camera.data.output[camera_index], backend="numpy")
+            single_cam_data = convert_dict_to_backend(
+                {k: v[camera_index] for k, v in camera.data.output.items()}, backend="numpy"
+            )
 
             # Extract the other information
             single_cam_info = camera.data.info[camera_index]