diff --git a/isaaclab_arena/assets/background_library.py b/isaaclab_arena/assets/background_library.py index 4e5639bcfb..7fcaf01c9a 100644 --- a/isaaclab_arena/assets/background_library.py +++ b/isaaclab_arena/assets/background_library.py @@ -173,7 +173,12 @@ class LightwheelKitchenBackground(LibraryBackground): initial_pose = Pose.identity() object_min_z = -0.2 - def __init__(self, layout_id: int = 1, style_id: int = 1): + def __init__( + self, + layout_id: int = 1, + style_id: int = 1, + **kwargs, + ): from lightwheel_sdk.loader import floorplan_loader # Lazily download the USD @@ -188,7 +193,7 @@ def __init__(self, layout_id: int = 1, style_id: int = 1): backend="robocasa", )[0] ) - super().__init__() + super().__init__(**kwargs) def get_viewer_cfg(self) -> ViewerCfg: # Looking in through the open front. diff --git a/isaaclab_arena/environments/relation_solver_interface.py b/isaaclab_arena/environments/relation_solver_interface.py index c8730b60ce..a0dc93974d 100644 --- a/isaaclab_arena/environments/relation_solver_interface.py +++ b/isaaclab_arena/environments/relation_solver_interface.py @@ -11,7 +11,7 @@ from isaaclab_arena.relations.collision_mode import CollisionMode, get_object_collision_mode from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams -from isaaclab_arena.relations.placement_events import get_pose_from_layout, solve_and_place_objects +from isaaclab_arena.relations.placement_events import PlacementPoolHandle, get_pose_from_layout, solve_and_place_objects from isaaclab_arena.relations.pooled_object_placer import PooledObjectPlacer from isaaclab_arena.relations.relations import get_anchor_objects from isaaclab_arena.utils.pose import PosePerEnv @@ -184,8 +184,7 @@ def _apply_dynamic_spawn_pose( func=solve_and_place_objects, mode="reset", params={ - "assets": assets, - "placement_pool": placement_pool, + "placement_pool": PlacementPoolHandle(placement_pool), }, ) diff --git a/isaaclab_arena/relations/placement_asset.py b/isaaclab_arena/relations/placement_asset.py index c659f54979..02a191ab88 100644 --- a/isaaclab_arena/relations/placement_asset.py +++ b/isaaclab_arena/relations/placement_asset.py @@ -27,16 +27,25 @@ class PlaceableAsset(Asset, ABC): """Asset whose root pose can be constrained by spatial relations.""" - def __init__(self, name: str, tags: list[str] | None = None, **kwargs) -> None: + def __init__( + self, + name: str, + tags: list[str] | None = None, + collision_mode: CollisionMode | str | None = None, + repair_collision_mesh_non_watertight: bool = True, + **kwargs, + ) -> None: super().__init__(name=name, tags=tags, **kwargs) self.initial_pose: Pose | PoseRange | PosePerEnv | None = None self._pose_event_cfg: EventTermCfg | None = None """Reset event restoring this asset's root pose; ``None`` until a pose with a reset event is set.""" self.relations: list[RelationBase] = [] # None delegates collision-mode selection to the solver. - self.collision_mode: CollisionMode | None = None + if collision_mode is not None: + collision_mode = CollisionMode(collision_mode) + self.collision_mode = collision_mode # Whether to replace a non-watertight collision mesh with its convex hull. - self.repair_collision_mesh_non_watertight = True + self.repair_collision_mesh_non_watertight = repair_collision_mesh_non_watertight def add_relation(self, relation: RelationBase) -> None: """Attach a relation to the asset.""" diff --git a/isaaclab_arena/relations/placement_events.py b/isaaclab_arena/relations/placement_events.py index bed1ca3d7e..49cf6315e4 100644 --- a/isaaclab_arena/relations/placement_events.py +++ b/isaaclab_arena/relations/placement_events.py @@ -26,6 +26,32 @@ PLACEMENT_RESET_EVENT_NAME = "placement_reset" +class PlacementPoolHandle: + """Opaque holder for a runtime placement pool to bypass EventTermCfg param deepcopy/validation errors. + + PooledObjectPlacer is used as an EventTermCfg param to set the initial spawn pose. Isaac Lab deep-copies + and validates the configclass param, leading to two crashes: deepcopy fails for the Warp GPU cache + wp.Mesh BVHs ("ctypes objects containing pointers cannot be pickled"); validation hits RecursionError + when recursively walking all dicts and reaches placement assets (including embodiments with cyclic + scene configs). + + This handle wraps PooledObjectPlacer with overrides for deepcopy and validation, while + PooledObjectPlacer itself stays a normal class. EventTermCfg params use this handle. + """ + + __slots__ = ("pool",) + """Store pool in a slot instead of an instance dictionary; ``hasattr(handle, "__dict__")`` is false, + so ``_validate(handle)`` stops traversing into PooledObjectPlacer .""" + + def __init__(self, pool: PooledObjectPlacer) -> None: + self.pool = pool + + def __deepcopy__(self, memo: dict[int, object]) -> PlacementPoolHandle: + """Share the live pool across ``copy.deepcopy`` to avoid deep-copying the Warp cache BVHs.""" + memo[id(self)] = self + return self + + def get_placement_pool(env) -> PooledObjectPlacer | None: """Return the pooled placer stored on the env reset event, or ``None`` when absent. @@ -39,7 +65,9 @@ def get_placement_pool(env) -> PooledObjectPlacer | None: term_cfg = env.unwrapped.event_manager.get_term_cfg(PLACEMENT_RESET_EVENT_NAME) except ValueError: return None - return term_cfg.params.get("placement_pool") + handle = term_cfg.params.get("placement_pool") + assert handle is not None, f"'{PLACEMENT_RESET_EVENT_NAME}' event is missing its placement_pool parameter." + return handle.pool def get_rotation_xyzw(asset: PlaceableAsset) -> tuple[float, float, float, float]: @@ -113,8 +141,7 @@ def write_layout_to_sim( def solve_and_place_objects( env: ManagerBasedEnv, env_ids: torch.Tensor | None, - assets: list[PlaceableAsset], - placement_pool: PooledObjectPlacer, + placement_pool: PlacementPoolHandle, ) -> None: """Coordinated reset event that draws layouts from the pool and writes poses. @@ -125,17 +152,19 @@ def solve_and_place_objects( Args: env: The Isaac Lab environment. env_ids: 1-D tensor of environment indices being reset. - assets: Assets participating in relation solving. - placement_pool: Runtime pool of solved placement layouts. + placement_pool: Opaque handle to the runtime pool of solved placement layouts. + Layout assets come from ``placement_pool.pool.objects``. """ + pool = placement_pool.pool if env_ids is None or len(env_ids) == 0: return + assets = pool.objects reset_env_ids = env_ids.tolist() num_scene_envs = env.scene.env_origins.shape[0] assert ( - placement_pool.num_envs == num_scene_envs - ), f"Placement pool has {placement_pool.num_envs} envs, but scene has {num_scene_envs} env origins." - results_by_env = placement_pool.sample_for_envs(reset_env_ids) + pool.num_envs == num_scene_envs + ), f"Placement pool has {pool.num_envs} envs, but scene has {num_scene_envs} env origins." + results_by_env = pool.sample_for_envs(reset_env_ids) anchor_assets = set(get_anchor_objects(assets)) base_rotations = get_base_rotation_per_asset(assets) diff --git a/isaaclab_arena/tests/test_placement_events.py b/isaaclab_arena/tests/test_placement_events.py index 29b7dfbeda..8ccd23051d 100644 --- a/isaaclab_arena/tests/test_placement_events.py +++ b/isaaclab_arena/tests/test_placement_events.py @@ -149,16 +149,11 @@ def scene_getitem(self, name: str) -> MagicMock: return env -def _solve_and_place_with_pool(env, env_ids, objects, pool): +def _solve_and_place_with_pool(env, env_ids, pool): """Call the reset event with the same runtime params EventTermCfg stores.""" - from isaaclab_arena.relations.placement_events import solve_and_place_objects + from isaaclab_arena.relations.placement_events import PlacementPoolHandle, solve_and_place_objects - return solve_and_place_objects( - env, - env_ids, - assets=objects, - placement_pool=pool, - ) + return solve_and_place_objects(env, env_ids, placement_pool=PlacementPoolHandle(pool)) def test_solve_and_place_objects_writes_poses_to_sim(): @@ -176,7 +171,7 @@ def test_solve_and_place_objects_writes_poses_to_sim(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) # Anchor (desk) should NOT have been written. assert "desk" not in env._assets, "Anchor pose should not be written to sim" @@ -191,7 +186,7 @@ def test_solve_and_place_objects_writes_poses_to_sim(): def test_solve_and_place_objects_uses_runtime_pool(): - from isaaclab_arena.relations.placement_events import solve_and_place_objects + from isaaclab_arena.relations.placement_events import PlacementPoolHandle, solve_and_place_objects from isaaclab_arena.relations.placement_result import PlacementResult from isaaclab_arena.tests.dummy_embodiment import DummyEmbodiment from isaaclab_arena.utils.bounding_box import AxisAlignedBoundingBox @@ -206,6 +201,7 @@ def test_solve_and_place_objects_uses_runtime_pool(): class Pool: num_envs = 1 + objects = [desk, robot] def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: assert env_ids == [0] @@ -221,8 +217,7 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: solve_and_place_objects( env, torch.tensor([0]), - assets=[desk, robot], - placement_pool=Pool(), + placement_pool=PlacementPoolHandle(Pool()), ) assert "desk" not in env._assets @@ -296,14 +291,14 @@ def test_reset_placement_asset_pose_per_env_requires_full_env_coverage(): def test_get_placement_pool_returns_runtime_pool(): - from isaaclab_arena.relations.placement_events import get_placement_pool + from isaaclab_arena.relations.placement_events import PlacementPoolHandle, get_placement_pool class Pool: pass pool = Pool() env = MagicMock() - env.unwrapped.event_manager.get_term_cfg.return_value.params = {"placement_pool": pool} + env.unwrapped.event_manager.get_term_cfg.return_value.params = {"placement_pool": PlacementPoolHandle(pool)} assert get_placement_pool(env) is pool @@ -328,7 +323,7 @@ def test_solve_and_place_objects_applies_random_yaw(): ) pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) # Anchor (desk) is never rotated or written, even with random yaw enabled. assert "desk" not in env._assets, "Anchor pose should not be written to sim" @@ -356,7 +351,7 @@ def test_solve_and_place_objects_skips_empty_env_ids(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=[desk, box1, box2], placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, torch.tensor([], dtype=torch.int64), [desk, box1, box2], pool) + _solve_and_place_with_pool(env, torch.tensor([], dtype=torch.int64), pool) assert len(env._assets) == 0, "No writes should occur for empty env_ids" @@ -373,7 +368,7 @@ def test_solve_and_place_objects_skips_none_env_ids(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=[desk, box1, box2], placer_params=placer_params, pool_size=10) - _solve_and_place_with_pool(env, None, [desk, box1, box2], pool) + _solve_and_place_with_pool(env, None, pool) assert len(env._assets) == 0, "No writes should occur for None env_ids" @@ -394,7 +389,7 @@ def test_solve_and_place_objects_handles_multiple_env_ids(): placer_params = ObjectPlacerParams(solver_params=solver_params) pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=12, num_envs=num_envs) - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) assert "desk" not in env._assets, "Anchor pose should not be written to sim" @@ -424,7 +419,7 @@ def test_solve_and_place_objects_partial_reset_homogeneous_pool_consumes_only_re pool = PooledObjectPlacer(objects=objects, placer_params=placer_params, pool_size=12, num_envs=num_envs) available_before = pool.total_remaining - _solve_and_place_with_pool(env, env_ids, objects, pool) + _solve_and_place_with_pool(env, env_ids, pool) available_after = pool.total_remaining assert available_before - available_after == len(env_ids) @@ -436,11 +431,11 @@ def test_solve_and_place_objects_writes_invalid_fallback_layout(capsys): from isaaclab_arena.relations.placement_result import PlacementResult desk, box1, box2 = _create_test_objects() - objects = [desk, box1, box2] env = _make_mock_env(num_envs=1) class InvalidPool: num_envs = 1 + objects = [desk, box1, box2] def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: assert env_ids == [0] @@ -453,7 +448,7 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: ) } - _solve_and_place_with_pool(env, torch.tensor([0]), objects, InvalidPool()) + _solve_and_place_with_pool(env, torch.tensor([0]), InvalidPool()) captured = capsys.readouterr() assert set(env._assets) == {box1.name, box2.name} @@ -466,12 +461,12 @@ def test_solve_and_place_objects_partial_reset_applies_absolute_env_origin(): from isaaclab_arena.relations.placement_result import PlacementResult desk, box1, box2 = _create_test_objects() - objects = [desk, box1, box2] env = _make_mock_env(num_envs=4) env.scene.env_origins[2] = torch.tensor([10.0, 0.0, 0.0]) class EnvIndexedPool: num_envs = 4 + objects = [desk, box1, box2] requested_env_ids = None def sample_without_replacement(self, count: int) -> list[PlacementResult]: @@ -493,7 +488,7 @@ def sample_for_envs(self, env_ids: list[int]) -> dict[int, PlacementResult]: } pool = EnvIndexedPool() - _solve_and_place_with_pool(env, torch.tensor([2]), objects, pool) + _solve_and_place_with_pool(env, torch.tensor([2]), pool) box1_pose = env._assets[box1.name].write_root_pose_to_sim.call_args[0][0] box2_pose = env._assets[box2.name].write_root_pose_to_sim.call_args[0][0] @@ -510,14 +505,14 @@ def test_solve_and_place_objects_asserts_env_indexed_pool_size_matches_scene(): """Env-indexed pool slots must line up with absolute Isaac Lab env ids.""" desk, box1, box2 = _create_test_objects() - objects = [desk, box1, box2] env = _make_mock_env(num_envs=2) class MismatchedEnvIndexedPool: num_envs = 1 + objects = [desk, box1, box2] with pytest.raises(AssertionError, match="scene has 2 env origins"): - _solve_and_place_with_pool(env, torch.tensor([0]), objects, MismatchedEnvIndexedPool()) + _solve_and_place_with_pool(env, torch.tensor([0]), MismatchedEnvIndexedPool()) def test_pooled_placer_sample_without_replacement_returns_different_layouts(): @@ -982,6 +977,13 @@ def test_solve_and_apply_relation_placement_drops_embodiment_from_event_params() assert params.reachability_config.embodiment is embodiment # ...while the pool the reset event captured no longer references the embodiment -- on the placer params # and on every built validator alike -- so configclass never deep-copies or recurses into it. - pool = event.params["placement_pool"] + from isaaclab.utils.configclass import _validate + + from isaaclab_arena.relations.placement_events import PlacementPoolHandle + + pool_handle = event.params["placement_pool"] + assert isinstance(pool_handle, PlacementPoolHandle) + _validate(event, prefix="") + pool = pool_handle.pool assert pool._placer.params.reachability_config.embodiment is None assert all(v._params.reachability_config.embodiment is None for v in pool._placer._validators) diff --git a/isaaclab_arena/tests/test_relation_solver_interface.py b/isaaclab_arena/tests/test_relation_solver_interface.py index 561a446770..a0af67fc02 100644 --- a/isaaclab_arena/tests/test_relation_solver_interface.py +++ b/isaaclab_arena/tests/test_relation_solver_interface.py @@ -34,8 +34,13 @@ def _make_box(name: str = "box"): class _FakePlacementPool: - def __init__(self, layouts) -> None: + def __init__(self, layouts, objects=None) -> None: self._layouts = layouts + self._objects = objects or [] + + @property + def objects(self): + return self._objects def sample_with_replacement(self, count: int): return self._layouts[:count] @@ -142,7 +147,10 @@ def test_dynamic_spawn_pose_event_params_use_runtime_assets(): desk = _make_desk() box = _make_box() - placement_pool = _FakePlacementPool([_fallback_layout(positions={box: (0.1, 0.2, 0.3)})]) + placement_pool = _FakePlacementPool( + [_fallback_layout(positions={box: (0.1, 0.2, 0.3)})], + objects=[desk, box], + ) event_cfg = _apply_dynamic_spawn_pose( assets=[desk, box], @@ -150,8 +158,48 @@ def test_dynamic_spawn_pose_event_params_use_runtime_assets(): anchor_assets={desk}, ) - assert [asset.name for asset in event_cfg.params["assets"]] == ["desk", "box"] assert "placement_pool" in event_cfg.params + from isaaclab_arena.relations.placement_events import PlacementPoolHandle + + assert isinstance(event_cfg.params["placement_pool"], PlacementPoolHandle) + assert [asset.name for asset in event_cfg.params["placement_pool"].pool.objects] == [ + "desk", + "box", + ] + + +def test_dynamic_spawn_pose_event_cfg_deepcopy_after_mesh_solve(): + """EventTermCfg deep-copies params; handle shares the pool across config copies.""" + import copy + import trimesh + + from isaaclab.managers import EventTermCfg + + from isaaclab_arena.environments.relation_solver_interface import solve_and_apply_relation_placement + from isaaclab_arena.relations.object_placer_params import ObjectPlacerParams + from isaaclab_arena.relations.relation_solver_params import CollisionMode, RelationSolverParams + from isaaclab_arena.relations.relations import On + + desk = _make_desk() + box = _make_box() + box.add_relation(On(desk, clearance_m=0.01)) + box.collision_mode = CollisionMode.MESH + box._collision_mesh = trimesh.creation.box(extents=(0.2, 0.2, 0.2)) + + params = ObjectPlacerParams( + placement_seed=17, + resolve_on_reset=True, + min_unique_layouts_per_env=1, + solver_params=RelationSolverParams(collision_mode=CollisionMode.MESH, max_iters=50), + ) + event_cfg = solve_and_apply_relation_placement([desk, box], num_envs=1, placer_params=params) + + assert event_cfg is not None + assert isinstance(event_cfg, EventTermCfg) + copy.deepcopy(event_cfg) + from isaaclab.utils.configclass import _validate + + _validate(event_cfg, prefix="") def test_static_embodiment_placement_stores_per_env_poses(): diff --git a/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml b/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml index 81827a2156..15d0836214 100644 --- a/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml +++ b/isaaclab_arena_environments/kitchen_bench/droid_pick_and_place_lightwheel_kitchen.yaml @@ -13,7 +13,9 @@ embodiment: background: id: kitchen registry_name: lightwheel_robocasa_kitchen - params: {} + params: + collision_mode: mesh + repair_collision_mesh_non_watertight: false objects: - id: mustard_bottle registry_name: mustard_bottle_hope_robolab