Support video frame output for Cosmos

docs/conf.py

@@ -127,6 +127,7 @@
 # Mock out modules that are not available on RTD
 autodoc_mock_imports = [
     "torch",
+    "torchvision",
     "numpy",
     "matplotlib",
     "scipy",

source/isaaclab/isaaclab/envs/manager_based_rl_env.py

@@ -72,14 +72,15 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **
             render_mode: The render mode for the environment. Defaults to None, which
                 is similar to ``"human"``.
         """
+        # -- counter for curriculum
+        self.common_step_counter = 0
+
         # initialize the base class to setup the scene.
         super().__init__(cfg=cfg)
         # store the render mode
         self.render_mode = render_mode
 
         # initialize data and constants
-        # -- counter for curriculum
-        self.common_step_counter = 0
         # -- init buffers
         self.episode_length_buf = torch.zeros(self.num_envs, device=self.device, dtype=torch.long)
         # -- set the framerate of the gym video recorder wrapper so that the playback speed of the produced video matches the simulation

source/isaaclab/isaaclab/envs/mdp/observations.py

@@ -11,7 +11,9 @@
 
 from __future__ import annotations
 
+import os
 import torch
+from torchvision.utils import save_image
 from typing import TYPE_CHECKING
 
 import isaaclab.utils.math as math_utils
@@ -237,6 +239,8 @@ def image(
     data_type: str = "rgb",
     convert_perspective_to_orthogonal: bool = False,
     normalize: bool = True,
+    save_image_to_file: bool = False,
+    image_path: str = "image",
 ) -> torch.Tensor:
     """Images of a specific datatype from the camera sensor.
 
@@ -276,6 +280,18 @@ def image(
             images -= mean_tensor
         elif "distance_to" in data_type or "depth" in data_type:
             images[images == float("inf")] = 0
+        elif data_type == "normals":
+            images = (images + 1.0) * 0.5
+
+    if save_image_to_file:
+        dir_path, _ = os.path.split(image_path)
+        if dir_path:
+            os.makedirs(dir_path, exist_ok=True)
+        if images.dtype == torch.uint8:
+            images = images.float() / 255.0
+        for tile in range(images.shape[0]):
+            tile_chw = torch.swapaxes(images[tile : tile + 1].unsqueeze(1), 1, -1).squeeze(-1)
+            save_image(tile_chw, f"{image_path}_{data_type}_{tile}_{env.common_step_counter}.png")
 
     return images.clone()
 

source/isaaclab/isaaclab/sensors/camera/camera.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import json
 import numpy as np
 import re
 import torch
@@ -455,7 +456,10 @@ def _initialize_impl(self):
                 #   if colorize is true, the data is mapped to colors and a uint8 4 channel image is returned.
                 #   if colorize is false, the data is returned as a uint32 image with ids as values.
                 if name == "semantic_segmentation":
-                    init_params = {"colorize": self.cfg.colorize_semantic_segmentation}
+                    init_params = {
+                        "colorize": self.cfg.colorize_semantic_segmentation,
+                        "mapping": json.dumps(self.cfg.semantic_segmentation_mapping),
+                    }
                 elif name == "instance_segmentation_fast":
                     init_params = {"colorize": self.cfg.colorize_instance_segmentation}
                 elif name == "instance_id_segmentation_fast":

source/isaaclab/isaaclab/sensors/camera/camera_cfg.py

@@ -114,3 +114,19 @@ class OffsetCfg:
     If True, instance segmentation is converted to an image where instance IDs are mapped to colors.
     and returned as a ``uint8`` 4-channel array. If False, the output is returned as a ``int32`` array.
     """
+
+    semantic_segmentation_mapping: dict = {}
+    """Dictionary mapping semantics to specific colours
+
+    Eg.
+    ```
+    {
+        "class:cube_1": (255, 36, 66, 255),
+        "class:cube_2": (255, 184, 48, 255),
+        "class:cube_3": (55, 255, 139, 255),
+        "class:table": (255, 237, 218, 255),
+        "class:ground": (100, 100, 100, 255),
+        "class:robot": (61, 178, 255, 255),
+    }
+    ```
+    """

source/isaaclab/isaaclab/sensors/camera/tiled_camera.py

@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import json
 import math
 import numpy as np
 import torch
@@ -209,7 +210,10 @@ def _initialize_impl(self):
             else:
                 init_params = None
                 if annotator_type == "semantic_segmentation":
-                    init_params = {"colorize": self.cfg.colorize_semantic_segmentation}
+                    init_params = {
+                        "colorize": self.cfg.colorize_semantic_segmentation,
+                        "mapping": json.dumps(self.cfg.semantic_segmentation_mapping),
+                    }
                 elif annotator_type == "instance_segmentation_fast":
                     init_params = {"colorize": self.cfg.colorize_instance_segmentation}
                 elif annotator_type == "instance_id_segmentation_fast":

source/isaaclab/isaaclab/utils/__init__.py

@@ -11,7 +11,6 @@
 from .dict import *
 from .interpolation import *
 from .modifiers import *
-from .notebook_widgets import *
 from .string import *
 from .timer import Timer
 from .types import *

source/isaaclab_mimic/isaaclab_mimic/envs/__init__.py

@@ -7,6 +7,7 @@
 
 import gymnasium as gym
 
+from .franka_stack_ik_rel_blueprint_mimic_env_cfg import FrankaCubeStackIKRelBlueprintMimicEnvCfg
 from .franka_stack_ik_rel_mimic_env import FrankaCubeStackIKRelMimicEnv
 from .franka_stack_ik_rel_mimic_env_cfg import FrankaCubeStackIKRelMimicEnvCfg
 
@@ -22,3 +23,12 @@
     },
     disable_env_checker=True,
 )
+
+gym.register(
+    id="Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-Mimic-v0",
+    entry_point="isaaclab_mimic.envs:FrankaCubeStackIKRelMimicEnv",
+    kwargs={
+        "env_cfg_entry_point": franka_stack_ik_rel_blueprint_mimic_env_cfg.FrankaCubeStackIKRelBlueprintMimicEnvCfg,
+    },
+    disable_env_checker=True,
+)

source/isaaclab_mimic/isaaclab_mimic/envs/franka_stack_ik_rel_blueprint_mimic_env_cfg.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2024-2025, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig
+from isaaclab.utils import configclass
+
+from isaaclab_tasks.manager_based.manipulation.stack.config.franka.stack_ik_rel_blueprint_env_cfg import (
+    FrankaCubeStackBlueprintEnvCfg,
+)
+
+
+@configclass
+class FrankaCubeStackIKRelBlueprintMimicEnvCfg(FrankaCubeStackBlueprintEnvCfg, MimicEnvCfg):
+    """
+    Isaac Lab Mimic environment config class for Franka Cube Stack IK Rel env.
+    """
+
+    def __post_init__(self):
+        # post init of parents
+        super().__post_init__()
+
+        # Override the existing values
+        self.datagen_config.name = "isaac_lab_franka_stack_ik_rel_blueprint_D0"
+        self.datagen_config.generation_guarantee = True
+        self.datagen_config.generation_keep_failed = True
+        self.datagen_config.generation_num_trials = 10
+        self.datagen_config.generation_select_src_per_subtask = True
+        self.datagen_config.generation_transform_first_robot_pose = False
+        self.datagen_config.generation_interpolate_from_last_target_pose = True
+        self.datagen_config.max_num_failures = 25
+        self.datagen_config.seed = 1
+
+        # The following are the subtask configurations for the stack task.
+        subtask_configs = []
+        subtask_configs.append(
+            SubTaskConfig(
+                # Each subtask involves manipulation with respect to a single object frame.
+                object_ref="cube_2",
+                # This key corresponds to the binary indicator in "datagen_info" that signals
+                # when this subtask is finished (e.g., on a 0 to 1 edge).
+                subtask_term_signal="grasp_1",
+                # Specifies time offsets for data generation when splitting a trajectory into
+                # subtask segments. Random offsets are added to the termination boundary.
+                subtask_term_offset_range=(10, 20),
+                # Selection strategy for the source subtask segment during data generation
+                selection_strategy="nearest_neighbor_object",
+                # Optional parameters for the selection strategy function
+                selection_strategy_kwargs={"nn_k": 3},
+                # Amount of action noise to apply during this subtask
+                action_noise=0.03,
+                # Number of interpolation steps to bridge to this subtask segment
+                num_interpolation_steps=5,
+                # Additional fixed steps for the robot to reach the necessary pose
+                num_fixed_steps=0,
+                # If True, apply action noise during the interpolation phase and execution
+                apply_noise_during_interpolation=False,
+            )
+        )
+        subtask_configs.append(
+            SubTaskConfig(
+                # Each subtask involves manipulation with respect to a single object frame.
+                object_ref="cube_1",
+                # Corresponding key for the binary indicator in "datagen_info" for completion
+                subtask_term_signal="stack_1",
+                # Time offsets for data generation when splitting a trajectory
+                subtask_term_offset_range=(10, 20),
+                # Selection strategy for source subtask segment
+                selection_strategy="nearest_neighbor_object",
+                # Optional parameters for the selection strategy function
+                selection_strategy_kwargs={"nn_k": 3},
+                # Amount of action noise to apply during this subtask
+                action_noise=0.03,
+                # Number of interpolation steps to bridge to this subtask segment
+                num_interpolation_steps=5,
+                # Additional fixed steps for the robot to reach the necessary pose
+                num_fixed_steps=0,
+                # If True, apply action noise during the interpolation phase and execution
+                apply_noise_during_interpolation=False,
+            )
+        )
+        subtask_configs.append(
+            SubTaskConfig(
+                # Each subtask involves manipulation with respect to a single object frame.
+                object_ref="cube_3",
+                # Corresponding key for the binary indicator in "datagen_info" for completion
+                subtask_term_signal="grasp_2",
+                # Time offsets for data generation when splitting a trajectory
+                subtask_term_offset_range=(10, 20),
+                # Selection strategy for source subtask segment
+                selection_strategy="nearest_neighbor_object",
+                # Optional parameters for the selection strategy function
+                selection_strategy_kwargs={"nn_k": 3},
+                # Amount of action noise to apply during this subtask
+                action_noise=0.03,
+                # Number of interpolation steps to bridge to this subtask segment
+                num_interpolation_steps=5,
+                # Additional fixed steps for the robot to reach the necessary pose
+                num_fixed_steps=0,
+                # If True, apply action noise during the interpolation phase and execution
+                apply_noise_during_interpolation=False,
+            )
+        )
+        subtask_configs.append(
+            SubTaskConfig(
+                # Each subtask involves manipulation with respect to a single object frame.
+                object_ref="cube_2",
+                # End of final subtask does not need to be detected
+                subtask_term_signal=None,
+                # No time offsets for the final subtask
+                subtask_term_offset_range=(0, 0),
+                # Selection strategy for source subtask segment
+                selection_strategy="nearest_neighbor_object",
+                # Optional parameters for the selection strategy function
+                selection_strategy_kwargs={"nn_k": 3},
+                # Amount of action noise to apply during this subtask
+                action_noise=0.03,
+                # Number of interpolation steps to bridge to this subtask segment
+                num_interpolation_steps=5,
+                # Additional fixed steps for the robot to reach the necessary pose
+                num_fixed_steps=0,
+                # If True, apply action noise during the interpolation phase and execution
+                apply_noise_during_interpolation=False,
+            )
+        )
+        self.subtask_configs["franka"] = subtask_configs

source/isaaclab_tasks/isaaclab_tasks/manager_based/manipulation/stack/config/franka/__init__.py

@@ -7,6 +7,7 @@
 
 from . import (
     agents,
+    stack_ik_rel_blueprint_env_cfg,
     stack_ik_rel_env_cfg,
     stack_ik_rel_instance_randomize_env_cfg,
     stack_joint_pos_env_cfg,
@@ -62,3 +63,12 @@
     },
     disable_env_checker=True,
 )
+
+gym.register(
+    id="Isaac-Stack-Cube-Franka-IK-Rel-Blueprint-v0",
+    entry_point="isaaclab.envs:ManagerBasedRLEnv",
+    kwargs={
+        "env_cfg_entry_point": stack_ik_rel_blueprint_env_cfg.FrankaCubeStackBlueprintEnvCfg,
+    },
+    disable_env_checker=True,
+)