init commit

nimbus/components/data/camera.py

@@ -0,0 +1,71 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+
+@dataclass
+class C2W:
+    """
+    Represents a camera-to-world transformation matrix.
+
+    Attributes:
+        matrix (List[float]): A list of 16 floats representing the 4x4 transformation matrix in row-major order.
+    """
+
+    matrix: List[float]
+
+
+@dataclass
+class Camera:
+    """
+    Represents a single camera pose in the trajectory.
+
+    Attributes:
+        trajectory (List[C2W]): List of C2W transformations for this camera pose.
+        intrinsic (Optional[List[float]]): 3x3 camera intrinsic matrix: [[fx, 0, cx], [0, fy, cy], [0, 0, 1]].
+        extrinsic (Optional[List[float]]): 4x4 tobase_extrinsic matrix representing the camera mounting offset
+            relative to the robot base (height + pitch).
+        length (Optional[int]): Length of the trajectory in number of frames.
+        depths (Optional[list[np.ndarray]]): List of depth images captured by this camera.
+        rgbs (Optional[list[np.ndarray]]): List of RGB images captured by this camera.
+        uv_tracks (Optional[Dict[str, Any]]): UV tracking data in the format
+            {mesh_name: {"per_frame": list, "width": W, "height": H}}.
+        uv_mesh_names (Optional[List[str]]): List of mesh names being tracked in the UV tracking data.
+    """
+
+    trajectory: List[C2W]
+    intrinsic: List[float] = None
+    extrinsic: List[float] = None
+    length: int = None
+    depths: list[np.ndarray] = None
+    rgbs: list[np.ndarray] = None
+    uv_tracks: Optional[Dict[str, Any]] = None
+    uv_mesh_names: Optional[List[str]] = None
+
+    def __len__(self):
+        if self.length is not None:
+            return self.length
+        self._check_length()
+        self.length = len(self.trajectory)
+        return len(self.trajectory)
+
+    def _check_length(self):
+        if self.depths is not None and len(self.depths) != len(self.trajectory):
+            raise ValueError("Length of depths does not match length of trajectory")
+        if self.rgbs is not None and len(self.rgbs) != len(self.trajectory):
+            raise ValueError("Length of rgbs does not match length of trajectory")
+        if self.uv_tracks is not None:
+            for mesh_name, track_data in self.uv_tracks.items():
+                if len(track_data["per_frame"]) != len(self.trajectory):
+                    raise ValueError(f"Length of uv_tracks for mesh {mesh_name} does not match length of trajectory")
+
+    def append_rgb(self, rgb_image: np.ndarray):
+        if self.rgbs is None:
+            self.rgbs = []
+        self.rgbs.append(rgb_image)
+
+    def append_depth(self, depth_image: np.ndarray):
+        if self.depths is None:
+            self.depths = []
+        self.depths.append(depth_image)

nimbus/components/data/iterator.py

@@ -0,0 +1,95 @@
+import logging
+import time
+from abc import abstractmethod
+from collections.abc import Iterator
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+# pylint: disable=E0102
+class Iterator(Iterator, Generic[T]):
+    def __init__(self, max_retry=3):
+        self._next_calls = 0.0
+        self._next_total_time = 0.0
+        self._init_time_costs = 0.0
+        self._init_times = 0
+        self._frame_compute_time = 0.0
+        self._frame_compute_frames = 0.0
+        self._frame_io_time = 0.0
+        self._frame_io_frames = 0.0
+        self._wait_time = 0.0
+        self._seq_num = 0.0
+        self._seq_time = 0.0
+        self.logger = logging.getLogger("de_logger")
+        self.max_retry = max_retry
+        self.retry_num = 0
+
+    def record_init_time(self, time_costs):
+        self._init_times += 1
+        self._init_time_costs += time_costs
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        start_time = time.time()
+        try:
+            result = self._next()
+        except StopIteration:
+            self._log_statistics()
+            raise
+        end_time = time.time()
+        self._next_calls += 1
+        self._next_total_time += end_time - start_time
+        return result
+
+    def collect_compute_frame_info(self, length, time_costs):
+        self._frame_compute_frames += length
+        self._frame_compute_time += time_costs
+
+    def collect_io_frame_info(self, length, time_costs):
+        self._frame_io_frames += length
+        self._frame_io_time += time_costs
+
+    def collect_wait_time_info(self, time_costs):
+        self._wait_time += time_costs
+
+    def collect_seq_info(self, length, time_costs):
+        self._seq_num += length
+        self._seq_time += time_costs
+
+    @abstractmethod
+    def _next(self):
+        raise NotImplementedError("Subclasses should implement this method.")
+
+    def _log_statistics(self):
+        class_name = self.__class__.__name__
+        self.logger.info(
+            f"{class_name}: Next method called {self._next_calls} times, total time:"
+            f" {self._next_total_time:.6f} seconds"
+        )
+        if self._init_time_costs > 0:
+            self.logger.info(
+                f"{class_name}: Init time: {self._init_time_costs:.6f} seconds, init {self._init_times} times"
+            )
+        if self._frame_compute_time > 0:
+            avg_compute_time = self._frame_compute_time / self._frame_compute_frames
+            self.logger.info(
+                f"{class_name}: compute frame num: {self._frame_compute_frames}, total time:"
+                f" {self._frame_compute_time:.6f} seconds, average time: {avg_compute_time:.6f} seconds per frame"
+            )
+        if self._frame_io_frames > 0:
+            avg_io_time = self._frame_io_time / self._frame_io_frames
+            self.logger.info(
+                f"{class_name}: io frame num: {self._frame_io_frames}, total time: {self._frame_io_time:.6f} seconds,"
+                f" average time: {avg_io_time:.6f} seconds per frame"
+            )
+        if self._wait_time > 0:
+            self.logger.info(f"{class_name}: wait time: {self._wait_time:.6f} seconds")
+        if self._seq_time > 0:
+            avg_seq_time = self._seq_time / self._seq_num
+            self.logger.info(
+                f"{class_name}: seq num: {self._seq_num:.6f}, total time: {self._seq_time:.6f} seconds, average time:"
+                f" {avg_seq_time:.6f} seconds per sequence"
+            )

nimbus/components/data/observation.py

@@ -0,0 +1,119 @@
+import os
+
+import cv2
+import imageio
+import numpy as np
+
+from nimbus.components.data.camera import Camera
+
+
+class Observations:
+    """
+    Represents a single observation of a scene, which may include multiple camera trajectories and associated data.
+    Each observation is identified by a unique name and index, and can contain multiple Camera items that capture
+    different viewpoints or modalities of the same scene.
+
+    Args:
+        scene_name (str): The name of the scene associated with this observation.
+        index (str): The index or ID of this observation within the scene.
+        length (int): Optional total length of the observation. Calculated from camera trajectories if not provided.
+        data (dict): Optional dictionary for storing additional arbitrary data, such as metadata or annotations.
+    """
+
+    def __init__(self, scene_name: str, index: str, length: int = None, data: dict = None):
+        self.scene_name = scene_name
+        self.obs_name = scene_name + "_" + index
+        self.index = index
+        self.cam_items = []
+        self.length = length
+        self.data = data
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+
+    def append_cam(self, item: Camera):
+        self.cam_items.append(item)
+
+    def __len__(self):
+        if self.length is not None:
+            return self.length
+        self.length = 0
+        for cam in self.cam_items:
+            self.length += len(cam)
+        return self.length
+
+    def get_length(self):
+        return len(self)
+
+    def flush_to_disk(self, path, video_fps=10):
+        path_to_save = os.path.join(path, "trajectory_" + self.index)
+        print(f"obs {self.obs_name} try to save path in {path_to_save}")
+        os.makedirs(path_to_save, exist_ok=True)
+
+        # Single camera: save in root directory
+        if len(self.cam_items) == 1:
+            cam = self.cam_items[0]
+            self._save_camera_data(path_to_save, cam, video_fps)
+        # Multiple cameras: save in camera_0/, camera_1/, etc.
+        else:
+            for idx, cam in enumerate(self.cam_items):
+                camera_dir = os.path.join(path_to_save, f"camera_{idx}")
+                os.makedirs(camera_dir, exist_ok=True)
+                self._save_camera_data(camera_dir, cam, video_fps)
+
+    def _save_camera_data(self, save_dir, cam: Camera, video_fps):
+        """Helper method to save camera visualization data (rgbs, depths) to a directory."""
+        # Save RGB and depth images if available
+        if cam.rgbs is not None and len(cam.rgbs) > 0:
+            rgb_images_path = os.path.join(save_dir, "rgb/")
+            os.makedirs(rgb_images_path, exist_ok=True)
+
+            fps_path = os.path.join(save_dir, "fps.mp4")
+
+            for idx, rgb_item in enumerate(cam.rgbs):
+                rgb_filename = os.path.join(rgb_images_path, f"{idx}.jpg")
+                cv2.imwrite(rgb_filename, cv2.cvtColor(rgb_item, cv2.COLOR_BGR2RGB))
+
+            imageio.mimwrite(fps_path, cam.rgbs, fps=video_fps)
+
+        if cam.depths is not None and len(cam.depths) > 0:
+            depth_images_path = os.path.join(save_dir, "depth/")
+            os.makedirs(depth_images_path, exist_ok=True)
+
+            depth_path = os.path.join(save_dir, "depth.mp4")
+
+            # Create a copy for video (8-bit version)
+            depth_video_frames = []
+            for idx, depth_item in enumerate(cam.depths):
+                depth_filename = os.path.join(depth_images_path, f"{idx}.png")
+                cv2.imwrite(depth_filename, depth_item)
+                depth_video_frames.append((depth_item >> 8).astype(np.uint8))
+
+            imageio.mimwrite(depth_path, depth_video_frames, fps=video_fps)
+
+        # Save UV tracking visualizations if available
+        if cam.uv_tracks is not None and cam.uv_mesh_names is not None and cam.rgbs is not None:
+            num_frames = len(cam.rgbs)
+            try:
+                from nimbus_extension.components.render.brpc_utils.point_tracking import (
+                    make_uv_overlays_and_video,
+                )
+            except ImportError as e:
+                raise ImportError(
+                    "UV tracking visualization requires nimbus_extension. "
+                    "Please add `import nimbus_extension` before running the pipeline."
+                ) from e
+
+            make_uv_overlays_and_video(
+                cam.rgbs,
+                cam.uv_tracks,
+                cam.uv_mesh_names,
+                start_frame=0,
+                end_frame=num_frames,
+                fps=video_fps,
+                path_to_save=save_dir,
+            )

nimbus/components/data/package.py

@@ -0,0 +1,39 @@
+import pickle
+
+
+class Package:
+    """
+    A class representing a data package that can be serialized and deserialized for pipeline.
+
+    Args:
+        data: The actual data contained in the package, which can be of any type.
+        task_id (int): The ID of the task associated with this package.
+        task_name (str): The name of the task associated with this package.
+        stop_sig (bool): Whether this package signals the pipeline to stop.
+    """
+
+    def __init__(self, data, task_id: int = -1, task_name: str = None, stop_sig: bool = False):
+        self.is_ser = False
+        self.data = data
+        self.task_id = task_id
+        self.task_name = task_name
+        self.stop_sig = stop_sig
+
+    def serialize(self):
+        assert self.is_ser is False, "data is already serialized"
+        self.data = pickle.dumps(self.data)
+        self.is_ser = True
+
+    def deserialize(self):
+        assert self.is_ser is True, "data is already deserialized"
+        self.data = pickle.loads(self.data)
+        self.is_ser = False
+
+    def is_serialized(self):
+        return self.is_ser
+
+    def get_data(self):
+        return self.data
+
+    def should_stop(self):
+        return self.stop_sig is True

nimbus/components/data/scene.py

@@ -0,0 +1,69 @@
+class Scene:
+    """
+    Represents a loaded scene in the simulation environment, holding workflow context and task execution state.
+
+    Args:
+        name (str): The name of the scene or task.
+        pcd: Point cloud data associated with the scene.
+        scale (float): Scale factor for the scene geometry.
+        materials: Material data for the scene.
+        textures: Texture data for the scene.
+        floor_heights: Floor height information for the scene.
+        wf: The task workflow instance managing this scene.
+        task_id (int): The index of the current task within the workflow.
+        task_exec_num (int): The execution count for the current task, used for task repetition tracking.
+        simulation_app: The Isaac Sim SimulationApp instance.
+    """
+
+    def __init__(
+        self,
+        name: str = None,
+        pcd=None,
+        scale: float = 1.0,
+        materials=None,
+        textures=None,
+        floor_heights=None,
+        wf=None,
+        task_id: int = None,
+        task_exec_num: int = 1,
+        simulation_app=None,
+    ):
+        self.name = name
+        self.pcd = pcd
+        self.materials = materials
+        self.textures = textures
+        self.floor_heights = floor_heights
+        self.scale = scale
+        self.wf = wf
+        self.simulation_app = simulation_app
+        self.task_id = task_id
+        self.plan_info = None
+        self.generate_success = False
+        self.task_exec_num = task_exec_num
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state["pcd"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self.pcd = None
+
+    def add_plan_info(self, plan_info):
+        self.plan_info = plan_info
+
+    def flush_to_disk(self, path):
+        pass
+
+    def load_from_disk(self, path):
+        pass
+
+    def update_generate_status(self, success):
+        self.generate_success = success
+
+    def get_generate_status(self):
+        return self.generate_success
+
+    def update_task_exec_num(self, num):
+        self.task_exec_num = num

nimbus/components/data/sequence.py

@@ -0,0 +1,145 @@
+import json
+import os
+
+import numpy as np
+import open3d as o3d
+
+from nimbus.components.data.camera import C2W, Camera
+
+
+class Sequence:
+    """
+    Represents a camera trajectory sequence with associated metadata.
+
+    Args:
+        scene_name (str): The name of the scene (e.g., room identifier).
+        index (str): The index or ID of this sequence within the scene.
+        length (int): Optional explicit sequence length. Calculated from camera trajectories if not provided.
+        data (dict): Optional additional arbitrary data associated with the sequence.
+    """
+
+    def __init__(self, scene_name: str, index: str, length: int = None, data: dict = None):
+        self.scene_name = scene_name
+        self.seq_name = scene_name + "_" + index
+        self.index = index
+        self.cam_items: list[Camera] = []
+        self.path_pcd = None
+        self.length = length
+        self.data = data
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["path_pcd_color"] = np.asarray(state["path_pcd"].colors)
+        state["path_pcd"] = o3d.io.write_point_cloud_to_bytes(state["path_pcd"], "mem::xyz")
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self.path_pcd = o3d.io.read_point_cloud_from_bytes(state["path_pcd"], "mem::xyz")
+        self.path_pcd.colors = o3d.utility.Vector3dVector(state["path_pcd_color"])
+
+    def __len__(self):
+        if self.length is not None:
+            return self.length
+        self.length = 0
+        for cam in self.cam_items:
+            self.length += len(cam)
+        return self.length
+
+    def append_cam(self, item: Camera):
+        self.cam_items.append(item)
+
+    def update_pcd(self, path_pcd):
+        self.path_pcd = path_pcd
+
+    def get_length(self):
+        return len(self)
+
+    def flush_to_disk(self, path):
+        path_to_save = os.path.join(path, "trajectory_" + self.index)
+        print(f"seq {self.seq_name} try to save path in {path_to_save}")
+        os.makedirs(path_to_save, exist_ok=True)
+        if self.path_pcd is not None:
+            pcd_path = os.path.join(path_to_save, "path.ply")
+            o3d.io.write_point_cloud(pcd_path, self.path_pcd)
+
+        # Single camera: save in root directory
+        if len(self.cam_items) == 1:
+            cam = self.cam_items[0]
+            camera_trajectory_list = [t.matrix for t in cam.trajectory]
+            save_dict = {
+                "camera_intrinsic": cam.intrinsic if cam.intrinsic is not None else None,
+                "camera_extrinsic": cam.extrinsic if cam.extrinsic is not None else None,
+                "camera_trajectory": camera_trajectory_list,
+            }
+            traj_path = os.path.join(path_to_save, "data.json")
+            json_object = json.dumps(save_dict, indent=4)
+            with open(traj_path, "w", encoding="utf-8") as outfile:
+                outfile.write(json_object)
+        # Multiple cameras: save in camera_0/, camera_1/, etc.
+        else:
+            for idx, cam in enumerate(self.cam_items):
+                camera_dir = os.path.join(path_to_save, f"camera_{idx}")
+                os.makedirs(camera_dir, exist_ok=True)
+                camera_trajectory_list = [t.matrix for t in cam.trajectory]
+                save_dict = {
+                    "camera_intrinsic": cam.intrinsic if cam.intrinsic is not None else None,
+                    "camera_extrinsic": cam.extrinsic if cam.extrinsic is not None else None,
+                    "camera_trajectory": camera_trajectory_list,
+                }
+                traj_path = os.path.join(camera_dir, "data.json")
+                json_object = json.dumps(save_dict, indent=4)
+                with open(traj_path, "w", encoding="utf-8") as outfile:
+                    outfile.write(json_object)
+
+    def load_from_disk(self, path):
+        print(f"seq {self.seq_name} try to load path from {path}")
+
+        pcd_path = os.path.join(path, "path.ply")
+        if os.path.exists(pcd_path):
+            self.path_pcd = o3d.io.read_point_cloud(pcd_path)
+
+        # Clear existing camera items
+        self.cam_items = []
+
+        # Check if single camera format (data.json in root)
+        traj_path = os.path.join(path, "data.json")
+        if os.path.exists(traj_path):
+            with open(traj_path, "r", encoding="utf-8") as infile:
+                data = json.load(infile)
+
+            camera_trajectory_list = []
+            for trajectory in data["camera_trajectory"]:
+                camera_trajectory_list.append(C2W(matrix=trajectory))
+
+            cam = Camera(
+                trajectory=camera_trajectory_list,
+                intrinsic=data.get("camera_intrinsic"),
+                extrinsic=data.get("camera_extrinsic"),
+            )
+            self.cam_items.append(cam)
+        else:
+            # Multiple camera format (camera_0/, camera_1/, etc.)
+            idx = 0
+            while True:
+                camera_dir = os.path.join(path, f"camera_{idx}")
+                camera_json = os.path.join(camera_dir, "data.json")
+                if not os.path.exists(camera_json):
+                    break
+
+                with open(camera_json, "r", encoding="utf-8") as infile:
+                    data = json.load(infile)
+
+                camera_trajectory_list = []
+                for trajectory in data["camera_trajectory"]:
+                    camera_trajectory_list.append(C2W(matrix=trajectory))
+
+                cam = Camera(
+                    trajectory=camera_trajectory_list,
+                    intrinsic=data.get("camera_intrinsic"),
+                    extrinsic=data.get("camera_extrinsic"),
+                )
+                self.cam_items.append(cam)
+                idx += 1
+
+            assert len(self.cam_items) > 0, f"No camera data found in {path}"

nimbus/components/dedump/__init__.py

@@ -0,0 +1,7 @@
+from nimbus.components.data.iterator import Iterator
+
+dedumper_dict = {}
+
+
+def register(type_name: str, cls: Iterator):
+    dedumper_dict[type_name] = cls

nimbus/components/dump/__init__.py

@@ -0,0 +1,7 @@
+from .base_dumper import BaseDumper
+
+dumper_dict = {}
+
+
+def register(type_name: str, cls: BaseDumper):
+    dumper_dict[type_name] = cls

nimbus/components/dump/base_dumper.py

@@ -0,0 +1,82 @@
+import time
+from abc import abstractmethod
+
+from pympler import asizeof
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.package import Package
+from nimbus.utils.utils import unpack_iter_data
+
+
+class BaseDumper(Iterator):
+    def __init__(self, data_iter, output_queue, max_queue_num=1):
+        super().__init__()
+        self.data_iter = data_iter
+        self.scene = None
+        self.output_queue = output_queue
+        self.total_case = 0
+        self.success_case = 0
+        self.max_queue_num = max_queue_num
+
+    def __iter__(self):
+        return self
+
+    def _next(self):
+        try:
+            data = next(self.data_iter)
+            scene, seq, obs = unpack_iter_data(data)
+            self.total_case += 1
+            if scene is not None:
+                if self.scene is not None and (
+                    scene.task_id != self.scene.task_id
+                    or scene.name != self.scene.name
+                    or scene.task_exec_num != self.scene.task_exec_num
+                ):
+                    self.logger.info(
+                        f"Scene {self.scene.name} generate finish, success rate: {self.success_case}/{self.total_case}"
+                    )
+                    self.total_case = 1
+                    self.success_case = 0
+                self.scene = scene
+            if obs is None and seq is None:
+                self.logger.info(f"generate failed, skip once! success rate: {self.success_case}/{self.total_case}")
+                if self.scene is not None:
+                    self.scene.update_generate_status(success=False)
+                return None
+            io_start_time = time.time()
+            if self.output_queue is not None:
+                obj = self.dump(seq, obs)
+                pack = Package(obj, task_id=scene.task_id, task_name=scene.name)
+                pack.serialize()
+
+                wait_time = time.time()
+                while self.output_queue.qsize() >= self.max_queue_num:
+                    time.sleep(1)
+                end_time = time.time()
+                self.collect_wait_time_info(end_time - wait_time)
+
+                st = time.time()
+                self.output_queue.put(pack)
+                ed = time.time()
+                self.logger.info(f"put time: {ed - st}, data size: {asizeof.asizeof(obj)}")
+            else:
+                obj = self.dump(seq, obs)
+            self.success_case += 1
+            self.scene.update_generate_status(success=True)
+            self.collect_seq_info(1, time.time() - io_start_time)
+        except StopIteration:
+            if self.output_queue is not None:
+                pack = Package(None, stop_sig=True)
+                self.output_queue.put(pack)
+            if self.scene is not None:
+                self.logger.info(
+                    f"Scene {self.scene.name} generate finish, success rate: {self.success_case}/{self.total_case}"
+                )
+            raise StopIteration("no data")
+        except Exception as e:
+            self.logger.exception(f"Error during data dumping: {e}")
+            raise e
+
+    @abstractmethod
+    def dump(self, seq, obs):
+        raise NotImplementedError("This method should be overridden by subclasses")

nimbus/components/load/__init__.py

@@ -0,0 +1,16 @@
+# flake8: noqa: F401
+# pylint: disable=C0413
+
+from .base_randomizer import LayoutRandomizer
+from .base_scene_loader import SceneLoader
+
+scene_loader_dict = {}
+layout_randomizer_dict = {}
+
+
+def register_loader(type_name: str, cls: SceneLoader):
+    scene_loader_dict[type_name] = cls
+
+
+def register_randomizer(type_name: str, cls: LayoutRandomizer):
+    layout_randomizer_dict[type_name] = cls

nimbus/components/load/base_randomizer.py

@@ -0,0 +1,72 @@
+import sys
+import time
+from abc import abstractmethod
+from typing import Optional
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.scene import Scene
+from nimbus.daemon.decorators import status_monitor
+
+
+class LayoutRandomizer(Iterator):
+    """
+    Base class for layout randomization in a scene. This class defines the structure for randomizing scenes and
+    tracking the randomization process. It manages the current scene, randomization count, and provides hooks for
+    subclasses to implement specific randomization logic.
+
+    Args:
+        scene_iter (Iterator): An iterator that provides scenes to be randomized.
+        random_num (int): The number of randomizations to perform for each scene before moving to the next one.
+        strict_mode (bool): If True, the randomizer will check the generation status of the current scene and retry
+                            randomization if it was not successful. This ensures that only successfully generated
+                            scenes are counted towards the randomization limit.
+    """
+
+    def __init__(self, scene_iter: Iterator, random_num: int, strict_mode: bool = False):
+        super().__init__()
+        self.scene_iter = scene_iter
+        self.random_num = random_num
+        self.strict_mode = strict_mode
+        self.cur_index = sys.maxsize
+        self.scene: Optional[Scene] = None
+
+    def reset(self, scene):
+        self.cur_index = 0
+        self.scene = scene
+
+    def _fetch_next_scene(self):
+        scene = next(self.scene_iter)
+        self.reset(scene)
+
+    @status_monitor()
+    def _randomize_with_status(self, scene) -> Scene:
+        scene = self.randomize_scene(self.scene)
+        return scene
+
+    def _next(self) -> Scene:
+        try:
+            if self.strict_mode and self.scene is not None:
+                if not self.scene.get_generate_status():
+                    self.logger.info("strict_mode is open, retry the randomization to generate sequence.")
+                    st = time.time()
+                    scene = self._randomize_with_status(self.scene)
+                    self.collect_seq_info(1, time.time() - st)
+                    return scene
+            if self.cur_index >= self.random_num:
+                self._fetch_next_scene()
+            if self.cur_index < self.random_num:
+                st = time.time()
+                scene = self._randomize_with_status(self.scene)
+                self.collect_seq_info(1, time.time() - st)
+                self.cur_index += 1
+            return scene
+        except StopIteration:
+            raise StopIteration("No more scenes to randomize.")
+        except Exception as e:
+            self.logger.exception(f"Error during scene idx {self.cur_index} randomization: {e}")
+            self.cur_index += 1
+            raise e
+
+    @abstractmethod
+    def randomize_scene(self, scene) -> Scene:
+        raise NotImplementedError("This method should be overridden by subclasses")

nimbus/components/load/base_scene_loader.py

@@ -0,0 +1,41 @@
+from abc import abstractmethod
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.scene import Scene
+
+
+class SceneLoader(Iterator):
+    """
+    Base class for scene loading in a simulation environment. This class defines the structure for loading scenes
+    and tracking the loading process. It manages the current package iterator and provides hooks for subclasses
+    to implement specific scene loading logic.
+
+    Args:
+        pack_iter (Iterator): An iterator that provides packages containing scene information to be loaded.
+    """
+
+    def __init__(self, pack_iter):
+        super().__init__()
+        self.pack_iter = pack_iter
+
+    @abstractmethod
+    def load_asset(self) -> Scene:
+        """
+        Abstract method to load and initialize a scene.
+
+        Subclasses must implement this method to define the specific logic for creating and configuring
+        a scene object based on the current state of the iterator.
+
+        Returns:
+            Scene: A fully initialized Scene object.
+        """
+        raise NotImplementedError("This method must be implemented by subclasses")
+
+    def _next(self) -> Scene:
+        try:
+            return self.load_asset()
+        except StopIteration:
+            raise StopIteration("No more scenes to load.")
+        except Exception as e:
+            self.logger.exception(f"Error during scene loading: {e}")
+            raise e

nimbus/components/plan_with_render/__init__.py

@@ -0,0 +1,7 @@
+from nimbus.components.data.iterator import Iterator
+
+plan_with_render_dict = {}
+
+
+def register(type_name: str, cls: Iterator):
+    plan_with_render_dict[type_name] = cls

nimbus/components/planner/__init__.py

@@ -0,0 +1,7 @@
+from .base_seq_planner import SequencePlanner
+
+seq_planner_dict = {}
+
+
+def register(type_name: str, cls: SequencePlanner):
+    seq_planner_dict[type_name] = cls

nimbus/components/planner/base_seq_planner.py

@@ -0,0 +1,102 @@
+import sys
+import time
+from abc import abstractmethod
+from typing import Optional
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.scene import Scene
+from nimbus.components.data.sequence import Sequence
+from nimbus.daemon.decorators import status_monitor
+from nimbus.utils.flags import is_debug_mode
+from nimbus.utils.types import ARGS, TYPE
+
+from .planner import path_planner_dict
+
+
+class SequencePlanner(Iterator):
+    """
+    A base class for sequence planning in a simulation environment. This class defines the structure for generating
+    sequences based on scenes and tracking the planning process. It manages the current scene, episode count
+    and provides hooks for subclasses to implement specific sequence generation logic.
+
+    Args:
+        scene_iter (Iterator): An iterator that provides scenes to be processed for sequence planning.
+        planner_cfg (dict): A dictionary containing configuration parameters for the planner,
+            such as the type of planner to use and its arguments.
+        episodes (int): The number of episodes to generate for each scene before moving to the next one. Default is 1.
+    """
+
+    def __init__(self, scene_iter: Iterator[Scene], planner_cfg: dict, episodes: int = 1):
+        super().__init__()
+        self.scene_iter = scene_iter
+        self.planner_cfg = planner_cfg
+        self.episodes = episodes
+        self.current_episode = sys.maxsize
+        self.scene: Optional[Scene] = None
+
+    @status_monitor()
+    def _plan_with_status(self) -> Optional[Sequence]:
+        seq = self.generate_sequence()
+        return seq
+
+    def _next(self) -> tuple[Scene, Sequence]:
+        try:
+            if self.scene is None or self.current_episode >= self.episodes:
+                try:
+                    self.scene = next(self.scene_iter)
+                    self.current_episode = 0
+                    if self.scene is None:
+                        return None, None
+                    self.initialize(self.scene)
+                except StopIteration:
+                    raise StopIteration("No more scene to process.")
+                except Exception as e:
+                    self.logger.exception(f"Error loading next scene: {e}")
+                    if is_debug_mode():
+                        raise e
+                    self.current_episode = sys.maxsize
+                    return None, None
+
+            while True:
+                compute_start_time = time.time()
+                seq = self._plan_with_status()
+                compute_end_time = time.time()
+                self.current_episode += 1
+
+                if seq is not None:
+                    self.collect_compute_frame_info(seq.get_length(), compute_end_time - compute_start_time)
+                    return self.scene, seq
+
+                if self.current_episode >= self.episodes:
+                    return self.scene, None
+
+                self.logger.info(f"Generate seq failed and retry. Current episode id is {self.current_episode}")
+
+        except StopIteration:
+            raise StopIteration("No more scene to process.")
+        except Exception as e:
+            scene_name = getattr(self.scene, "name", "<unknown>")
+            self.logger.exception(
+                f"Error during idx {self.current_episode} sequence generation for scene {scene_name}: {e}"
+            )
+            if is_debug_mode():
+                raise e
+            self.current_episode += 1
+            return self.scene, None
+
+    @abstractmethod
+    def generate_sequence(self) -> Optional[Sequence]:
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    def _initialize(self, scene):
+        if self.planner_cfg is not None:
+            self.logger.info(f"init {self.planner_cfg[TYPE]} planner in seq_planner")
+            self.planner = path_planner_dict[self.planner_cfg[TYPE]](scene, **self.planner_cfg.get(ARGS, {}))
+        else:
+            self.planner = None
+            self.logger.info("planner config is None in seq_planner and skip initialize")
+
+    def initialize(self, scene):
+        init_start_time = time.time()
+        self._initialize(scene)
+        self.record_init_time(time.time() - init_start_time)

nimbus/components/planner/planner/__init__.py

@@ -0,0 +1,5 @@
+path_planner_dict = {}
+
+
+def register(type_name: str, cls):
+    path_planner_dict[type_name] = cls

nimbus/components/render/__init__.py

@@ -0,0 +1,7 @@
+from .base_renderer import BaseRenderer
+
+renderer_dict = {}
+
+
+def register(type_name: str, cls: BaseRenderer):
+    renderer_dict[type_name] = cls

nimbus/components/render/base_renderer.py

@@ -0,0 +1,80 @@
+import time
+from abc import abstractmethod
+from typing import Optional
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.observation import Observations
+from nimbus.components.data.scene import Scene
+from nimbus.components.data.sequence import Sequence
+from nimbus.daemon.decorators import status_monitor
+
+
+class BaseRenderer(Iterator):
+    """
+    Base class for rendering in a simulation environment. This class defines the structure for rendering scenes and
+    tracking the rendering process. It manages the current scene and provides hooks for subclasses to implement
+    specific rendering logic.
+
+    Args:
+        scene_seq_iter (Iterator): An iterator that provides pairs of scenes and sequences to be rendered. Each item
+                                  from the iterator should be a tuple containing a scene and its corresponding sequence.
+    """
+
+    def __init__(self, scene_seq_iter: Iterator[tuple[Scene, Sequence]]):
+        super().__init__()
+        self.scene_seq_iter = scene_seq_iter
+        self.scene: Optional[Scene] = None
+
+    @status_monitor()
+    def _generate_obs_with_status(self, seq) -> Optional[Observations]:
+        compute_start_time = time.time()
+        obs = self.generate_obs(seq)
+        end_start_time = time.time()
+        if obs is not None:
+            self.collect_compute_frame_info(len(obs), end_start_time - compute_start_time)
+        return obs
+
+    def _next(self):
+        try:
+            scene, seq = next(self.scene_seq_iter)
+            if scene is not None:
+                if self.scene is None:
+                    self.reset(scene)
+                elif scene.task_id != self.scene.task_id or scene.name != self.scene.name:
+                    self.logger.info(f"Scene changed: {self.scene.name} -> {scene.name}")
+                    self.reset(scene)
+            if seq is None:
+                return scene, None, None
+            obs = self._generate_obs_with_status(seq)
+            if obs is None:
+                return scene, None, None
+            return scene, seq, obs
+        except StopIteration:
+            raise StopIteration("No more sequences to process.")
+        except Exception as e:
+            self.logger.exception(f"Error during rendering: {e}")
+            raise e
+
+    @abstractmethod
+    def generate_obs(self, seq) -> Optional[Observations]:
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    @abstractmethod
+    def _lazy_init(self):
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    @abstractmethod
+    def _close_resource(self):
+        raise NotImplementedError("This method should be overridden by subclasses")
+
+    def reset(self, scene):
+        try:
+            self.scene = scene
+            self._close_resource()
+            init_start_time = time.time()
+            self._lazy_init()
+            self.record_init_time(time.time() - init_start_time)
+        except Exception as e:
+            self.logger.exception(f"Error initializing renderer: {e}")
+            self.scene = None
+            raise e

nimbus/components/store/__init__.py

@@ -0,0 +1,7 @@
+from .base_writer import BaseWriter
+
+writer_dict = {}
+
+
+def register(type_name: str, cls: BaseWriter):
+    writer_dict[type_name] = cls

nimbus/components/store/base_writer.py

@@ -0,0 +1,163 @@
+import time
+from abc import abstractmethod
+from concurrent.futures import ThreadPoolExecutor
+from copy import copy
+
+from nimbus.components.data.iterator import Iterator
+from nimbus.components.data.observation import Observations
+from nimbus.components.data.scene import Scene
+from nimbus.components.data.sequence import Sequence
+from nimbus.daemon import ComponentStatus, StatusReporter
+from nimbus.utils.flags import is_debug_mode
+from nimbus.utils.utils import unpack_iter_data
+
+
+def run_batch(func, args):
+    for arg in args:
+        func(*arg)
+
+
+class BaseWriter(Iterator):
+    """
+    A base class for writing generated sequences and observations to disk. This class defines the structure for
+    writing data and tracking the writing process. It manages the current scene, success and total case counts,
+    and provides hooks for subclasses to implement specific data writing logic. The writer supports both synchronous
+    and asynchronous batch writing modes, allowing for efficient data handling in various scenarios.
+
+    Args:
+        data_iter (Iterator): An iterator that provides data to be written, typically containing scenes,
+            sequences, and observations.
+        seq_output_dir (str): The directory where generated sequences will be saved. Can be None
+            if sequence output is not needed.
+        obs_output_dir (str): The directory where generated observations will be saved. Can be None
+            if observation output is not needed.
+        batch_async (bool): If True, the writer will use asynchronous batch writing to improve performance
+            when handling large amounts of data. Default is True.
+        async_threshold (int): The maximum number of asynchronous write operations that can be in progress
+            at the same time. If the threshold is reached, the writer will wait for the oldest operation
+            to complete before starting a new one. Default is 1.
+        batch_size (int): The number of data items to write in each batch when using asynchronous writing.
+            Default is 2, and it will be capped at 8 to prevent potential issues with too many concurrent operations.
+    """
+
+    def __init__(
+        self,
+        data_iter: Iterator[tuple[Scene, Sequence, Observations]],
+        seq_output_dir: str,
+        obs_output_dir: str,
+        batch_async: bool = True,
+        async_threshold: int = 1,
+        batch_size: int = 2,
+    ):
+        super().__init__()
+        assert (
+            seq_output_dir is not None or obs_output_dir is not None
+        ), "At least one output directory must be provided"
+        self.data_iter = data_iter
+        self.seq_output_dir = seq_output_dir
+        self.obs_output_dir = obs_output_dir
+        self.scene = None
+        self.async_mode = batch_async
+        self.batch_size = batch_size if batch_size <= 8 else 8
+        if batch_async and batch_size > self.batch_size:
+            self.logger.info("Batch size is larger than 8(probably cause program hang), batch size will be set to 8")
+        self.async_threshold = async_threshold
+        self.flush_executor = ThreadPoolExecutor(max_workers=max(1, 64 // self.batch_size))
+        self.flush_threads = []
+        self.data_buffer = []
+        self.logger.info(
+            f"Batch Async Write Mode: {self.async_mode}, async threshold: {self.async_threshold}, batch size:"
+            f" {self.batch_size}"
+        )
+        self.total_case = 0
+        self.success_case = 0
+        self.last_scene_key = None
+        self.status_reporter = StatusReporter(self.__class__.__name__)
+
+    def _next(self):
+        try:
+            data = next(self.data_iter)
+            scene, seq, obs = unpack_iter_data(data)
+
+            new_key = (scene.task_id, scene.name, scene.task_exec_num) if scene is not None else None
+
+            self.scene = scene
+
+            if new_key != self.last_scene_key:
+                if self.scene is not None and self.last_scene_key is not None:
+                    self.logger.info(
+                        f"Scene {self.scene.name} generate finish, success rate: {self.success_case}/{self.total_case}"
+                    )
+                self.success_case = 0
+                self.total_case = 0
+                self.last_scene_key = new_key
+
+            if self.scene is None:
+                return None
+
+            self.total_case += 1
+
+            self.status_reporter.update_status(ComponentStatus.RUNNING)
+            if seq is None and obs is None:
+                self.logger.info(f"generate failed, skip once! success rate: {self.success_case}/{self.total_case}")
+                self.scene.update_generate_status(success=False)
+                return None
+            scene_name = self.scene.name
+            io_start_time = time.time()
+            if self.async_mode:
+                cp_start_time = time.time()
+                cp = copy(self.scene.wf)
+                cp_end_time = time.time()
+                if self.scene.wf is not None:
+                    self.logger.info(f"Scene {scene_name} workflow copy time: {cp_end_time - cp_start_time:.2f}s")
+                self.data_buffer.append((cp, scene_name, seq, obs))
+                if len(self.data_buffer) >= self.batch_size:
+                    self.flush_threads = [t for t in self.flush_threads if not t.done()]
+
+                    if len(self.flush_threads) >= self.async_threshold:
+                        self.logger.info("Max async workers reached, waiting for the oldest thread to finish")
+                        self.flush_threads[0].result()
+                        self.flush_threads = self.flush_threads[1:]
+
+                    to_flush_buffer = self.data_buffer.copy()
+                    async_flush = self.flush_executor.submit(run_batch, self.flush_to_disk, to_flush_buffer)
+                    if is_debug_mode():
+                        async_flush.result()  # surface exceptions immediately in debug mode
+                    self.flush_threads.append(async_flush)
+                    self.data_buffer = []
+                flush_length = len(obs) if obs is not None else len(seq)
+            else:
+                flush_length = self.flush_to_disk(self.scene.wf, scene_name, seq, obs)
+            self.success_case += 1
+            self.scene.update_generate_status(success=True)
+            self.collect_io_frame_info(flush_length, time.time() - io_start_time)
+            self.status_reporter.update_status(ComponentStatus.COMPLETED)
+            return None
+        except StopIteration:
+            if self.async_mode:
+                if len(self.data_buffer) > 0:
+                    async_flush = self.flush_executor.submit(run_batch, self.flush_to_disk, self.data_buffer)
+                    self.flush_threads.append(async_flush)
+                for thread in self.flush_threads:
+                    thread.result()
+            if self.scene is not None:
+                self.logger.info(
+                    f"Scene {self.scene.name} generate finish, success rate: {self.success_case}/{self.total_case}"
+                )
+            raise StopIteration("no data")
+        except Exception as e:
+            self.logger.exception(f"Error during data writing: {e}")
+            raise e
+
+    def __del__(self):
+        for thread in self.flush_threads:
+            thread.result()
+        self.logger.info(f"Writer {len(self.flush_threads)} threads closed")
+        # Close the simulation app if it exists
+        if self.scene is not None and self.scene.simulation_app is not None:
+            self.logger.info("Closing simulation app")
+            self.scene.simulation_app.close()
+
+    @abstractmethod
+    def flush_to_disk(self, task, scene_name, seq, obs):
+        raise NotImplementedError("This method should be overridden by subclasses")