diff --git a/rlbench/backend/scene.py b/rlbench/backend/scene.py index 96fc1ddd1..564283814 100644 --- a/rlbench/backend/scene.py +++ b/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') diff --git a/rlbench/backend/spawn_boundary.py b/rlbench/backend/spawn_boundary.py index 846981b6e..7b0ef229b 100644 --- a/rlbench/backend/spawn_boundary.py +++ b/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) diff --git a/rlbench/environment.py b/rlbench/environment.py index 4e9cb1a23..2dc630e55 100644 --- a/rlbench/environment.py +++ b/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) diff --git a/rlbench/sim2real/domain_randomization_scene.py b/rlbench/sim2real/domain_randomization_scene.py index 9ea7cee39..904f5c6f6 100644 --- a/rlbench/sim2real/domain_randomization_scene.py +++ b/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 diff --git a/tests/unit/test_spawn_boundary.py b/tests/unit/test_spawn_boundary.py new file mode 100644 index 000000000..822186515 --- /dev/null +++ b/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()