Randomizer no longer randomize on each frame if cosmos is enabled

icar_sim_sdg/entrypoint.py

@@ -197,10 +197,10 @@ def close(self) -> None:
 
     randomizer = base_randomizer.Randomizer()
 
-    cam_randomizer = camera_randomizer.CameraRandomizer()
-    robot_anim_randomizer = robot_randomizer.RobotRandomizer()
-    scene_props_randomizer = props_randomizer.PropsRandomizer()
-    scene_lights_randomizer = lights_randomizer.LightsRandomizer()
+    cam_randomizer = camera_randomizer.CameraRandomizer(randomize_all_frames=not args.cosmos)
+    robot_anim_randomizer = robot_randomizer.RobotRandomizer(randomize_all_frames=not args.cosmos)
+    scene_props_randomizer = props_randomizer.PropsRandomizer(randomize_all_frames=not args.cosmos)
+    scene_lights_randomizer = lights_randomizer.LightsRandomizer(randomize_all_frames=not args.cosmos)
 
     randomizer.add_randomizer(cam_randomizer)
     randomizer.add_randomizer(robot_anim_randomizer)

icar_sim_sdg/randomizers/camera_randomizer.py

@@ -46,8 +46,17 @@
 class CameraRandomizer:
     """Randomizer for camera position and orientation in the scene."""
 
-    def __init__(self) -> None:
-        """Initialize the camera randomizer with configuration and camera setup."""
+    def __init__(self, randomize_all_frames: bool = True) -> None:
+        """Initialize the camera randomizer with configuration and camera setup.
+
+        Args:
+        ----
+            randomize_all_frames: If True, randomizes camera pose on every trigger.
+                If False, randomizes camera pose once during initialization.
+
+        """
+        self.randomize_all_frames = randomize_all_frames
+
         self.config = get_config("camera")
         self.camera, camera_prim_path = create_camera(
             self.config.get("position", DEFAULT_CAMERA_POSITION),
@@ -76,9 +85,13 @@ def __init__(self) -> None:
         self.altitude_range = self.config.get("camera_altitude", [0.0, 60.0])  # in degrees
         self.max_sampling_attempts = 5  # Max attempts to find valid position
 
+        if not self.randomize_all_frames:
+            self.update_camera_pose()
+
     def randomize(self) -> None:
         """Execute camera randomization by updating its pose."""
-        self.update_camera_pose()
+        if self.randomize_all_frames:
+            self.update_camera_pose()
 
     def update_camera_pose(
         self,

icar_sim_sdg/randomizers/lights_randomizer.py

@@ -28,38 +28,64 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import numpy as np
+import omni
 import omni.replicator.core as rep
+from omni.replicator.core.functional.get import light as f_light
+from omni.replicator.core.functional.modify import attribute as f_attribute
+from pxr import Gf
 from randomizers.utils import get_config
 
 
 class LightsRandomizer:
-    """Randomizer for robot animation and pose."""
+    """Randomizer for scene lighting attributes."""
 
-    def __init__(self) -> None:
-        """Initialize the lights randomizer with animation and pose randomization."""
+    def __init__(self, randomize_all_frames: bool = True) -> None:
+        """Initialize the lights randomizer with lighting configuration.
+
+        Args:
+        ----
+            randomize_all_frames: If True, randomizes light attributes on every trigger.
+                If False, randomizes light attributes once during initialization.
+
+        """
         lights_config = get_config("lights")
 
         if lights_config.get("enable_lights_randomization", False) is False:
             return
 
-        # Get spawned lights
-        self.lights = rep.get.light()
-
         # Configuration parameters
         self.enable_color_temperature = lights_config.get("enable_color_temperature", True)
         self.min_color = lights_config.get("min_color", [0.0, 0.0, 0.0])
         self.max_color = lights_config.get("max_color", [1.0, 1.0, 1.0])
         self.min_exposure = lights_config.get("min_exposure", 0.0)
         self.max_exposure = lights_config.get("max_exposure", 0.5)
 
-        def randomize_lights() -> None:
-            """Randomize lights attributes."""
-            with self.lights:
-                rep.modify.attribute("inputs:color", rep.distribution.uniform(self.min_color, self.max_color))
-                rep.modify.attribute("inputs:enableColorTemperature", self.enable_color_temperature)
-                rep.modify.attribute("inputs:exposure", rep.distribution.uniform(self.min_exposure, self.max_exposure))
+        if randomize_all_frames:
+            # Get spawned lights
+            self.lights = rep.get.light()
+
+            def randomize_lights() -> None:
+                """Randomize lights attributes."""
+                with self.lights:
+                    rep.modify.attribute("inputs:color", rep.distribution.uniform(self.min_color, self.max_color))
+                    rep.modify.attribute("inputs:enableColorTemperature", self.enable_color_temperature)
+                    rep.modify.attribute(
+                        "inputs:exposure", rep.distribution.uniform(self.min_exposure, self.max_exposure)
+                    )
 
-        rep.randomizer.register(randomize_lights)
+            rep.randomizer.register(randomize_lights)
 
-        with rep.trigger.on_frame():
-            randomize_lights()
+            with rep.trigger.on_frame():
+                randomize_lights()
+        else:
+            stage = omni.usd.get_context().get_stage()
+            lights = f_light(stage=stage)
+            for light in lights:
+                f_attribute(
+                    light, "inputs:color", Gf.Vec3f(*np.random.uniform(low=self.min_color, high=self.max_color))
+                )
+                f_attribute(light, "inputs:enableColorTemperature", self.enable_color_temperature)
+                f_attribute(
+                    light, "inputs:exposure", float(np.random.uniform(low=self.min_exposure, high=self.max_exposure))
+                )

icar_sim_sdg/randomizers/props_randomizer.py

@@ -30,7 +30,11 @@
 
 import pathlib
 
+import numpy as np
 import omni.replicator.core as rep
+from omni.replicator.core.functional.create import plane as f_plane
+from omni.replicator.core.functional.create import reference as f_reference
+from omni.replicator.core.functional.randomizer import scatter_2d as f_scatter_2d
 from randomizers.utils import get_config
 
 PROPS_PATH = pathlib.Path(__file__).parent.parent.parent / "assets" / "props"
@@ -39,9 +43,16 @@
 class PropsRandomizer:
     """Randomizer for scattering props on a surface in the scene."""
 
-    def __init__(self) -> None:
-        """Initialize the props randomizer with configuration."""
-        self.surface = rep.create.plane(position=(0, 0, 0), scale=0.75, visible=False)
+    def __init__(self, randomize_all_frames: bool = True) -> None:
+        """Initialize the props randomizer with configuration.
+
+        Args:
+        ----
+            randomize_all_frames: If True, randomizes prop positions and rotations on every trigger.
+                If False, randomizes prop positions and rotations once during initialization.
+
+        """
+        self.randomize_all_frames = randomize_all_frames
 
         self.config = get_config("props")
         self.num_props = self.config.get("num_props", 10)
@@ -70,15 +81,23 @@ def __init__(self) -> None:
         self.rotation_min, self.rotation_max = (self.rotation_range["z"][0],), (self.rotation_range["z"][1],)
         self.usds = rep.utils.get_usd_files(str(PROPS_PATH), recursive=True)
 
-        def randomize_props() -> None:
-            """Randomize prop positions and rotations."""
-            with rep.randomizer.instantiate(self.usds, size=self.num_props, semantics=[("class", "prop")]):
-                rep.randomizer.scatter_2d(
-                    check_for_collisions=True,
-                    surface_prims=self.surface,
-                )
+        if self.randomize_all_frames:
+            surface = rep.create.plane(position=(0, 0, 0), scale=0.75, visible=False)
+
+            def randomize_props() -> None:
+                """Randomize prop positions and rotations."""
+                with rep.randomizer.instantiate(self.usds, size=self.num_props, semantics=[("class", "prop")]):
+                    rep.randomizer.scatter_2d(
+                        check_for_collisions=True,
+                        surface_prims=surface,
+                    )
 
-        rep.randomizer.register(randomize_props)
+            rep.randomizer.register(randomize_props)
 
-        with rep.trigger.on_frame():
-            randomize_props()
+            with rep.trigger.on_frame():
+                randomize_props()
+        else:
+            surface = f_plane(position=(0, 0, 0), scale=0.75, visible=False)
+            random_idx = np.random.randint(0, len(self.usds) - 1, size=self.num_props)
+            props_prims = [f_reference(usd_path=self.usds[idx], semantics={"class": "prop"}) for idx in random_idx]
+            f_scatter_2d(prims=props_prims, surface_prims=surface, check_for_collisions=True)

icar_sim_sdg/randomizers/robot_randomizer.py

@@ -30,27 +30,61 @@
 
 import pathlib
 
+import numpy as np
 import omni.replicator.core as rep
+import omni.timeline
+from omni.replicator.core.functional.create import reference as f_reference
 
 ROBOT_PATH = pathlib.Path(__file__).parent.parent.parent / "assets" / "ur10e" / "pick_animation" / "playback.usda"
+FIRST_MOVING_FRAME = 80
 NUM_FRAMES = 642  # Animation frames
 
 
 class RobotRandomizer:
     """Randomizer for robot animation and pose."""
 
-    def __init__(self) -> None:
-        """Initialize the robot randomizer with animation and pose randomization."""
-        # Spawn robot pick animation
-        self.robot = rep.create.from_usd(str(ROBOT_PATH), semantics={"class": "robot"})
+    def __init__(self, randomize_all_frames: bool = True) -> None:
+        """Initialize the robot randomizer with animation and pose randomization.
 
-        def randomize_robot_pose() -> None:
-            """Randomize robot animation frame and rotation."""
-            with self.robot:
-                rep.modify.timeline(rep.distribution.uniform(lower=0, upper=NUM_FRAMES), modify_type="frame")
-                rep.modify.pose(rotation_z=rep.distribution.uniform(0.0, 360.0))
+        Args:
+        ----
+            randomize_all_frames: If True, randomizes animation frame and rotation on every trigger.
+                If False, incrementally advances through animation frames.
 
-        rep.randomizer.register(randomize_robot_pose)
+        """
+        self.randomize_all_frames = randomize_all_frames
+        self.timeline_iface = omni.timeline.acquire_timeline_interface()
+        self.current_frame = FIRST_MOVING_FRAME
 
-        with rep.trigger.on_frame():
-            randomize_robot_pose()
+        if self.randomize_all_frames:
+            # Spawn robot pick animation
+            self.robot = rep.create.from_usd(str(ROBOT_PATH), semantics={"class": "robot"})
+
+            def randomize_robot_pose() -> None:
+                """Randomize robot animation frame and rotation."""
+                with self.robot:
+                    rep.modify.timeline(rep.distribution.uniform(lower=0, upper=NUM_FRAMES), modify_type="frame")
+                    rep.modify.pose(rotation_z=rep.distribution.uniform(0.0, 360.0))
+
+            rep.randomizer.register(randomize_robot_pose)
+
+            with rep.trigger.on_frame():
+                randomize_robot_pose()
+        else:
+            random_rotation = np.random.uniform(0.0, 360.0)
+            f_reference(
+                usd_path=str(ROBOT_PATH),
+                position=(0, 0, 0),
+                rotation=(0, 0, random_rotation),
+                semantics={"class": "robot"},
+            )
+
+    def randomize(self) -> None:
+        """Apply randomization to the robot."""
+        if self.randomize_all_frames:
+            return
+        fps = float(self.timeline_iface.get_time_codes_per_seconds())
+        self.timeline_iface.set_current_time(self.current_frame / fps)
+        self.current_frame += 1
+        if self.current_frame >= NUM_FRAMES:
+            self.current_frame = FIRST_MOVING_FRAME