init commit

nimbus_extension/components/load/__init__.py

@@ -0,0 +1,8 @@
+import nimbus.components.load as _load
+
+from .env_loader import EnvLoader
+from .env_randomizer import EnvRandomizer
+
+_load.register_loader("env_loader", EnvLoader)
+
+_load.register_randomizer("env_randomizer", EnvRandomizer)

nimbus_extension/components/load/env_loader.py

@@ -0,0 +1,180 @@
+import time
+from fractions import Fraction
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.package import Package
+from nimbus.components.data.scene import Scene
+from nimbus.components.load import SceneLoader
+from nimbus.daemon import ComponentStatus, StatusReporter
+from nimbus.daemon.decorators import status_monitor
+from nimbus.utils.flags import get_random_seed
+from workflows.base import create_workflow
+
+
+class EnvLoader(SceneLoader):
+    """
+    Environment loader that initializes Isaac Sim and loads scenes based on workflow configurations.
+
+    This loader integrates with the workflow system to manage scene loading and task execution.
+    It supports two operating modes:
+    - Standalone mode (pack_iter=None): Loads tasks directly from workflow configuration
+    - Pipeline mode (pack_iter provided): Loads tasks from a package iterator
+
+    It also supports task repetition for data augmentation across different random seeds.
+
+    Args:
+        pack_iter (Iterator[Package]): An iterator from the previous component. None for standalone.
+        cfg_path (str): Path to the workflow configuration file.
+        workflow_type (str): Type of workflow to create (e.g., 'SimBoxDualWorkFlow').
+        simulator (dict): Simulator configuration including physics_dt, rendering_dt, headless, etc.
+        task_repeat (int): How many times to repeat each task before advancing (-1 means single execution).
+        need_preload (bool): Whether to preload assets on scene initialization.
+        scene_info (str): Configuration key for scene information in the workflow config.
+    """
+
+    def __init__(
+        self,
+        pack_iter: Iterator[Package],
+        cfg_path: str,
+        workflow_type: str,
+        simulator: dict,
+        task_repeat: int = -1,
+        need_preload: bool = False,
+        scene_info: str = "dining_room_scene_info",
+    ):
+        init_start_time = time.time()
+        super().__init__(pack_iter)
+
+        self.status_reporter = StatusReporter(self.__class__.__name__)
+        self.status_reporter.update_status(ComponentStatus.IDLE)
+        self.need_preload = need_preload
+        self.task_repeat_cnt = task_repeat
+        self.task_repeat_idx = 0
+        self.workflow_type = workflow_type
+
+        # Parse simulator config
+        physics_dt = simulator.get("physics_dt", "1/30")
+        rendering_dt = simulator.get("rendering_dt", "1/30")
+        if isinstance(physics_dt, str):
+            physics_dt = float(Fraction(physics_dt))
+        if isinstance(rendering_dt, str):
+            rendering_dt = float(Fraction(rendering_dt))
+
+        from isaacsim import SimulationApp
+
+        self.simulation_app = SimulationApp(
+            {
+                "headless": simulator.get("headless", True),
+                "anti_aliasing": simulator.get("anti_aliasing", 3),
+                "multi_gpu": simulator.get("multi_gpu", True),
+                "renderer": simulator.get("renderer", "RayTracedLighting"),
+            }
+        )
+
+        self.logger.info(f"simulator params: physics dt={physics_dt}, rendering dt={rendering_dt}")
+        from omni.isaac.core import World
+
+        world = World(
+            physics_dt=physics_dt,
+            rendering_dt=rendering_dt,
+            stage_units_in_meters=simulator.get("stage_units_in_meters", 1.0),
+        )
+
+        # Import workflow extensions and create workflow
+        from workflows import import_extensions
+
+        import_extensions(workflow_type)
+        self.workflow = create_workflow(
+            workflow_type,
+            world,
+            cfg_path,
+            scene_info=scene_info,
+            random_seed=get_random_seed(),
+        )
+
+        self.scene = None
+        self.task_finish = False
+        self.cur_index = 0
+        self.record_init_time(time.time() - init_start_time)
+
+        self.status_reporter.update_status(ComponentStatus.READY)
+
+    @status_monitor()
+    def _init_next_task(self):
+        """
+        Internal helper method to initialize and return the next task as a Scene object.
+
+        Handles task repetition logic and advances the task index when all repetitions are complete.
+
+        Returns:
+            Scene: Initialized scene object for the next task.
+
+        Raises:
+            StopIteration: When all tasks have been exhausted.
+        """
+        if self.scene is not None and self.task_repeat_cnt > 0 and self.task_repeat_idx < self.task_repeat_cnt:
+            self.logger.info(f"Task execute times {self.task_repeat_idx + 1}/{self.task_repeat_cnt}")
+            self.workflow.init_task(self.cur_index - 1, self.need_preload)
+            self.task_repeat_idx += 1
+            scene = Scene(
+                name=self.workflow.get_task_name(),
+                wf=self.workflow,
+                task_id=self.cur_index - 1,
+                task_exec_num=self.task_repeat_idx,
+                simulation_app=self.simulation_app,
+            )
+            return scene
+        if self.cur_index >= len(self.workflow.task_cfgs):
+            self.logger.info("No more tasks to load, stopping iteration.")
+            raise StopIteration
+        self.logger.info(f"Loading task {self.cur_index + 1}/{len(self.workflow.task_cfgs)}")
+        self.workflow.init_task(self.cur_index, self.need_preload)
+        self.task_repeat_idx = 1
+        scene = Scene(
+            name=self.workflow.get_task_name(),
+            wf=self.workflow,
+            task_id=self.cur_index,
+            task_exec_num=self.task_repeat_idx,
+            simulation_app=self.simulation_app,
+        )
+        self.cur_index += 1
+        return scene
+
+    def load_asset(self) -> Scene:
+        """
+        Load and initialize the next scene from workflow.
+
+        Supports two modes:
+        - Standalone: Iterates through workflow tasks directly
+        - Pipeline: Synchronizes with incoming packages and applies plan info to scene
+
+        Returns:
+            Scene: The loaded and initialized Scene object.
+
+        Raises:
+            StopIteration: When no more scenes are available.
+        """
+        try:
+            # Standalone mode: load tasks directly from workflow
+            if self.pack_iter is None:
+                self.scene = self._init_next_task()
+            # Pipeline mode: load tasks from package iterator
+            else:
+                package = next(self.pack_iter)
+                self.cur_index = package.task_id
+
+                # Initialize scene if this is the first package or a new task
+                if self.scene is None:
+                    self.scene = self._init_next_task()
+                elif self.cur_index > self.scene.task_id:
+                    self.scene = self._init_next_task()
+
+                # Apply plan information from package to scene
+                package.data = self.scene.wf.dedump_plan_info(package.data)
+                self.scene.add_plan_info(package.data)
+
+            return self.scene
+        except StopIteration:
+            raise StopIteration
+        except Exception as e:
+            raise e

nimbus_extension/components/load/env_randomizer.py

@@ -0,0 +1,46 @@
+import os
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.scene import Scene
+from nimbus.components.load import LayoutRandomizer
+
+
+class EnvRandomizer(LayoutRandomizer):
+    """
+    Environment randomizer that extends the base layout randomizer to include additional randomization
+    capabilities specific to the simulation environment.
+    This class can be used to randomize various aspects of the environment, such as object placements,
+    textures, lighting conditions, and other scene parameters, based on the provided configuration.
+    The randomization process can be controlled through the number of randomizations to perform and
+    whether to operate in strict mode.
+
+    Args:
+        scene_iter (Iterator[Scene]): An iterator that yields scenes to be randomized.
+        random_num (int): How many randomizations to perform for each scene.
+        strict_mode (bool): Whether to operate in strict mode, which enforces certain constraints
+            on the randomization process.
+        input_dir (str): Directory from which to load additional randomization data such as object
+            placements or textures. If None, randomization is performed without loading additional data.
+    """
+
+    def __init__(
+        self, scene_iter: Iterator[Scene], random_num: int = 1, strict_mode: bool = False, input_dir: str = None
+    ):
+        super().__init__(scene_iter, random_num, strict_mode)
+        assert self.random_num > 0, "random_num must be greater than 0"
+        self.input_dir = input_dir
+        if self.input_dir is not None:
+            self.paths_names = os.listdir(self.input_dir)
+            self.random_num = len(self.paths_names)
+
+    def randomize_scene(self, scene) -> Scene:
+        if scene.plan_info is None:
+            path = None
+            if self.input_dir is not None:
+                path = os.path.join(self.input_dir, self.paths_names[self.cur_index])
+            if not scene.wf.randomization(path):
+                return None
+        else:
+            if not scene.wf.randomization_from_mem(scene.plan_info):
+                return None
+        return scene