feat: support optional rng for deterministic scene placement sampling

rlbench/backend/scene.py

@@ -31,7 +31,8 @@ def __init__(self,
                  pyrep: PyRep,
                  robot: Robot,
                  obs_config: ObservationConfig = ObservationConfig(),
-                 robot_setup: str = 'panda'):
+                 robot_setup: str = 'panda',
+                 rng=None):
         self.pyrep = pyrep
         self.robot = robot
         self.robot_setup = robot_setup
@@ -40,8 +41,9 @@ def __init__(self,
         self._initial_task_state = None
         self._start_arm_joint_pos = robot.arm.get_joint_positions()
         self._starting_gripper_joint_pos = robot.gripper.get_joint_positions()
+        self._rng = rng
         self._workspace = Shape('workspace')
-        self._workspace_boundary = SpawnBoundary([self._workspace])
+        self._workspace_boundary = SpawnBoundary([self._workspace], rng=rng)
         self._cam_over_shoulder_left = VisionSensor('cam_over_shoulder_left')
         self._cam_over_shoulder_right = VisionSensor('cam_over_shoulder_right')
         self._cam_overhead = VisionSensor('cam_overhead')

rlbench/backend/spawn_boundary.py

@@ -49,10 +49,13 @@ def within_boundary(self, boundary: 'BoundingBox', is_plane: bool) -> bool:
 
 class BoundaryObject(object):
 
-    def __init__(self, boundary: Object):
+    def __init__(self, boundary: Object, rng=None):
         self._boundary = boundary
         self._is_plane = False
         self._contained_objects = []
+        # An optional source of randomness exposing the numpy RandomState /
+        # Generator interface (uniform, choice). Defaults to np.random.
+        self._rng = rng if rng is not None else np.random
 
         if boundary.is_model():
             (minx, maxx, miny,
@@ -69,16 +72,16 @@ def __init__(self, boundary: Object):
 
     def _get_position_within_boundary(self, obj: Object, obj_bbox: BoundingBox
                                       ) -> List[float]:
-        x = np.random.uniform(
+        x = self._rng.uniform(
             self._boundary_bbox.min_x + np.abs(obj_bbox.min_x),
             self._boundary_bbox.max_x - np.abs(obj_bbox.max_x))
-        y = np.random.uniform(
+        y = self._rng.uniform(
             self._boundary_bbox.min_y + np.abs(obj_bbox.min_y),
             self._boundary_bbox.max_y - np.abs(obj_bbox.max_y))
         if self._is_plane:
             _, _, z = obj.get_position(self._boundary)
         else:
-            z = np.random.uniform(
+            z = self._rng.uniform(
                 self._boundary_bbox.min_z + np.abs(obj_bbox.min_z),
                 self._boundary_bbox.max_z - np.abs(obj_bbox.max_z))
         return [x, y, z]
@@ -100,7 +103,7 @@ def add(self, obj: Object, ignore_collisions: bool = False,
         else:
             bb = obj.get_bounding_box()
         obj_bbox = BoundingBox(*bb)
-        rotation = np.random.uniform(list(min_rotation), list(max_rotation))
+        rotation = self._rng.uniform(list(min_rotation), list(max_rotation))
         obj_bbox = obj_bbox.rotate(rotation)
 
         if not obj_bbox.within_boundary(self._boundary_bbox, self._is_plane):
@@ -135,11 +138,21 @@ class SpawnBoundary(object):
 
     MAX_SAMPLES = 100
 
-    def __init__(self, boundaries: List[Object]):
+    def __init__(self, boundaries: List[Object], rng=None):
+        """Creates a spawner for sampling object placements within boundaries.
+
+        :param boundaries: The boundaries to sample placements within.
+        :param rng: Optional source of randomness exposing the numpy
+            RandomState/Generator interface (e.g. `np.random.default_rng(seed)`).
+            If given, all sampling draws are taken from it, making placements
+            reproducible independent of the global numpy random state.
+            Defaults to `np.random`, preserving the previous behavior.
+        """
+        self._rng = rng if rng is not None else np.random
         self._boundaries = []
         areas = []
         for b in boundaries:
-            bo = BoundaryObject(b)
+            bo = BoundaryObject(b, rng=self._rng)
             areas.append(bo.get_area())
             self._boundaries.append(bo)
         self._probabilities = np.array(areas) / np.sum(areas)
@@ -156,7 +169,7 @@ def sample(self, obj: Object, ignore_collisions=False,
         """
         collision_fails = boundary_fails = self.MAX_SAMPLES
         while collision_fails > 0 and boundary_fails > 0:
-            sampled_boundary = np.random.choice(self._boundaries,
+            sampled_boundary = self._rng.choice(self._boundaries,
                                                 p=self._probabilities)
             result = sampled_boundary.add(
                 obj, ignore_collisions, min_rotation, max_rotation, min_distance)

rlbench/environment.py

@@ -42,6 +42,7 @@ def __init__(self,
                  shaped_rewards: bool = False,
                  arm_max_velocity: float = 1.0,
                  arm_max_acceleration: float = 4.0,
+                 rng=None,
                  ):
 
         self._dataset_root = dataset_root
@@ -50,6 +51,10 @@ def __init__(self,
         self._headless = headless
         self._static_positions = static_positions
         self._robot_setup = robot_setup.lower()
+        # Optional source of randomness (e.g. np.random.default_rng(seed))
+        # used by the scene when sampling object placements. If None, the
+        # global numpy random state is used, as before.
+        self._rng = rng
 
         self._randomize_every = randomize_every
         self._frequency = frequency
@@ -123,13 +128,15 @@ def launch(self):
         self._robot = Robot(arm, gripper)
         if self._randomize_every is None:
             self._scene = Scene(
-                self._pyrep, self._robot, self._obs_config, self._robot_setup)
+                self._pyrep, self._robot, self._obs_config, self._robot_setup,
+                rng=self._rng)
         else:
             self._scene = DomainRandomizationScene(
                 self._pyrep, self._robot, self._obs_config, self._robot_setup,
                 self._randomize_every, self._frequency,
                 self._visual_randomization_config,
-                self._dynamics_randomization_config)
+                self._dynamics_randomization_config,
+                rng=self._rng)
 
         self._action_mode.arm_action_mode.set_control_mode(self._robot)
 

rlbench/sim2real/domain_randomization_scene.py

@@ -29,8 +29,9 @@ def __init__(self,
                  randomize_every: RandomizeEvery=RandomizeEvery.EPISODE,
                  frequency: int=1,
                  visual_randomization_config=None,
-                 dynamics_randomization_config=None):
-        super().__init__(pyrep, robot, obs_config, robot_setup)
+                 dynamics_randomization_config=None,
+                 rng=None):
+        super().__init__(pyrep, robot, obs_config, robot_setup, rng=rng)
         self._randomize_every = randomize_every
         self._frequency = frequency
         self._visual_rand_config = visual_randomization_config

tests/unit/test_spawn_boundary.py

@@ -0,0 +1,207 @@
+"""Unit tests for rlbench.backend.spawn_boundary.
+
+These tests do not require CoppeliaSim/PyRep: the module under test only
+uses pyrep for type annotations, so the `pyrep` package is stubbed when it
+is not importable, and `spawn_boundary.py` / `exceptions.py` are loaded
+directly from the repository. This keeps the tests runnable in any
+environment with numpy installed.
+"""
+import importlib.util
+import os
+import sys
+import types
+import unittest
+
+import numpy as np
+
+REPO_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(
+    os.path.abspath(__file__))))
+BOUNDARY_PATH = os.path.join(
+    REPO_ROOT, 'rlbench', 'backend', 'spawn_boundary.py')
+EXCEPTIONS_PATH = os.path.join(
+    REPO_ROOT, 'rlbench', 'backend', 'exceptions.py')
+
+# Boundary with bbox (min_x, max_x, min_y, max_y, min_z, max_z) => not a plane.
+BOUNDARY_BBOX = (-1.0, 1.0, -1.0, 1.0, 0.0, 1.0)
+# A small object that always fits within the boundary, including the worst
+# case where any rotation about the origin is applied to it.
+OBJ_BBOX = (-0.05, 0.05, -0.05, 0.05, 0.48, 0.52)
+# An object that can never fit within the boundary.
+HUGE_OBJ_BBOX = (-10.0, 10.0, -10.0, 10.0, -10.0, 10.0)
+
+
+def _load_from_file(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _stub_missing_dependencies():
+    """Make spawn_boundary.py importable without pyrep/gymnasium.
+
+    Real modules are always preferred: we only register lightweight stubs
+    when the real imports fail (i.e. PyRep/CoppeliaSim is not available),
+    so CI machines with PyRep installed are unaffected.
+    """
+    try:
+        import pyrep  # noqa: F401
+    except Exception:
+        pyrep_mod = types.ModuleType('pyrep')
+        objects_mod = types.ModuleType('pyrep.objects')
+        object_mod = types.ModuleType('pyrep.objects.object')
+
+        class Object(object):
+            pass
+
+        object_mod.Object = Object
+        pyrep_mod.objects = objects_mod
+        objects_mod.object = object_mod
+        sys.modules.setdefault('pyrep', pyrep_mod)
+        sys.modules.setdefault('pyrep.objects', objects_mod)
+        sys.modules.setdefault('pyrep.objects.object', object_mod)
+
+    try:
+        import rlbench.backend.exceptions  # noqa: F401
+    except Exception:
+        # Loads the real exceptions.py standalone and registers it so that
+        # `from rlbench.backend.exceptions import BoundaryError` resolves
+        # without importing the rlbench package (which needs gymnasium).
+        _load_from_file('rlbench.backend.exceptions', EXCEPTIONS_PATH)
+
+
+_stub_missing_dependencies()
+_spawn_boundary = _load_from_file(
+    'rlbench_backend_spawn_boundary_under_test', BOUNDARY_PATH)
+SpawnBoundary = _spawn_boundary.SpawnBoundary
+BoundaryError = _spawn_boundary.BoundaryError
+
+
+class FakeObject(object):
+    """Duck-typed stand-in for a pyrep Object (boundary or placed object)."""
+
+    def __init__(self, bbox, is_model=False, collides=False):
+        self._bbox = bbox
+        self._is_model = is_model
+        self._collides = collides
+        self.position = None
+        self.rotation = None
+
+    def is_model(self):
+        return self._is_model
+
+    def get_model_bounding_box(self):
+        return self._bbox
+
+    def get_bounding_box(self):
+        return self._bbox
+
+    def get_position(self, relative_to=None):
+        return np.array([0.0, 0.0, 0.0])
+
+    def set_position(self, position, relative_to=None):
+        self.position = np.array(position, dtype=float)
+
+    def rotate(self, rotation):
+        self.rotation = list(rotation)
+
+    def get_objects_in_tree(self, exclude_base=True,
+                            first_generation_only=False):
+        return [self]
+
+    def check_collision(self, other):
+        return self._collides
+
+
+class TestSpawnBoundaryDeterminism(unittest.TestCase):
+
+    def test_same_seed_produces_identical_placements(self):
+        boundary_a, boundary_b = FakeObject(BOUNDARY_BBOX), FakeObject(
+            BOUNDARY_BBOX)
+        sb_a = SpawnBoundary([boundary_a], rng=np.random.default_rng(7))
+        sb_b = SpawnBoundary([boundary_b], rng=np.random.default_rng(7))
+        obj_a, obj_b = FakeObject(OBJ_BBOX), FakeObject(OBJ_BBOX)
+
+        sb_a.sample(obj_a)
+        sb_b.sample(obj_b)
+
+        self.assertIsNotNone(obj_a.position)
+        np.testing.assert_allclose(obj_a.position, obj_b.position)
+        np.testing.assert_allclose(obj_a.rotation, obj_b.rotation)
+
+    def test_same_seed_identical_across_multiple_boundaries(self):
+        boundaries = [FakeObject(BOUNDARY_BBOX),
+                      FakeObject((0.0, 1.0, 0.0, 1.0, 0.0, 1.0))]
+        sb_a = SpawnBoundary(
+            [FakeObject(b._bbox) for b in boundaries],
+            rng=np.random.default_rng(3))
+        sb_b = SpawnBoundary(
+            [FakeObject(b._bbox) for b in boundaries],
+            rng=np.random.default_rng(3))
+        obj_a, obj_b = FakeObject(OBJ_BBOX), FakeObject(OBJ_BBOX)
+
+        sb_a.sample(obj_a)
+        sb_b.sample(obj_b)
+
+        np.testing.assert_allclose(obj_a.position, obj_b.position)
+        np.testing.assert_allclose(obj_a.rotation, obj_b.rotation)
+
+    def test_different_seed_produces_different_placements(self):
+        sb_a = SpawnBoundary(
+            [FakeObject(BOUNDARY_BBOX)], rng=np.random.default_rng(7))
+        sb_b = SpawnBoundary(
+            [FakeObject(BOUNDARY_BBOX)], rng=np.random.default_rng(8))
+        obj_a, obj_b = FakeObject(OBJ_BBOX), FakeObject(OBJ_BBOX)
+
+        sb_a.sample(obj_a)
+        sb_b.sample(obj_b)
+
+        self.assertFalse(np.allclose(obj_a.position, obj_b.position))
+
+    def test_random_state_is_supported(self):
+        sb_a = SpawnBoundary(
+            [FakeObject(BOUNDARY_BBOX)], rng=np.random.RandomState(5))
+        sb_b = SpawnBoundary(
+            [FakeObject(BOUNDARY_BBOX)], rng=np.random.RandomState(5))
+        obj_a, obj_b = FakeObject(OBJ_BBOX), FakeObject(OBJ_BBOX)
+
+        sb_a.sample(obj_a)
+        sb_b.sample(obj_b)
+
+        np.testing.assert_allclose(obj_a.position, obj_b.position)
+        np.testing.assert_allclose(obj_a.rotation, obj_b.rotation)
+
+    def test_no_rng_keeps_global_numpy_behavior(self):
+        np.random.seed(11)
+        sb_a = SpawnBoundary([FakeObject(BOUNDARY_BBOX)])
+        obj_a = FakeObject(OBJ_BBOX)
+        sb_a.sample(obj_a)
+
+        np.random.seed(11)
+        sb_b = SpawnBoundary([FakeObject(BOUNDARY_BBOX)])
+        obj_b = FakeObject(OBJ_BBOX)
+        sb_b.sample(obj_b)
+
+        self.assertIsNotNone(obj_a.position)
+        np.testing.assert_allclose(obj_a.position, obj_b.position)
+        np.testing.assert_allclose(obj_a.rotation, obj_b.rotation)
+
+    def test_positions_stay_within_boundary(self):
+        sb = SpawnBoundary(
+            [FakeObject(BOUNDARY_BBOX)], rng=np.random.default_rng(0))
+        for _ in range(20):
+            obj = FakeObject(OBJ_BBOX)
+            sb.sample(obj)
+            self.assertTrue(np.all(obj.position >= np.array([-1., -1., 0.])))
+            self.assertTrue(np.all(obj.position <= np.array([1., 1., 1.])))
+
+    def test_boundary_error_still_raised_for_oversized_object(self):
+        sb = SpawnBoundary(
+            [FakeObject(BOUNDARY_BBOX)], rng=np.random.default_rng(0))
+        with self.assertRaises(BoundaryError):
+            sb.sample(FakeObject(HUGE_OBJ_BBOX))
+
+
+if __name__ == '__main__':
+    unittest.main()