issacdataengine/nimbus_extension/components/load/env_loader.py

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))

        try:
            from isaacsim import SimulationApp  # IS 5.x
        except ImportError:
            from omni.isaac.kit import SimulationApp  # IS 4.x

        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