WIP

Author: duburcqa

.github/workflows/production.yml

@@ -22,7 +22,7 @@ jobs:
       WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
       HF_HUB_DOWNLOAD_TIMEOUT: 60
-      GENESIS_IMAGE_VER: "1_0"
+      GENESIS_IMAGE_VER: "1_1"
       TIMEOUT_MINUTES: 180
 
     steps:

examples/rigid/single_franka_batch_render.py

@@ -42,15 +42,15 @@ def main():
         pos=(1.5, 0.5, 1.5),
         lookat=(0.0, 0.0, 0.5),
         fov=45,
-        GUI=True,
+        GUI=args.vis,
     )
     cam_0.attach(franka.links[6], trans_to_T(np.array([0.0, 0.5, 0.0])))
     cam_1 = scene.add_camera(
         res=(512, 512),
         pos=(1.5, -0.5, 1.5),
         lookat=(0.0, 0.0, 0.5),
         fov=45,
-        GUI=True,
+        GUI=args.vis,
     )
     scene.add_light(
         pos=[0.0, 0.0, 1.5],
@@ -81,8 +81,8 @@ def main():
             rgba, depth, _, _ = scene.render_all_cameras(rgb=True, depth=True)
             exporter.export_frame_all_cameras(i, rgb=rgba, depth=depth)
         else:
-            rgba, depth, _, _ = cam_0.render(rgb=True, depth=True)
-            exporter.export_frame_single_camera(i, cam_0.idx, rgb=rgba, depth=depth)
+            rgba, depth, _, _ = cam_1.render(rgb=True, depth=True)
+            exporter.export_frame_single_camera(i, cam_1.idx, rgb=rgba, depth=depth)
 
 
 if __name__ == "__main__":

genesis/engine/entities/rigid_entity/rigid_link.py

@@ -1,9 +1,9 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-from numpy.typing import ArrayLike
 import taichi as ti
 import torch
+from numpy.typing import ArrayLike
 
 import genesis as gs
 import trimesh

genesis/engine/scene.py

@@ -1,10 +1,11 @@
 import os
+import pickle
+import time
 
 import numpy as np
 import torch
-import pickle
-import time
 import taichi as ti
+from numpy.typing import ArrayLike
 
 import genesis as gs
 import genesis.utils.geom as gu
@@ -434,95 +435,84 @@ def link_entities(
 
         if child_link._parent_idx != -1:
             gs.logger.warning(
-                "Child entity already has a parent link. This may cause the entity to break into parts. Make sure this operation is intended."
+                "Child entity already has a parent link. This may cause the entity to break into parts. Make sure "
+                "this operation is intended."
             )
         child_link._parent_idx = parent_link.idx
         parent_link._child_idxs.append(child_link.idx)
 
     @gs.assert_unbuilt
     def add_light(
         self,
-        morph: Morph,
-        color=(1.0, 1.0, 1.0, 1.0),
-        intensity=20.0,
-        revert_dir=False,
-        double_sided=False,
-        beam_angle=180.0,
+        *,
+        morph: Morph | None = None,
+        color: ArrayLike | None = (1.0, 1.0, 1.0, 1.0),
+        intensity: float = 20.0,
+        revert_dir: bool | None = False,
+        double_sided: bool | None = False,
+        beam_angle: float | None = 180.0,
+        pos: ArrayLike | None = None,
+        dir: ArrayLike | None = None,
+        directional: bool | None = None,
+        castshadow: bool | None = None,
+        cutoff: float | None = None,
     ):
         """
-        Add a light to the scene. Note that lights added this way can be instantiated from morphs
-        (supporting `gs.morphs.Primitive` or `gs.morphs.Mesh`), and will only be used by the RayTracer renderer.
+        Add a light to the scene.
+
+        Warning
+        -------
+        The signature of this method is different depending on the renderer being used, i.e.:
+        - RayTracer: 'add_light(self, morph, color, intensity, revert_dir, double_sided, beam_angle)'
+        - BatchRender: 'add_ligth(self, pos, dir, intensity, directional, castshadow, cutoff)'
+        - Rasterizer: **Unsupported**
 
         Parameters
         ----------
         morph : gs.morphs.Morph
-            The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`.
+            The morph of the light. Must be an instance of `gs.morphs.Primitive` or `gs.morphs.Mesh`. Only supported by
+            RayTracer.
         color : tuple of float, shape (3,)
-            The color of the light, specified as (r, g, b).
+            The color of the light, specified as (r, g, b). Only supported by RayTracer.
         intensity : float
             The intensity of the light.
         revert_dir : bool
             Whether to revert the direction of the light. If True, the light will be emitted towards the mesh's inside.
+            Only supported by RayTracer.
         double_sided : bool
-            Whether to emit light from both sides of surface.
+            Whether to emit light from both sides of surface. Only supported by RayTracer.
         beam_angle : float
-            The beam angle of the light.
-        """
-        if isinstance(self.renderer_options, gs.renderers.BatchRenderer):
-            gs.logger.warning(
-                "This add_light() function is only supported when NOT using BatchRenderer."
-                "Please use add_light(self, pos, dir, intensity, directional, castshadow, cutoff) instead."
-            )
-            return
-
-        if self.visualizer.raytracer is None:
-            gs.logger.warning("Light is only supported by RayTracer renderer.")
-            return
-
-        if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)):
-            gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.")
-
-        mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False))
-        self.visualizer.raytracer.add_mesh_light(
-            mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, beam_angle
-        )
-
-    @gs.assert_unbuilt
-    def add_light(
-        self,
-        pos,
-        dir,
-        intensity,
-        directional,
-        castshadow,
-        cutoff,
-    ):
-        """
-        Add a light to the scene for batch renderer.
-
-        Parameters
-        ----------
+            The beam angle of the light. Only supported by RayTracer.
         pos : tuple of float, shape (3,)
-            The position of the light, specified as (x, y, z).
+            The position of the light, specified as (x, y, z). Only supported by BatchRenderer.
         dir : tuple of float, shape (3,)
-            The direction of the light, specified as (x, y, z).
+            The direction of the light, specified as (x, y, z). Only supported by BatchRenderer.
         intensity : float
-            The intensity of the light.
+            The intensity of the light. Only supported by BatchRenderer.
         directional : bool
-            Whether the light is directional.
+            Whether the light is directional. Only supported by BatchRenderer.
         castshadow : bool
-            Whether the light casts shadows.
+            Whether the light casts shadows. Only supported by BatchRenderer.
         cutoff : float
-            The cutoff angle of the light in degrees.
+            The cutoff angle of the light in degrees. Only supported by BatchRenderer.
         """
-        if not isinstance(self.renderer_options, gs.renderers.BatchRenderer):
-            gs.logger.warning(
-                "This add_light() function is only supported when using BatchRenderer."
-                "Please use add_light(self, morph, color, intensity, revert_dir, double_sided, beam_angle) instead."
+        if self._visualizer.batch_renderer is not None:
+            if any(map(lambda e: e is None, (pos, dir, intensity, directional, castshadow, cutoff))):
+                gs.raise_exception("Input arguments do not complain with expected signature when using 'BatchRenderer'")
+
+            self.visualizer.add_light(pos, dir, intensity, directional, castshadow, cutoff)
+        elif self.visualizer.raytracer is not None:
+            if any(map(lambda e: e is None, (morph, color, intensity, revert_dir, double_sided, beam_angle))):
+                gs.raise_exception("Input arguments do not complain with expected signature when using 'RayTracer'")
+            if not isinstance(morph, (gs.morphs.Primitive, gs.morphs.Mesh)):
+                gs.raise_exception("Light morph only supports `gs.morphs.Primitive` or `gs.morphs.Mesh`.")
+
+            mesh = gs.Mesh.from_morph_surface(morph, gs.surfaces.Plastic(smooth=False))
+            self.visualizer.raytracer.add_mesh_light(
+                mesh, color, intensity, morph.pos, morph.quat, revert_dir, double_sided, beam_angle
             )
-            return
-
-        self.visualizer.add_light(pos, dir, intensity, directional, castshadow, cutoff)
+        else:
+            gs.raise_exception("Adding lights is only supported by 'RayTracer' and 'BatchRenderer'.")
 
     @gs.assert_unbuilt
     def add_camera(
@@ -1110,7 +1100,7 @@ def render_all_cameras(self, rgb=True, depth=False, normal=False, segmentation=F
         Returns:
             A tuple of tensors of shape (n_envs, H, W, 3) if rgb is not None,
             otherwise a list of tensors of shape (n_envs, H, W, 1) if depth is not None.
-            If n_envs ==0, the first dimension of the tensor is squeezed.
+            If n_envs == 0, the first dimension of the tensor is squeezed.
         """
         return self._visualizer.batch_renderer.render(rgb, depth, normal, segmentation, force_render)
 

genesis/utils/image_exporter.py

@@ -9,17 +9,22 @@
 from genesis.utils.misc import tensor_to_array
 
 
-class FrameImageExporter:
-    @staticmethod
-    def _export_frame_rgb_camera(i_env, export_dir, i_cam, i_step, rgb):
-        # Take the rgb channel in case the rgb tensor has RGBA channel.
-        rgb = np.flip(tensor_to_array(rgb[i_env, ..., :3]), axis=-1)
-        cv2.imwrite(f"{export_dir}/rgb_cam{i_cam}_env{i_env}_{i_step:03d}.png", rgb)
+def _export_frame_rgb_camera(i_env, export_dir, i_cam, i_step, rgb):
+    # Take the rgb channel in case the rgb tensor has RGBA channel.
+    rgb = np.flip(tensor_to_array(rgb[i_env, ..., :3]), axis=-1)
+    cv2.imwrite(f"{export_dir}/rgb_cam{i_cam}_env{i_env}_{i_step:03d}.png", rgb)
+
 
-    @staticmethod
-    def _export_frame_depth_camera(i_env, export_dir, i_cam, i_step, depth):
-        depth = tensor_to_array(depth[i_env])
-        cv2.imwrite(f"{export_dir}/depth_cam{i_cam}_env{i_env}_{i_step:03d}.png", depth)
+def _export_frame_depth_camera(i_env, export_dir, i_cam, i_step, depth):
+    depth = tensor_to_array(depth[i_env])
+    cv2.imwrite(f"{export_dir}/depth_cam{i_cam}_env{i_env}_{i_step:03d}.png", depth)
+
+
+class FrameImageExporter:
+    """
+    This class enables exporting images from all cameras and all environments in batch and in parallel, unlike
+    `Camera.(start|stop)_recording` API, which only allows for exporting images from a single camera and environment.
+    """
 
     def __init__(self, export_dir, depth_clip_max=100, depth_scale="log"):
         self.export_dir = export_dir
@@ -38,7 +43,7 @@ def _normalize_depth(self, depth):
             Normalized depth tensor as uint8
         """
         # Clip depth values
-        depth = depth.clamp(0, self.depth_clip_max)
+        depth = depth.clamp(0.0, self.depth_clip_max)
 
         # Apply scaling if specified
         if self.depth_scale == "log":
@@ -49,7 +54,9 @@ def _normalize_depth(self, depth):
         depth_max = depth.amax(dim=(-3, -2), keepdim=True)
 
         # Normalize to 0-255 range
-        return ((depth - depth_min) / (depth_max - depth_min) * 255).to(torch.uint8)
+        return torch.where(
+            depth_max - depth_min > gs.EPS, ((depth_max - depth) / (depth_max - depth_min) * 255).to(torch.uint8), 0
+        )
 
     def export_frame_all_cameras(self, i_step, camera_idx=None, rgb=None, depth=None):
         """
@@ -58,23 +65,25 @@ def export_frame_all_cameras(self, i_step, camera_idx=None, rgb=None, depth=None
         Args:
             i_step: The current step index.
             camera_idx: array of indices of cameras to export. If None, all cameras are exported.
-            rgb: rgb image is a tuple of tensors of shape (n_envs, H, W, 3).
-            depth: Depth image is a tuple of tensors of shape (n_envs, H, W).
+            rgb: rgb image is a sequence of tensors of shape (n_envs, H, W, 3).
+            depth: Depth image is a sequence of tensors of shape (n_envs, H, W).
         """
         if rgb is None and depth is None:
-            print("No rgb or depth to export")
+            gs.logger.info("No rgb or depth images to export")
             return
-        if rgb is not None:
-            assert isinstance(rgb, tuple) and len(rgb) > 0, "rgb must be a tuple of tensors with length > 0"
-        if depth is not None:
-            assert isinstance(depth, tuple) and len(depth) > 0, "depth must be a tuple of tensors with length > 0"
+        if rgb is not None and (not isinstance(rgb, (tuple, list)) or not rgb):
+            gs.raise_exception("'rgb' must be a non-empty sequence of tensors.")
+        if depth is not None and (not isinstance(depth, (tuple, list)) or not depth):
+            gs.raise_exception("'depth' must be a non-empty sequence of tensors.")
         if camera_idx is None:
-            camera_idx = range(len(depth if rgb is None else rgb))
+            camera_idx = range(len(depth or rgb))
         for i_cam in camera_idx:
-            rgb_cam = rgb[i_cam] if rgb is not None and i_cam < len(rgb) else None
-            depth_cam = depth[i_cam] if depth is not None and i_cam < len(depth) else None
-            if rgb_cam is not None or depth_cam is not None:
-                self.export_frame_single_camera(i_step, i_cam, rgb_cam, depth_cam)
+            rgb_cam, depth_cam = None, None
+            if rgb is not None:
+                rgb_cam = rgb[i_cam]
+            if depth is not None:
+                depth_cam = depth[i_cam]
+            self.export_frame_single_camera(i_step, i_cam, rgb_cam, depth_cam)
 
     def export_frame_single_camera(self, i_step, i_cam, rgb=None, depth=None):
         """
@@ -92,10 +101,11 @@ def export_frame_single_camera(self, i_step, i_cam, rgb=None, depth=None):
             # Unsqueeze rgb to (n_envs, H, W, 3)
             if rgb.ndim == 3:
                 rgb = rgb.unsqueeze(0)
-            assert rgb.ndim == 4, "rgb must be of shape (n_envs, H, W, 3)"
+            if rgb.ndim != 4 or rgb.shape[-1] != 3:
+                gs.raise_exception("'rgb' must be a tensor of shape (n_envs, H, W, 3)")
 
             rgb_job = partial(
-                FrameImageExporter._export_frame_rgb_camera,
+                _export_frame_rgb_camera,
                 export_dir=self.export_dir,
                 i_cam=i_cam,
                 i_step=i_step,
@@ -112,12 +122,13 @@ def export_frame_single_camera(self, i_step, i_cam, rgb=None, depth=None):
             if depth.ndim == 3:
                 depth = depth.unsqueeze(0)
             elif depth.ndim == 2:
-                depth = depth.reshape(1, depth.shape[0], depth.shape[1], 1)
+                depth = depth.reshape((1, *depth.shape, 1))
             depth = self._normalize_depth(depth)
-            assert depth.ndim == 4, "depth must be of shape (n_envs, H, W, 1)"
+            if depth.ndim != 4 or depth.shape[-1] != 1:
+                gs.raise_exception("'rgb' must be a tensor of shape (n_envs, H, W, 1)")
 
             depth_job = partial(
-                FrameImageExporter._export_frame_depth_camera,
+                _export_frame_depth_camera,
                 export_dir=self.export_dir,
                 i_cam=i_cam,
                 i_step=i_step,

genesis/utils/misc.py

@@ -205,9 +205,8 @@ def get_device(backend: gs_backend):
             total_mem = device_property.total_memory / 1024**3
         else:  # pytorch tensors on cpu
             # logger may not be configured at this point
-            getattr(gs, "logger", LOGGER).warning(
-                "No Intel XPU device available. Falling back to CPU for torch device."
-            )
+            logger = getattr(gs, "logger", None) or LOGGER
+            logger.warning("No Intel XPU device available. Falling back to CPU for torch device.")
             device, device_name, total_mem, _ = get_device(gs_backend.cpu)
 
     elif backend == gs_backend.gpu: