181 lines
6.9 KiB
Python
181 lines
6.9 KiB
Python
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
|