diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 91a55ad81..bdd9ac55b 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -160,16 +160,42 @@ def __init__( self._configure_timing() + # Phase 1 only declares scene topology. Spawn-backed assets intentionally + # remain metadata-light until the single prepare boundary below. self._setup_scene(**kwargs) # Keep the established env._profiler API while sharing the single # profiler instance owned by SimulationManager. self._profiler = self.sim.profiler - if self.sim.is_default_backend and self.sim.is_use_gpu_physics: - self.sim.init_gpu_physics() - elif self.sim.is_newton_backend: - self.sim.finalize_newton_physics() + # Materialize every physical declaration in one transaction. DexSim's + # articulation adapter parses each source while finalizing, then the + # resulting handles bind the existing EmbodiChain facades in place. + self.sim.prepare() + + # Phase 2 may now consume link/joint metadata, construct action spaces, + # and create render-only resources such as CameraGroup instances. + configured_robot = self._setup_robot(**kwargs) + if configured_robot is not None: + self.robot = configured_robot + + if self.robot is None: + logger.log_error( + f"The robot instance must be initialized in :meth:`_setup_robot` function." + ) + if len(self.active_joint_ids) == 0: + self.active_joint_ids = self.robot.active_joint_ids + if self.single_action_space is None: + logger.log_error( + f":attr:`single_action_space` must be defined in the :meth:`_setup_robot` function." + ) + + self.sensors = self._setup_sensors(**kwargs) + self._camera_group_ids = [ + sensor.group_id + for sensor in self.sensors.values() + if isinstance(sensor, Camera) + ] if not self.sim_cfg.headless: self.sim.open_window() @@ -483,8 +509,9 @@ def add_camera_group_id(self, group_id: int) -> None: self._camera_group_ids.append(group_id) def _setup_scene(self, **kwargs): - # Init sim manager. - # we want to open gui window when the scene is setup, so init sim manager in headless mode first. + """Declare physical scene topology without consuming runtime metadata.""" + # Init sim manager. We want to open the GUI window after the scene is + # materialized, so construct the manager in headless mode first. headless = self.sim_cfg.headless self.sim_cfg.headless = True self.sim = SimulationManager(self.sim_cfg) @@ -494,35 +521,35 @@ def _setup_scene(self, **kwargs): f"Initializing {self.num_envs} environments on {self.sim_cfg.device}." ) - self.robot = self._setup_robot(**kwargs) - if len(self.active_joint_ids) == 0: - self.active_joint_ids = self.robot.active_joint_ids - - if self.robot is None: - logger.log_error( - f"The robot instance must be initialized in :meth:`_setup_robot` function." - ) - if self.single_action_space is None: - logger.log_error( - f":attr:`single_action_space` must be defined in the :meth:`_setup_robot` function." - ) + # Config-driven environments can declare their robot here while + # deferring all link/joint queries until the post-prepare phase. Generic + # BaseEnv subclasses may keep returning None and add a runtime robot in + # _setup_robot() for backwards compatibility. + self.robot = self._declare_robot(**kwargs) self._prepare_scene(**kwargs) - self.sensors = self._setup_sensors(**kwargs) + def _declare_robot(self, **kwargs) -> Robot | None: + """Optionally declare a robot before the scene prepare boundary. + + Config-driven environments should override this hook and call + :meth:`SimulationManager.add_robot` without querying link/joint data. + The returned facade is bound in place by :meth:`SimulationManager.prepare`. - # Setup camera groups for rendering. - self._camera_group_ids: List[int] = [] - for sensor in self.sensors.values(): - if isinstance(sensor, Camera): - self._camera_group_ids.append(sensor.group_id) + Generic subclasses that only implement the historical + :meth:`_setup_robot` hook remain supported: their robot is added after + the initial prepare boundary and is prepared immediately by the manager. + """ + del kwargs + return None def _setup_robot(self, **kwargs) -> Robot: - """Load the robot agent, setup the controller and action space. + """Configure the bound robot, controller, and action space. Note: - 1. The fuction must return the robot instance. - 2. The self.single_action_space should be defined. + This hook runs after :meth:`SimulationManager.prepare`, so link, + joint, and limit metadata are available. It must return the robot + instance and define ``self.single_action_space``. """ # TODO: single_action_space may be configured in config? diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index 1f42df648..4878acee7 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -789,9 +789,9 @@ def _extend_reward( return rewards def _prepare_scene(self, **kwargs) -> None: - self._setup_lights() self._setup_background() self._setup_interactive_objects() + self._setup_lights() def _update_sim_state(self, **kwargs) -> None: """Perform the simulation step and apply events if configured. @@ -1657,8 +1657,15 @@ def _postprocess_action(self, action): return self.action_manager.process_action(action, mode="post") return super()._postprocess_action(action) + def _declare_robot(self, **kwargs) -> Robot: + """Declare the configured robot without reading articulation metadata.""" + del kwargs + if self.cfg.robot is None: + logger.log_error("Robot configuration is not provided.") + return self.sim.add_robot(self.cfg.robot) + def _setup_robot(self, **kwargs) -> Robot: - """Setup the robot in the environment. + """Configure the finalized robot interface for the environment. Currently, only joint position control is supported. Would be extended to support joint velocity and torque control in the future. @@ -1666,11 +1673,10 @@ def _setup_robot(self, **kwargs) -> Robot: Returns: Robot: The robot instance added to the scene. """ - if self.cfg.robot is None: - logger.log_error("Robot configuration is not provided.") - - # Initialize the robot based on the configuration. - robot: Robot = self.sim.add_robot(self.cfg.robot) + del kwargs + robot = self.robot + if robot is None: + logger.log_error("Robot was not declared before simulation prepare.") # Setup active joints for robot to control. if self.cfg.control_parts: diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index a55d2ee34..f30e65a6b 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -20,9 +20,10 @@ import dexsim import numpy as np +from copy import deepcopy from dataclasses import dataclass from functools import cached_property -from typing import List, Sequence, Dict, Union, Tuple, Optional +from typing import TYPE_CHECKING, List, Sequence, Dict, Union, Tuple, Optional from dexsim.engine import Articulation as _Articulation from dexsim.types import ( @@ -55,8 +56,10 @@ from embodichain.lab.sim.objects.backends import ( DefaultArticulationView, NewtonArticulationView, + SpawnArticulationView, is_newton_scene, ) +from embodichain.lab.sim.objects.backends.base import ArticulationViewBase from embodichain.utils.math import ( matrix_from_quat, quat_from_matrix, @@ -69,13 +72,20 @@ ) from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedArticulation + @dataclass class ArticulationData: """GPU data manager for articulation.""" def __init__( - self, entities: List[_Articulation], ps: PhysicsScene, device: torch.device + self, + entities: Sequence[_Articulation | SpawnedArticulation], + ps: PhysicsScene | None, + device: torch.device, + articulation_view: ArticulationViewBase | None = None, ) -> None: """Initialize the ArticulationData. @@ -88,7 +98,9 @@ def __init__( self.ps = ps self.num_instances = len(entities) self.device = device - if is_newton_scene(ps): + if articulation_view is not None: + self.articulation_view = articulation_view + elif is_newton_scene(ps): self.articulation_view = NewtonArticulationView( entities=entities, scene=ps, device=device ) @@ -100,9 +112,14 @@ def __init__( # Backward-compatible alias for callers that use GPU/articulation ids. self.gpu_indices = self.articulation_view.articulation_ids_tensor - self.dof = self.entities[0].get_dof() - self.num_links = self.entities[0].get_links_num() - self.link_names = self.entities[0].get_link_names() + if isinstance(self.articulation_view, SpawnArticulationView): + self.dof = self.articulation_view.dof + self.num_links = self.articulation_view.num_links + self.link_names = self.articulation_view.link_names + else: + self.dof = self.entities[0].get_dof() + self.num_links = self.entities[0].get_links_num() + self.link_names = self.entities[0].get_link_names() self._root_pose = torch.zeros( (self.num_instances, 7), dtype=torch.float32, device=self.device @@ -114,11 +131,13 @@ def __init__( (self.num_instances, 3), dtype=torch.float32, device=self.device ) - max_num_links = ( - self.ps.gpu_get_articulation_max_link_count() - if self.device.type == "cuda" and not self.is_newton_backend - else self.num_links - ) + max_num_links = self.num_links + if ( + articulation_view is None + and self.device.type == "cuda" + and not self.is_newton_backend + ): + max_num_links = self.ps.gpu_get_articulation_max_link_count() self._body_link_pose = torch.zeros( (self.num_instances, max_num_links, 7), dtype=torch.float32, @@ -141,11 +160,13 @@ def __init__( device=self.device, ) - max_dof = ( - self.ps.gpu_get_articulation_max_dof() - if self.device.type == "cuda" and not self.is_newton_backend - else self.dof - ) + max_dof = self.dof + if ( + articulation_view is None + and self.device.type == "cuda" + and not self.is_newton_backend + ): + max_dof = self.ps.gpu_get_articulation_max_dof() self._target_qpos = torch.zeros( (self.num_instances, max_dof), dtype=torch.float32, device=self.device ) @@ -417,14 +438,45 @@ class Articulation(BatchEntity): def __init__( self, cfg: ArticulationCfg, - entities: List[_Articulation] = None, + entities: Sequence[_Articulation | SpawnedArticulation] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: - # Initialize world and physics scene - self._world = dexsim.default_world() - from embodichain.lab.sim.sim_manager import get_physics_scene + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared Articulation requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._world = None + self._ps = None + self._data = None + self._all_indices = torch.arange(declared_num_instances, dtype=torch.int32) + self._visual_material = [{} for _ in range(declared_num_instances)] + self.is_shared_visual_material = False + self._has_collision_visible_node_dict = {} + return + + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result + if spawn_result is None: + # Legacy initialization remains temporarily while SimulationManager + # migration is in progress. Spawn-bound facades never reach for a + # process-global World or raw PhysicsScene. + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene - self._ps = get_physics_scene() + self._ps = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = None self.cfg = cfg self._entities = entities @@ -433,10 +485,23 @@ def __init__( # Store all indices for batch operations self._all_indices = torch.arange(len(entities), dtype=torch.int32) - if device.type == "cuda" and not is_newton_scene(self._ps): + if ( + spawn_result is None + and device.type == "cuda" + and not is_newton_scene(self._ps) + ): self._world.update(0.001) - self._data = ArticulationData(entities=entities, ps=self._ps, device=device) + articulation_view = None + if spawn_result is not None: + batch = spawn_result.create_articulation_batch(entities) + articulation_view = SpawnArticulationView(spawn_result, batch, device) + self._data = ArticulationData( + entities=entities, + ps=self._ps, + device=device, + articulation_view=articulation_view, + ) self.cfg: ArticulationCfg if self.cfg.init_qpos is None: @@ -445,50 +510,7 @@ def __init__( # Get default masses. self.default_link_masses = self.get_mass() - # Determine if we should use USD properties or cfg properties. - if not self.cfg.use_usd_properties: - num_entities = len(entities) - dof = self._data.dof - default_cfg = JointDrivePropertiesCfg() - self.default_joint_damping = torch.full( - (num_entities, dof), - default_cfg.damping, - dtype=torch.float32, - device=device, - ) - self.default_joint_stiffness = torch.full( - (num_entities, dof), - default_cfg.stiffness, - dtype=torch.float32, - device=device, - ) - self.default_joint_max_effort = torch.full( - (num_entities, dof), - default_cfg.max_effort, - dtype=torch.float32, - device=device, - ) - self.default_joint_max_velocity = torch.full( - (num_entities, dof), - default_cfg.max_velocity, - dtype=torch.float32, - device=device, - ) - self.default_joint_friction = torch.full( - (num_entities, dof), - default_cfg.friction, - dtype=torch.float32, - device=device, - ) - self.default_joint_armature = torch.full( - (num_entities, dof), - default_cfg.armature, - dtype=torch.float32, - device=device, - ) - self._set_default_joint_drive() - else: - # Read current properties from USD-loaded entities + if self.cfg.use_usd_properties: self.default_joint_stiffness = self._data.joint_stiffness.clone() self.default_joint_damping = self._data.joint_damping.clone() self.default_joint_friction = self._data.joint_friction.clone() @@ -496,26 +518,49 @@ def __init__( self.default_joint_max_effort = self._data.qf_limits.clone() self.default_joint_max_velocity = self._data.qvel_limits.clone() - # Write the USD properties back to cfg - usd_drive_pros = self.cfg.drive_pros - usd_drive_pros.stiffness = ( - self.default_joint_stiffness[0].cpu().numpy().tolist() - ) - usd_drive_pros.damping = ( - self.default_joint_damping[0].cpu().numpy().tolist() - ) - usd_drive_pros.friction = ( - self.default_joint_friction[0].cpu().numpy().tolist() - ) - usd_drive_pros.armature = ( - self.default_joint_armature[0].cpu().numpy().tolist() - ) - usd_drive_pros.max_effort = ( - self.default_joint_max_effort[0].cpu().numpy().tolist() - ) - usd_drive_pros.max_velocity = ( - self.default_joint_max_velocity[0].cpu().numpy().tolist() - ) + if spawn_result is None: + usd_drive_pros = self.cfg.drive_pros + usd_drive_pros.stiffness = ( + self.default_joint_stiffness[0].cpu().numpy().tolist() + ) + usd_drive_pros.damping = ( + self.default_joint_damping[0].cpu().numpy().tolist() + ) + usd_drive_pros.friction = ( + self.default_joint_friction[0].cpu().numpy().tolist() + ) + usd_drive_pros.armature = ( + self.default_joint_armature[0].cpu().numpy().tolist() + ) + usd_drive_pros.max_effort = ( + self.default_joint_max_effort[0].cpu().numpy().tolist() + ) + usd_drive_pros.max_velocity = ( + self.default_joint_max_velocity[0].cpu().numpy().tolist() + ) + else: + default_cfg = JointDrivePropertiesCfg() + values = { + "default_joint_damping": default_cfg.damping, + "default_joint_stiffness": default_cfg.stiffness, + "default_joint_max_effort": default_cfg.max_effort, + "default_joint_max_velocity": default_cfg.max_velocity, + "default_joint_friction": default_cfg.friction, + "default_joint_armature": default_cfg.armature, + } + for name, value in values.items(): + setattr( + self, + name, + torch.full( + (len(self._entities), self._data.dof), + float(value), + dtype=torch.float32, + device=self.device, + ), + ) + if spawn_result is None: + self._set_default_joint_drive() # Apply configured qpos limits if provided. This replaces the asset # limits as the baseline and allows expanding the allowed range. @@ -543,10 +588,17 @@ def __init__( self.set_qpos_limits(qpos_limits) self.pk_chain = None - if self.cfg.build_pk_chain: + is_usd_source = str(self.cfg.fpath).lower().endswith((".usd", ".usda", ".usdc")) + if self.cfg.build_pk_chain and not is_usd_source: self.pk_chain = create_pk_chain( urdf_path=self.cfg.fpath, device=self.device ) + elif self.cfg.build_pk_chain: + logger.log_warning( + f"Articulation {self.uid!r} uses USD for simulation; skipping " + "the URDF-only pk_chain. Configure a solver with its matching " + "URDF when kinematics are required." + ) # For rendering purposes, each articulation can have multiple material instances associated with its links. self._visual_material: List[Dict[str, VisualMaterialInst]] = [ @@ -560,7 +612,11 @@ def __init__( self.active_joint_ids = [i for i in range(self.dof) if i not in self.mimic_ids] # TODO: very weird that we must call update here to make sure the GPU indices are valid. - if device.type == "cuda" and not is_newton_scene(self._ps): + if ( + spawn_result is None + and device.type == "cuda" + and not is_newton_scene(self._ps) + ): self._world.update(0.001) super().__init__(cfg, entities, device) @@ -568,14 +624,148 @@ def __init__( self._initialize_existing_visual_material() # set default collision filter - self._set_default_collision_filter() + if spawn_result is None: + self._set_default_collision_filter() # flag for collision visible node existence self._has_collision_visible_node_dict = dict() for link_name in self.link_names: self._has_collision_visible_node_dict[link_name] = False + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized SpawnResult.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its SpawnResult binding.""" + return self._spawn_result is None and len(self._entities) == 0 + + @property + def num_instances(self) -> int: + if self._entities: + return len(self._entities) + return self._declared_num_instances + + def bind_spawn( + self, + result: SpawnResult, + entities: Sequence[SpawnedArticulation], + ) -> None: + """Initialize this declared facade from Spawn articulation handles.""" + if self.is_spawn_bound: + raise RuntimeError(f"Articulation {self.uid!r} is already Spawn-bound.") + if not self.is_declared: + raise RuntimeError( + f"Articulation {self.uid!r} was not created as a Spawn declaration." + ) + if len(entities) != self._declared_num_instances: + raise ValueError( + f"Articulation {self.uid!r} expected " + f"{self._declared_num_instances} Spawn handles, got {len(entities)}." + ) + + cfg = self.cfg + device = self.device + type(self).__init__( + self, + cfg, + list(entities), + device, + spawn_result=result, + ) + self._apply_spawn_config() + + def _apply_spawn_config(self) -> None: + """Apply config values that require finalized source metadata. + + The source file is loaded only by the DexSim Spawn adapter. This + method runs after binding, when canonical link and active-joint names + are available. + """ + is_usd = str(self.cfg.fpath).lower().endswith((".usd", ".usda", ".usdc")) + use_source_properties = is_usd and self.cfg.use_usd_properties + if use_source_properties: + return + + self._set_default_joint_drive() + self._apply_configured_link_masses() + + if self.cfg.compute_uv: + for entity in self._entities: + for link_name in self.link_names: + render_body = entity.get_render_body(link_name) + if render_body is not None: + render_body.set_projective_uv() + + logger.log_warning( + f"Spawn articulation {self.uid!r}: TODO: non-mass link physics " + "attributes are not exposed by DexSim SpawnedArticulation." + ) + + def _apply_configured_link_masses(self) -> None: + """Apply configured masses after source link names are available.""" + base_mass = self.cfg.attrs.mass + groups = self.cfg.link_attrs or {} + if base_mass is None and not any( + group.attrs.mass is not None for group in groups.values() + ): + return + if self.body_data.is_newton_backend: + logger.log_warning( + f"Spawn articulation {self.uid!r}: Newton link-mass overrides " + "require retained-desc support and were not applied." + ) + return + + masses = self.get_mass() + mass_changed = False + if base_mass is not None: + if base_mass == 0: + logger.log_warning( + f"Spawn articulation {self.uid!r}: density-derived mass is " + "not exposed by the Spawn facade and was not applied." + ) + else: + masses.fill_(float(base_mass)) + mass_changed = True + + claimed: set[str] = set() + for group in groups.values(): + if group.attrs.mass is None: + continue + if group.attrs.mass == 0: + logger.log_warning( + f"Spawn articulation {self.uid!r}: density-derived per-link " + "mass is not exposed by the Spawn facade and was not applied." + ) + continue + matched_indices, matched_names = resolve_matching_names( + keys=group.link_names_expr, + list_of_strings=self.link_names, + ) + overlap = claimed.intersection(matched_names) + if overlap: + raise ValueError( + "Articulation link mass override groups overlap for links " + f"{sorted(overlap)}." + ) + claimed.update(matched_names) + masses[:, matched_indices] = float(group.attrs.mass) + mass_changed = True + + if mass_changed: + self.set_mass(masses, self.link_names) + self.default_link_masses = self.get_mass() + def __str__(self) -> str: + if self.is_declared: + parent_str = ( + f"{self.__class__}: declared {self.num_instances} Spawn " + f"articulations | uid: {self.uid} | device: {self.device}" + ) + return parent_str parent_str = super().__str__() return parent_str + f" | dof: {self.dof} | num_links: {self.num_links}" @@ -1315,7 +1505,9 @@ def set_mass( for i, env_idx in enumerate(local_env_ids): for j, name in enumerate(link_names): - if self._data.is_newton_backend: + if self.is_spawn_bound: + self._entities[env_idx].set_link_mass(name, mass[i, j].item()) + elif self._data.is_newton_backend: local_name = self._entity_link_name(env_idx, name) self._entities[env_idx].set_link_mass(local_name, mass[i, j].item()) else: @@ -1353,7 +1545,15 @@ def get_mass( ) for i, env_idx in enumerate(local_env_ids): for j, name in enumerate(link_names): - if self._data.is_newton_backend: + if self.is_spawn_bound: + status, values = self._entities[env_idx].get_link_mass(name) + if status < 0 or name not in values: + raise RuntimeError( + f"Spawn articulation {self.uid!r} did not expose " + f"mass for link {name!r} in row {env_idx}." + ) + mass_tensor[i, j] = values[name] + elif self._data.is_newton_backend: local_name = self._entity_link_name(env_idx, name) mass_tensor[i, j] = self._entities[env_idx].get_link_mass( local_name @@ -1501,6 +1701,34 @@ def _drive_arg(value: torch.Tensor, index: int) -> float | np.ndarray: return result.item() if result.size == 1 else result for i, env_idx in enumerate(local_env_ids): + if self.is_spawn_bound and self.body_data.is_newton_backend: + if drive_type == "acceleration": + raise NotImplementedError( + "Newton Spawn does not have an exact equivalent of " + "DexSim's acceleration drive. Use drive_type='force' " + "or provide a Newton-native drive descriptor." + ) + if drive_type not in {"force", "none"}: + raise ValueError(f"Unsupported joint drive type {drive_type!r}.") + drive_args = { + "target_mode": 3 if drive_type == "force" else 0, + "joint_ids": local_joint_ids, + } + if stiffness is not None: + drive_args["target_ke"] = _drive_arg(stiffness, i) + if damping is not None: + drive_args["target_kd"] = _drive_arg(damping, i) + if max_effort is not None: + drive_args["effort_limit"] = _drive_arg(max_effort, i) + if max_velocity is not None: + drive_args["velocity_limit"] = _drive_arg(max_velocity, i) + if friction is not None: + drive_args["friction"] = _drive_arg(friction, i) + if armature is not None: + drive_args["armature"] = _drive_arg(armature, i) + self._entities[env_idx].set_newton_drive(**drive_args) + continue + drive_args = { "drive_type": get_dexsim_drive_type(drive_type), "joint_ids": local_joint_ids, @@ -1753,24 +1981,36 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: self.restore_visual_material(env_ids=local_env_ids) - pos = torch.as_tensor( - self.cfg.init_pos, dtype=torch.float32, device=self.device - ) - rot = ( - torch.as_tensor(self.cfg.init_rot, dtype=torch.float32, device=self.device) - * torch.pi - / 180.0 - ) - pos = pos.unsqueeze(0).repeat(num_instances, 1) - rot = rot.unsqueeze(0).repeat(num_instances, 1) - mat = matrix_from_euler(rot, "XYZ") - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze(0) - .repeat(num_instances, 1, 1) - ) - pose[:, :3, 3] = pos - pose[:, :3, :3] = mat + if self.cfg.init_local_pose is not None: + pose = ( + torch.as_tensor( + self.cfg.init_local_pose, + dtype=torch.float32, + device=self.device, + ) + .reshape(1, 4, 4) + .repeat(num_instances, 1, 1) + ) + else: + pos = torch.as_tensor( + self.cfg.init_pos, dtype=torch.float32, device=self.device + ) + rot = ( + torch.as_tensor( + self.cfg.init_rot, dtype=torch.float32, device=self.device + ) + * torch.pi + / 180.0 + ) + pos = pos.unsqueeze(0).repeat(num_instances, 1) + rot = rot.unsqueeze(0).repeat(num_instances, 1) + pose = ( + torch.eye(4, dtype=torch.float32, device=self.device) + .unsqueeze(0) + .repeat(num_instances, 1, 1) + ) + pose[:, :3, 3] = pos + pose[:, :3, :3] = matrix_from_euler(rot, "XYZ") self.set_local_pose(pose, env_ids=local_env_ids) qpos = torch.as_tensor( @@ -1787,11 +2027,17 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: if self.device.type == "cpu" and not self._data.is_newton_backend: self._world.update(0.001) - def _set_default_joint_drive(self) -> None: + def _set_default_joint_drive( + self, + drive_pros: JointDrivePropertiesCfg | dict | None = None, + ) -> None: """Set default joint drive parameters based on the configuration.""" import numbers from embodichain.utils.string import resolve_matching_names_values + if drive_pros is None: + drive_pros = self.cfg.drive_pros + drive_props = [ ("damping", self.default_joint_damping), ("stiffness", self.default_joint_stiffness), @@ -1802,7 +2048,11 @@ def _set_default_joint_drive(self) -> None: ] for prop_name, default_array in drive_props: - value = getattr(self.cfg.drive_pros, prop_name, None) + value = ( + drive_pros.get(prop_name) + if isinstance(drive_pros, dict) + else getattr(drive_pros, prop_name, None) + ) if value is None: continue if isinstance(value, numbers.Number): @@ -1818,7 +2068,6 @@ def _set_default_joint_drive(self) -> None: except Exception as e: logger.log_error(f"Failed to set {prop_name}: {e}") - drive_pros = self.cfg.drive_pros if isinstance(drive_pros, dict): drive_type = drive_pros.get("drive_type", "none") else: @@ -2335,6 +2584,9 @@ def set_self_collision( ) def destroy(self) -> None: + if self.is_declared or self.is_spawn_bound: + # SpawnResult is the sole owner of native lifetime. + return env = self._world.get_env() arenas = env.get_all_arenas() if len(arenas) == 0: diff --git a/embodichain/lab/sim/objects/backends/__init__.py b/embodichain/lab/sim/objects/backends/__init__.py index 538afeb1b..3d039017d 100644 --- a/embodichain/lab/sim/objects/backends/__init__.py +++ b/embodichain/lab/sim/objects/backends/__init__.py @@ -23,6 +23,7 @@ apply_collision_filter_for_envs, is_newton_scene, ) +from .spawn import SpawnArticulationView, SpawnRigidBodyView __all__ = [ "ArticulationViewBase", @@ -34,4 +35,6 @@ "apply_collision_filter_for_entities", "apply_collision_filter_for_envs", "is_newton_scene", + "SpawnArticulationView", + "SpawnRigidBodyView", ] diff --git a/embodichain/lab/sim/objects/backends/newton.py b/embodichain/lab/sim/objects/backends/newton.py index 735d68fcb..0b1e3c39f 100644 --- a/embodichain/lab/sim/objects/backends/newton.py +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -15,12 +15,11 @@ # ---------------------------------------------------------------------------- from __future__ import annotations -from typing import Sequence +from typing import TYPE_CHECKING, Any, Sequence import numpy as np import torch from dexsim.models import MeshObject -from dexsim.engine.newton_physics import NewtonPhysicsScene from embodichain.lab.sim.objects.backends.base import ( ArticulationViewBase, RigidBodyViewBase, @@ -28,6 +27,11 @@ from embodichain.utils import logger from embodichain.utils.math import matrix_from_quat, quat_from_matrix +if TYPE_CHECKING: + from dexsim.engine.newton_physics.newton_physics_scene import NewtonPhysicsScene +else: + NewtonPhysicsScene = Any + __all__ = [ "NewtonRigidBodyView", "NewtonArticulationView", diff --git a/embodichain/lab/sim/objects/backends/spawn.py b/embodichain/lab/sim/objects/backends/spawn.py new file mode 100644 index 000000000..1a9c3552e --- /dev/null +++ b/embodichain/lab/sim/objects/backends/spawn.py @@ -0,0 +1,633 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""EmbodiChain tensor-layout adapters for :mod:`dexsim.spawn` batches. + +The classes in this module deliberately know nothing about PhysX scenes or +Newton runtime objects. Backend selection, handle rebinding, and topology +revision tracking remain owned by DexSim's ``SpawnResult`` and batch classes. +EmbodiChain only adapts logical row selections and its public pose convention +``(x, y, z, qx, qy, qz, qw)``. + +DexSim does not yet expose lightweight row/DOF/link selections on its public +batches. Until that API lands, partial writes use a correctness-first +read/modify/write fallback. The fallback is kept here, at the boundary, so it +can be deleted without changing object or environment APIs. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Sequence + +import torch + +from .base import ArticulationViewBase, RigidBodyViewBase + +if TYPE_CHECKING: + from dexsim.spawn import ArticulationBatch, RigidBodyBatch, SpawnResult + +__all__ = ["SpawnArticulationView", "SpawnRigidBodyView"] + + +def _rows( + selection: Sequence[int] | torch.Tensor | None, + count: int, + device: torch.device, +) -> torch.Tensor: + if selection is None: + return torch.arange(count, dtype=torch.long, device=device) + result = torch.as_tensor(selection, dtype=torch.long, device=device).reshape(-1) + if torch.any(result < 0) or torch.any(result >= count): + raise IndexError(f"Batch row selection is outside [0, {count}).") + return result + + +def _spawn_pose(data: torch.Tensor) -> torch.Tensor: + """Convert rigid-body ``xyz+xyzw`` poses to Spawn ``xyzw+xyz``.""" + result = torch.empty_like(data, dtype=torch.float32) + result[..., 0:4] = data[..., 3:7] + result[..., 4:7] = data[..., 0:3] + return result + + +def _embodichain_pose(data: torch.Tensor) -> torch.Tensor: + """Convert Spawn ``xyzw+xyz`` poses to rigid-body ``xyz+xyzw``.""" + result = torch.empty_like(data, dtype=torch.float32) + result[..., 0:3] = data[..., 4:7] + result[..., 3:7] = data[..., 0:4] + return result + + +def _spawn_articulation_pose(data: torch.Tensor) -> torch.Tensor: + """Convert articulation ``xyz+wxyz`` poses to Spawn ``xyzw+xyz``.""" + result = torch.empty_like(data, dtype=torch.float32) + result[..., 0:3] = data[..., 4:7] + result[..., 3] = data[..., 3] + result[..., 4:7] = data[..., 0:3] + return result + + +def _embodichain_articulation_pose(data: torch.Tensor) -> torch.Tensor: + """Convert Spawn ``xyzw+xyz`` poses to articulation ``xyz+wxyz``.""" + result = torch.empty_like(data, dtype=torch.float32) + result[..., 0:3] = data[..., 4:7] + result[..., 3] = data[..., 3] + result[..., 4:7] = data[..., 0:3] + return result + + +class _SpawnSelectionAdapter: + """Shared correctness-first selection support for fixed-size Spawn batches.""" + + def __init__(self, batch: Any, device: torch.device, row_count: int) -> None: + self._batch = batch + self.device = device + self._row_count = row_count + + def _fetch_rows( + self, + method_name: str, + out: torch.Tensor, + selection: Sequence[int] | torch.Tensor | None, + tail_shape: tuple[int, ...], + ) -> torch.Tensor: + rows = _rows(selection, self._row_count, self.device) + full = torch.empty( + (self._row_count, *tail_shape), + dtype=torch.float32, + device=self.device, + ) + getattr(self._batch, method_name)(full) + selected = full.index_select(0, rows) + out.copy_(selected.to(device=out.device, dtype=out.dtype)) + return out + + def _apply_rows( + self, + method_name: str, + values: torch.Tensor, + selection: Sequence[int] | torch.Tensor, + tail_shape: tuple[int, ...], + *, + fetch_method_name: str | None, + ) -> None: + rows = _rows(selection, self._row_count, self.device) + values = values.to(device=self.device, dtype=torch.float32) + expected_shape = (len(rows), *tail_shape) + if tuple(values.shape) != expected_shape: + raise ValueError( + f"Expected selected data shape {expected_shape}, got " + f"{tuple(values.shape)}." + ) + + if fetch_method_name is None: + full = torch.zeros( + (self._row_count, *tail_shape), + dtype=torch.float32, + device=self.device, + ) + else: + full = torch.empty( + (self._row_count, *tail_shape), + dtype=torch.float32, + device=self.device, + ) + getattr(self._batch, fetch_method_name)(full) + full.index_copy_(0, rows, values) + getattr(self._batch, method_name)(full) + + +class SpawnRigidBodyView(_SpawnSelectionAdapter, RigidBodyViewBase): + """Backend-neutral rigid-body view backed by ``RigidBodyBatch``.""" + + def __init__( + self, + result: SpawnResult, + batch: RigidBodyBatch, + device: torch.device, + ) -> None: + super().__init__(batch, device, len(batch)) + self.result = result + self.batch = batch + self._body_ids_tensor = torch.arange( + len(batch), dtype=torch.int32, device=device + ) + + @property + def is_ready(self) -> bool: + return True + + @property + def is_newton_backend(self) -> bool: + return self.result.backend == "newton" + + @property + def body_ids(self) -> list[int]: + return list(range(self._row_count)) + + @property + def body_ids_tensor(self) -> torch.Tensor: + return self._body_ids_tensor + + def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> torch.Tensor: + return self._body_ids_tensor[indices] + + def fetch_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + spawn = torch.empty((len(data), 7), dtype=torch.float32, device=self.device) + self._fetch_rows("fetch_pose", spawn, body_ids, (7,)) + data.copy_(_embodichain_pose(spawn).to(data.device, data.dtype)) + + def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_pose", + _spawn_pose(pose.to(self.device, torch.float32)), + body_ids, + (7,), + fetch_method_name="fetch_pose", + ) + + def fetch_com_local_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + spawn = torch.empty((len(data), 7), dtype=torch.float32, device=self.device) + self._fetch_rows("fetch_com_local_pose", spawn, body_ids, (7,)) + data.copy_(_embodichain_pose(spawn).to(data.device, data.dtype)) + + def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_com_local_pose", + _spawn_pose(data.to(self.device, torch.float32)), + body_ids, + (7,), + fetch_method_name="fetch_com_local_pose", + ) + + def fetch_linear_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_linear_velocity", data, body_ids, (3,)) + + def fetch_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_angular_velocity", data, body_ids, (3,)) + + def apply_linear_velocity(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_linear_velocity", + data, + body_ids, + (3,), + fetch_method_name="fetch_linear_velocity", + ) + + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_rows( + "apply_angular_velocity", + data, + body_ids, + (3,), + fetch_method_name="fetch_angular_velocity", + ) + + def fetch_linear_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_linear_acceleration", data, body_ids, (3,)) + + def fetch_angular_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_angular_acceleration", data, body_ids, (3,)) + + def apply_force(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_force", data, body_ids, (3,), fetch_method_name=None) + + def apply_torque(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_torque", data, body_ids, (3,), fetch_method_name=None) + + def fetch_mass( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_mass", data, body_ids, (1,)) + + def apply_mass(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_mass", data, body_ids, (1,), fetch_method_name="fetch_mass" + ) + + def fetch_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_inertia_diagonal", data, body_ids, (3,)) + + def apply_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_rows( + "apply_inertia_diagonal", + data, + body_ids, + (3,), + fetch_method_name="fetch_inertia_diagonal", + ) + + @staticmethod + def _unsupported_property(name: str) -> None: + raise NotImplementedError( + f"DexSim Spawn RigidBodyBatch does not expose the {name} property yet. " + "Extend the public Spawn batch instead of accessing backend internals." + ) + + def fetch_friction( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + del data, body_ids + self._unsupported_property("friction") + + def apply_friction(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + del data, body_ids + self._unsupported_property("friction") + + def fetch_restitution( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + del data, body_ids + self._unsupported_property("restitution") + + def apply_restitution(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + del data, body_ids + self._unsupported_property("restitution") + + def fetch_contact_offset( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + del data, body_ids + self._unsupported_property("contact_offset") + + def apply_contact_offset(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + del data, body_ids + self._unsupported_property("contact_offset") + + +class SpawnArticulationView(_SpawnSelectionAdapter, ArticulationViewBase): + """Backend-neutral articulation state view backed by ``ArticulationBatch``. + + Joint selections currently require one scalar DOF per selected joint. The + public DexSim layout already describes multi-DOF joints; supporting them + without ambiguity requires a DOF-selection API in DexSim and is therefore + kept as an explicit boundary rather than guessed here. + """ + + def __init__( + self, + result: SpawnResult, + batch: ArticulationBatch, + device: torch.device, + ) -> None: + super().__init__(batch, device, len(batch)) + self.result = result + self.batch = batch + self._validate_homogeneous_layout() + self._articulation_ids = torch.arange( + len(batch), dtype=torch.int32, device=device + ) + + def _validate_homogeneous_layout(self) -> None: + """Require the uniform topology promised by one EC Articulation.""" + dof_counts = tuple(self.batch.dof_counts) + link_counts = tuple(self.batch.link_counts) + joint_names = tuple(self.batch.joint_names_per_articulation) + link_names = tuple(self.batch.link_names_per_articulation) + if dof_counts and len(set(dof_counts)) != 1: + raise ValueError( + "One EmbodiChain Articulation cannot bind heterogeneous Spawn " + f"DOF counts: {dof_counts}." + ) + if link_counts and len(set(link_counts)) != 1: + raise ValueError( + "One EmbodiChain Articulation cannot bind heterogeneous Spawn " + f"link counts: {link_counts}." + ) + if joint_names and any(names != joint_names[0] for names in joint_names[1:]): + raise ValueError( + "One EmbodiChain Articulation requires identical active-joint " + "ordering in every Spawn row." + ) + if link_names and any(names != link_names[0] for names in link_names[1:]): + raise ValueError( + "One EmbodiChain Articulation requires identical link ordering " + "in every Spawn row." + ) + layouts = tuple(self.batch.joint_layouts_per_articulation) + if layouts and any(layout.dof_count != 1 for layout in layouts[0]): + raise NotImplementedError( + "EmbodiChain's Articulation API currently indexes joints and " + "scalar DOFs interchangeably. Spawn multi-DOF joints require " + "an explicit DOF-selection API before they can be bound safely." + ) + + @property + def dof(self) -> int: + """Scalar DOF width shared by every articulation row.""" + return self.batch.dof_width + + @property + def num_links(self) -> int: + """Link count shared by every articulation row.""" + return self.batch.link_width + + @property + def joint_names(self) -> list[str]: + """Active joints in public flattened-DOF order.""" + rows = self.batch.joint_names_per_articulation + return [] if not rows else list(rows[0]) + + @property + def link_names(self) -> list[str]: + """Links in public link-buffer order.""" + rows = self.batch.link_names_per_articulation + return [] if not rows else list(rows[0]) + + @property + def is_ready(self) -> bool: + return True + + @property + def is_newton_backend(self) -> bool: + return self.result.backend == "newton" + + @property + def articulation_ids_tensor(self) -> torch.Tensor: + return self._articulation_ids + + def select_articulation_ids( + self, env_ids: Sequence[int] | torch.Tensor + ) -> torch.Tensor: + return self._articulation_ids[env_ids] + + def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: + spawn = torch.empty_like(data, dtype=torch.float32, device=self.device) + self.batch.fetch_root_pose(spawn) + data.copy_(_embodichain_articulation_pose(spawn).to(data.device, data.dtype)) + return data + + def fetch_root_linear_velocity(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_root_linear_velocity(data) + return data + + def fetch_root_angular_velocity(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_root_angular_velocity(data) + return data + + def fetch_qpos(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_joint_position(data) + return data + + def fetch_target_qpos(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_joint_target_position(data) + return data + + def fetch_qvel(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_joint_velocity(data) + return data + + def fetch_target_qvel(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_joint_target_velocity(data) + return data + + def fetch_qacc(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_joint_acceleration(data) + return data + + def fetch_qf(self, data: torch.Tensor) -> torch.Tensor: + self.batch.fetch_joint_force(data) + return data + + def fetch_link_pose(self, data: torch.Tensor) -> torch.Tensor: + spawn = torch.empty_like(data, dtype=torch.float32, device=self.device) + self.batch.fetch_link_pose(spawn) + data.copy_(_embodichain_articulation_pose(spawn).to(data.device, data.dtype)) + return data + + def fetch_link_velocity( + self, + data: torch.Tensor, + linear_data: torch.Tensor, + angular_data: torch.Tensor, + ) -> torch.Tensor: + self.batch.fetch_link_linear_velocity(linear_data) + self.batch.fetch_link_angular_velocity(angular_data) + data[..., 0:3] = linear_data + data[..., 3:6] = angular_data + return data + + def apply_root_pose( + self, pose: torch.Tensor, env_ids: Sequence[int] | torch.Tensor + ) -> None: + self._apply_rows( + "apply_root_pose", + _spawn_articulation_pose(pose.to(self.device, torch.float32)), + env_ids, + (7,), + fetch_method_name="fetch_root_pose", + ) + + def _joint_columns(self, joint_ids: Sequence[int] | torch.Tensor) -> torch.Tensor: + ids = torch.as_tensor(joint_ids, dtype=torch.long, device=self.device) + layouts = self.batch.joint_layouts_per_articulation + if not layouts: + return ids + reference = layouts[0] + columns: list[int] = [] + for joint_id in ids.detach().cpu().tolist(): + layout = reference[joint_id] + if layout.dof_count != 1: + raise NotImplementedError( + "SpawnArticulationView needs DexSim DOF selection for " + f"multi-DOF joint {layout.name!r}." + ) + columns.append(layout.dof_start) + return torch.as_tensor(columns, dtype=torch.long, device=self.device) + + def _apply_joint_selection( + self, + values: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + apply_method: str, + fetch_method: str | None, + ) -> None: + rows = _rows(env_ids, self._row_count, self.device) + columns = self._joint_columns(joint_ids) + values = values.to(device=self.device, dtype=torch.float32) + expected = (len(rows), len(columns)) + if tuple(values.shape) != expected: + raise ValueError( + f"Expected selected joint data shape {expected}, got " + f"{tuple(values.shape)}." + ) + width = self.batch.dof_width + if fetch_method is None: + full = torch.zeros( + (self._row_count, width), + dtype=torch.float32, + device=self.device, + ) + else: + full = torch.empty( + (self._row_count, width), + dtype=torch.float32, + device=self.device, + ) + getattr(self.batch, fetch_method)(full) + full[rows[:, None], columns] = values + getattr(self.batch, apply_method)(full) + + def apply_qpos( + self, + qpos: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + self._apply_joint_selection( + qpos, + env_ids, + joint_ids, + apply_method=( + "apply_joint_target_position" if target else "apply_joint_position" + ), + fetch_method=( + "fetch_joint_target_position" if target else "fetch_joint_position" + ), + ) + + def apply_qvel( + self, + qvel: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + self._apply_joint_selection( + qvel, + env_ids, + joint_ids, + apply_method=( + "apply_joint_target_velocity" if target else "apply_joint_velocity" + ), + fetch_method=( + "fetch_joint_target_velocity" if target else "fetch_joint_velocity" + ), + ) + + def apply_qf( + self, + qf: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + ) -> None: + self._apply_joint_selection( + qf, + env_ids, + joint_ids, + apply_method="apply_joint_force", + fetch_method=None, + ) + + def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + rows = _rows(env_ids, self._row_count, self.device) + zeros = torch.zeros( + (len(rows), self.batch.dof_width), + dtype=torch.float32, + device=self.device, + ) + self._apply_rows( + "apply_joint_velocity", + zeros, + rows, + (self.batch.dof_width,), + fetch_method_name="fetch_joint_velocity", + ) + self._apply_rows( + "apply_joint_target_velocity", + zeros, + rows, + (self.batch.dof_width,), + fetch_method_name="fetch_joint_target_velocity", + ) + self._apply_rows( + "apply_joint_force", + zeros, + rows, + (self.batch.dof_width,), + fetch_method_name=None, + ) + + def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + # DexSim currently refreshes the complete batch. Since this operation + # only propagates already-authored state, that is equivalent to a row + # selection and keeps selection details out of EmbodiChain. + del env_ids + if self.batch.compute_kinematics() < 0: + raise RuntimeError("DexSim Spawn articulation kinematics update failed.") diff --git a/embodichain/lab/sim/objects/cloth_object.py b/embodichain/lab/sim/objects/cloth_object.py index 6cbef6a8e..c61bbdfc7 100644 --- a/embodichain/lab/sim/objects/cloth_object.py +++ b/embodichain/lab/sim/objects/cloth_object.py @@ -19,10 +19,11 @@ import torch import dexsim import numpy as np +from copy import deepcopy from functools import cached_property from dataclasses import dataclass -from typing import List, Sequence, Union +from typing import Any, List, Sequence, TYPE_CHECKING, Union from dexsim.models import MeshObject from dexsim.engine import ClothBody, PhysicsScene @@ -47,6 +48,9 @@ ) from embodichain.utils.math import xyz_quat_to_4x4_matrix +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult + __all__ = ["ClothBodyData", "ClothObject", "ClothObjectCfg"] @@ -86,7 +90,7 @@ def __init__( dtype=torch.float32, ) for i, cloth_body in enumerate(self.cloth_bodies): - self._rest_position_buffer[i] = cloth_body.get_position_inv_mass_buffer() + self._rest_position_buffer[i] = cloth_body.get_rest_position_buffer() self._vertex_position = torch.zeros( (self.num_instances, self.n_vertices, 3), @@ -126,21 +130,52 @@ class ClothObject(BatchEntity): def __init__( self, cfg: ClothObjectCfg, - entities: List[MeshObject] = None, + entities: Sequence[Any] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: - self._world: dexsim.World = dexsim.default_world() - from embodichain.lab.sim.sim_manager import get_physics_scene + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared ClothObject requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._world = None + self._ps = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._visual_material = [None] * declared_num_instances + self.is_shared_visual_material = False + return - self._ps = get_physics_scene() + entities = list(entities) + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result + if spawn_result is None: + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = self._world.get_physics_scene() self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() self._data = ClothBodyData(entities=entities, ps=self._ps, device=device) - self._world.update(0.001) + if spawn_result is None: + self._world.update(0.001) self._surface_triangles = self._build_surface_triangles( entities[0], self._data.rest_vertices[0].detach().cpu().numpy(), + self._data.cloth_bodies[0].get_initial_transform(), ) self._visual_material: List[VisualMaterialInst | None] = [None] * len(entities) @@ -153,10 +188,49 @@ def __init__( self._set_default_collision_filter() + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized SpawnResult.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its SpawnResult binding.""" + return self._spawn_result is None and len(self._entities) == 0 + + @property + def num_instances(self) -> int: + return len(self._entities) if self._entities else self._declared_num_instances + + def bind_spawn(self, result: SpawnResult, entities: Sequence[Any]) -> None: + """Bind a declared facade to finalized cloth handles in place.""" + if len(entities) != self._declared_num_instances: + raise ValueError( + f"ClothObject {self.uid!r} expected {self._declared_num_instances} " + f"Spawn handles, got {len(entities)}." + ) + bound = ClothObject( + self.cfg, + entities, + self.device, + spawn_result=result, + ) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) + + def __str__(self) -> str: + if self.is_declared: + return ( + f"{self.__class__}: declared {self.num_instances} Spawn cloth " + f"objects | uid: {self.uid} | device: {self.device}" + ) + return super().__str__() + @staticmethod def _build_surface_triangles( entity: MeshObject, rest_vertices: np.ndarray, + initial_transform: np.ndarray, ) -> np.ndarray: """Map render triangles onto DexSim's welded cloth vertex buffer.""" render_body = entity.get_render_body() @@ -178,6 +252,10 @@ def _build_surface_triangles( vertices = np.concatenate(render_vertices, axis=0) triangles = np.concatenate(render_triangles, axis=0) + initial_transform = np.asarray(initial_transform, dtype=np.float32).reshape( + 4, 4 + ) + vertices = vertices @ initial_transform[:3, :3].T + initial_transform[:3, 3] distances, cloth_vertex_ids = cKDTree(rest_vertices).query(vertices) scale = max(float(np.ptp(rest_vertices, axis=0).max()), 1.0) if float(distances.max(initial=0.0)) > scale * 1.0e-5: @@ -386,20 +464,26 @@ def set_local_pose( arena_offsets = sim.arena_offsets for i, env_idx in enumerate(local_env_ids): # TODO: cloth body cannot directly set by `set_local_pose` currently. - rest_vertices = self.body_data.rest_vertices[i] + cloth_body: ClothBody = self._entities[env_idx].get_physical_body() + rest_vertices = self.body_data.rest_vertices[env_idx] + initial_transform = torch.as_tensor( + cloth_body.get_initial_transform(), + dtype=torch.float32, + device=self.device, + ) + rest_vertices_local = ( + rest_vertices - initial_transform[:3, 3] + ) @ initial_transform[:3, :3] rotation = pose4x4[i][:3, :3] translation = pose4x4[i][:3, 3] - # apply transformation to local rest vertices and back - rest_vertices_local = rest_vertices - arena_offsets[i] transformed_vertices = rest_vertices_local @ rotation.T + translation - transformed_vertices = transformed_vertices + arena_offsets[i] + transformed_vertices = transformed_vertices + arena_offsets[env_idx] - cloth_body: ClothBody = self._entities[env_idx].get_physical_body() position_buffer = cloth_body.get_position_inv_mass_buffer() velocity_buffer = cloth_body.get_velocity_buffer() position_buffer[:, :3] = transformed_vertices - velocity_buffer[:, 3:] = 0.0 + velocity_buffer[:, :3] = 0.0 cloth_body.mark_dirty(ClothBodyGPUAPIReadWriteType.ALL) # TODO: currently cloth body has no wake up interface, use set_wake_counter and pass in a positive value to wake it up @@ -448,6 +532,8 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: self.set_local_pose(pose, env_ids=local_env_ids) def destroy(self) -> None: + if self.is_spawn_bound: + return # TODO: not tested yet env = self._world.get_env() arenas = env.get_all_arenas() diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index 49841a419..572061e3c 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -20,8 +20,9 @@ import dexsim import numpy as np +from copy import deepcopy from dataclasses import dataclass, MISSING -from typing import List, Sequence, Union +from typing import TYPE_CHECKING, List, Sequence, Union from functools import cached_property from dexsim.models import MeshObject @@ -56,6 +57,9 @@ from embodichain.utils.math import matrix_from_quat, quat_from_matrix, matrix_from_euler from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedObject + _UINT64_MAX = (1 << 64) - 1 __all__ = ["RigidBodyData", "RigidObject", "RigidObjectCfg"] @@ -69,7 +73,11 @@ class RigidBodyData: """ def __init__( - self, entities: List[MeshObject], ps: PhysicsScene, device: torch.device + self, + entities: List[MeshObject], + ps: PhysicsScene | None, + device: torch.device, + body_view: RigidBodyViewBase | None = None, ) -> None: """Initialize the RigidBodyData. @@ -84,7 +92,9 @@ def __init__( self.device = device # Create the appropriate backend view. - if is_newton_scene(ps): + if body_view is not None: + self.body_view = body_view + elif is_newton_scene(ps): self.body_view: RigidBodyViewBase = NewtonRigidBodyView( entities=entities, scene=ps, device=device ) @@ -133,7 +143,13 @@ def __init__( @property def is_newton_backend(self) -> bool: - return isinstance(self.body_view, NewtonRigidBodyView) + return bool( + getattr( + self.body_view, + "is_newton_backend", + isinstance(self.body_view, NewtonRigidBodyView), + ) + ) @property def gpu_indices(self) -> torch.Tensor: @@ -227,27 +243,68 @@ def __init__( cfg: RigidObjectCfg, entities: List[MeshObject] = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared RigidObject requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self.body_type = cfg.body_type + self._entities = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._ps = None + self._world = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._visual_material = [None] * declared_num_instances + self.is_shared_visual_material = False + self._has_collision_visible_node = False + return + + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result self.body_type = cfg.body_type - self._world = dexsim.default_world() - from embodichain.lab.sim.sim_manager import get_physics_scene + if spawn_result is None: + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene - self._ps = get_physics_scene() + self._ps = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = None self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() # data for managing body data (only for dynamic and kinematic bodies) on GPU. self._data: RigidBodyData | None = None if self.is_static is False: - self._data = RigidBodyData(entities=entities, ps=self._ps, device=device) + body_view = None + if spawn_result is not None: + from embodichain.lab.sim.objects.backends import SpawnRigidBodyView + + batch = spawn_result.create_rigid_body_batch(entities) + body_view = SpawnRigidBodyView(spawn_result, batch, device) + self._data = RigidBodyData( + entities=entities, + ps=self._ps, + device=device, + body_view=body_view, + ) # For rendering purposes, each instance can have its own material. self._visual_material: List[VisualMaterialInst] = [None] * len(entities) self.is_shared_visual_material = False # Determine if we should use USD properties or cfg properties. - if not cfg.use_usd_properties: + if spawn_result is None and not cfg.use_usd_properties: for entity in entities: entity.set_body_scale(*cfg.body_scale) if is_newton_scene(self._ps): @@ -256,7 +313,7 @@ def __init__( # set_physical_attr() is still default-backend only. continue entity.set_physical_attr(cfg.attrs.attr()) - else: + elif spawn_result is None: # Read current properties from USD-loaded entities and write back to cfg # Use first entity as reference first_entity: MeshObject = entities[0] @@ -271,7 +328,8 @@ def __init__( self._initialize_existing_visual_material() # set default collision filter - self._set_default_collision_filter() + if spawn_result is None: + self._set_default_collision_filter() self._apply_initial_state() @@ -281,15 +339,66 @@ def __init__( # TODO: Must be called after setting all attributes. # May be improved in the future. - if cfg.attrs.enable_collision is False: + if spawn_result is None and cfg.attrs.enable_collision is False: flag = torch.zeros(len(entities), dtype=torch.bool) self.enable_collision(flag) # reserve flag for collision visible node existence self._has_collision_visible_node = False + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized SpawnResult.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its SpawnResult binding.""" + return self._spawn_result is None and len(self._entities) == 0 + + @property + def num_instances(self) -> int: + if self._entities: + return len(self._entities) + return self._declared_num_instances + + def bind_spawn( + self, + result: SpawnResult, + entities: Sequence[SpawnedObject], + ) -> None: + """Bind a declared facade to stable Spawn handles in place.""" + if self.is_spawn_bound: + raise RuntimeError(f"RigidObject {self.uid!r} is already Spawn-bound.") + if len(entities) != self._declared_num_instances: + raise ValueError( + f"RigidObject {self.uid!r} expected {self._declared_num_instances} " + f"Spawn handles, got {len(entities)}." + ) + cfg = self.cfg + device = self.device + # Construct the bound state off to the side. Batch creation may fail + # (for example when a backend/device capability is unavailable); the + # public declaration facade must remain retryable rather than becoming + # half-bound. Replacing the dictionary also drops declaration-time + # cached_property values such as the empty user-id cache. + bound = RigidObject( + cfg, + list(entities), + device, + spawn_result=result, + ) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) + def __str__(self) -> str: - parent_str = super().__str__() + if self.is_declared: + parent_str = ( + f"{self.__class__}: declared {self.num_instances} Spawn objects " + f"| uid: {self.uid} | device: {self.device}" + ) + else: + parent_str = super().__str__() max_hull = self.cfg.max_convex_hull_num if max_hull is MISSING: if isinstance(self.cfg.shape, MeshCfg): @@ -482,6 +591,12 @@ def set_collision_filter( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." ) + if self.is_spawn_bound: + raise NotImplementedError( + "DexSim Spawn does not expose rigid-body collision-filter batch " + "updates yet. The filter must remain in the birth descriptor." + ) + if is_newton_scene(self._ps): if self._data is not None and isinstance( self._data.body_view, NewtonRigidBodyView @@ -746,6 +861,13 @@ def set_attrs( """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self.is_spawn_bound: + raise NotImplementedError( + "RigidObject.set_attrs() needs the remaining typed Spawn property " + "batch APIs (friction/restitution/contact offset). Use the " + "supported set_mass/set_inertia/set_com_pose methods meanwhile." + ) + if isinstance(attrs, List) and len(local_env_ids) != len(attrs): logger.log_error( f"Length of env_ids {len(local_env_ids)} does not match attrs length {len(attrs)}." @@ -956,6 +1078,11 @@ def set_damping( """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self.is_spawn_bound: + raise NotImplementedError( + "DexSim Spawn does not expose rigid-body damping yet." + ) + if len(local_env_ids) != len(damping): logger.log_error( f"Length of env_ids {len(local_env_ids)} does not match damping length {len(damping)}." @@ -990,6 +1117,11 @@ def get_damping(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self.is_spawn_bound: + raise NotImplementedError( + "DexSim Spawn does not expose rigid-body damping yet." + ) + dampings = [] for _, env_idx in enumerate(local_env_ids): if is_newton_scene(self._ps): @@ -1363,6 +1495,12 @@ def set_body_type(self, body_type: str) -> None: """ from dexsim.types import ActorType + if self.is_spawn_bound: + raise NotImplementedError( + "Changing actor topology after Spawn binding requires a public " + "descriptor mutation transaction and is not implemented yet." + ) + if is_newton_scene(self._ps): logger.log_warning( "Newton backend does not support changing RigidObject body type at " @@ -1518,6 +1656,13 @@ def set_physical_visible( if len(rgba) != 4: logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.") + if self.is_spawn_bound: + color = np.asarray(rgba, dtype=np.float32) + for entity in self._entities: + self._spawn_result.set_physical_visible(entity, color, visible) + self._has_collision_visible_node = True + return + # create collision visible node if not exist if visible: if not self._has_collision_visible_node: @@ -1550,6 +1695,16 @@ def set_visible(self, visible: bool = True) -> None: def _build_cfg_init_pose(self, env_ids: Sequence[int]) -> torch.Tensor: """Build initial root poses from cfg as ``(N, 4, 4)`` matrices.""" num_instances = len(env_ids) + if self.cfg.init_local_pose is not None: + return ( + torch.as_tensor( + self.cfg.init_local_pose, + dtype=torch.float32, + device=self.device, + ) + .reshape(1, 4, 4) + .repeat(num_instances, 1, 1) + ) pos = torch.as_tensor( self.cfg.init_pos, dtype=torch.float32, device=self.device ) @@ -1577,6 +1732,19 @@ def _apply_initial_state(self) -> None: ``BUILDER`` via the scene batch API; velocities are cleared after finalization through :meth:`SimulationManager.finalize_newton_physics`. """ + if self.is_spawn_bound: + if self._spawn_result.backend == "dexsim": + # PhysX Direct GPU readiness performs native warm-up updates. + # Re-apply the authored state after the batch becomes usable + # so prepare() itself is not an observable simulation step. + self.reset() + else: + # Newton finalization materializes the descriptor pose without + # advancing simulation; only one-step dynamics buffers need + # clearing after batch binding. + self.clear_dynamics() + return + if is_newton_scene(self._ps): if self._newton_lifecycle_state() == "BUILDER": self.set_local_pose( @@ -1594,8 +1762,9 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: self.restore_visual_material(env_ids=local_env_ids) - # TODO: support attributes setter for newton. - if not is_newton_scene(self._ps): + # Spawn descriptors and their live property APIs are the canonical + # physical configuration; reset changes state only. + if not self.is_spawn_bound and not is_newton_scene(self._ps): self.set_attrs(self.cfg.attrs, env_ids=local_env_ids) self.clear_dynamics(env_ids=local_env_ids) @@ -1605,6 +1774,10 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: ) def destroy(self) -> None: + if self.is_declared or self.is_spawn_bound: + # SimulationManager owns topology removal and SpawnResult lifetime. + # Direct facade destruction must never bypass that owner. + return env = self._world.get_env() arenas = env.get_all_arenas() if len(arenas) == 0: diff --git a/embodichain/lab/sim/objects/rigid_object_group.py b/embodichain/lab/sim/objects/rigid_object_group.py index 0f6192d28..1d10e0c2d 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/embodichain/lab/sim/objects/rigid_object_group.py @@ -16,249 +16,168 @@ from __future__ import annotations -import torch -import dexsim -import numpy as np +from copy import deepcopy +from typing import TYPE_CHECKING, Sequence -from dataclasses import dataclass -from typing import List, Sequence, Union +import numpy as np +import torch -from dexsim.models import MeshObject -from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType -from dexsim.engine import CudaArray, PhysicsScene -from embodichain.lab.sim.cfg import ( - RigidObjectGroupCfg, - RigidBodyAttributesCfg, +from embodichain.lab.sim import BatchEntity +from embodichain.lab.sim.cfg import RigidObjectGroupCfg +from embodichain.lab.sim.material import VisualMaterial +from embodichain.lab.sim.objects.backends.spawn import SpawnRigidBodyView +from embodichain.utils.math import ( + convert_quat, + matrix_from_euler, + matrix_from_quat, + quat_from_matrix, ) -from embodichain.lab.sim import ( - BatchEntity, -) -from embodichain.lab.sim.material import VisualMaterial, VisualMaterialInst -from ._mesh_utils import ( - get_combined_triangles, - get_combined_vertices, -) -from embodichain.utils.math import convert_quat -from embodichain.utils.math import matrix_from_quat, quat_from_matrix, matrix_from_euler -from embodichain.utils import logger + +from ._mesh_utils import get_combined_triangles, get_combined_vertices + +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedObject __all__ = ["RigidBodyGroupData", "RigidObjectGroup", "RigidObjectGroupCfg"] -@dataclass class RigidBodyGroupData: - """Data manager for rigid body group with body type of dynamic or kinematic.""" + """Expose one flat Spawn rigid-body batch as ``[env, object, ...]`` tensors.""" def __init__( - self, entities: List[List[MeshObject]], ps: PhysicsScene, device: torch.device + self, + body_view: SpawnRigidBodyView, + *, + num_instances: int, + num_objects: int, + device: torch.device, ) -> None: - """Initialize the RigidBodyGroupData. - - Args: - entities (List[List[MeshObject]]): List of List MeshObjects representing the rigid body group. - ps (PhysicsScene): The physics scene. - device (torch.device): The device to use for the rigid body group data. - """ - self.entities = entities - self.ps = ps - self.num_instances = len(entities) - self.num_objects = len(entities[0]) + self.body_view = body_view + self.num_instances = num_instances + self.num_objects = num_objects self.device = device - - # get gpu indices for the rigid bodies with shape of (num_instances, num_objects) - self.gpu_indices = ( - torch.as_tensor( - [ - [entity.get_gpu_index() for entity in instance] - for instance in entities - ], - dtype=torch.int32, - device=self.device, - ) - if self.device.type == "cuda" - else None - ) - - # Initialize rigid body group data tensors. Shape of (num_instances, num_objects, data_dim) - self._pose = torch.zeros( - (self.num_instances, self.num_objects, 7), - dtype=torch.float32, - device=self.device, + self._pose = torch.empty( + (num_instances, num_objects, 7), dtype=torch.float32, device=device ) - self._lin_vel = torch.zeros( - (self.num_instances, self.num_objects, 3), - dtype=torch.float32, - device=self.device, - ) - self._ang_vel = torch.zeros( - (self.num_instances, self.num_objects, 3), - dtype=torch.float32, - device=self.device, + self._lin_vel = torch.empty( + (num_instances, num_objects, 3), dtype=torch.float32, device=device ) + self._ang_vel = torch.empty_like(self._lin_vel) @property def pose(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch pose from CPU entities - xyzs = torch.as_tensor( - [ - [entity.get_location() for entity in instance] - for instance in self.entities - ], - device=self.device, - ) - quats = torch.as_tensor( - [ - [entity.get_rotation_quat() for entity in instance] - for instance in self.entities - ], - device=self.device, - ) - quats = convert_quat(quats.reshape(-1, 4), to="wxyz").reshape( - -1, self.num_objects, 4 - ) - return torch.cat((xyzs, quats), dim=-1) - else: - pose = self._pose.reshape(-1, 7) - self.ps.gpu_fetch_rigid_body_data( - data=pose, - gpu_indices=self.gpu_indices.flatten(), - data_type=RigidBodyGPUAPIReadType.POSE, - ) - pose = convert_quat(pose[:, :4], to="wxyz") - pose = pose[:, [4, 5, 6, 0, 1, 2, 3]] - return self._pose + """Local poses in the legacy Group layout ``xyz + wxyz``.""" + flat = self._pose.reshape(-1, 7) + self.body_view.fetch_pose(flat) + flat[:, 3:7] = convert_quat(flat[:, 3:7], to="wxyz") + return self._pose @property def lin_vel(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch linear velocity from CPU entities - self._lin_vel = torch.as_tensor( - [ - [entity.get_linear_velocity() for entity in instance] - for instance in self.entities - ], - dtype=torch.float32, - device=self.device, - ) - else: - lin_vel = self._lin_vel.reshape(-1, 3) - self.ps.gpu_fetch_rigid_body_data( - data=lin_vel, - gpu_indices=self.gpu_indices.flatten(), - data_type=RigidBodyGPUAPIReadType.LINEAR_VELOCITY, - ) + self.body_view.fetch_linear_velocity(self._lin_vel.reshape(-1, 3)) return self._lin_vel @property def ang_vel(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch angular velocity from CPU entities - self._ang_vel = torch.as_tensor( - [ - [entity.get_angular_velocity() for entity in instance] - for instance in self.entities - ], - dtype=torch.float32, - device=self.device, - ) - else: - ang_vel = self._ang_vel.reshape(-1, 3) - self.ps.gpu_fetch_rigid_body_data( - data=ang_vel, - gpu_indices=self.gpu_indices.flatten(), - data_type=RigidBodyGPUAPIReadType.ANGULAR_VELOCITY, - ) + self.body_view.fetch_angular_velocity(self._ang_vel.reshape(-1, 3)) return self._ang_vel @property def vel(self) -> torch.Tensor: - """Get the linear and angular velocities of the rigid bodies. - - Returns: - torch.Tensor: The linear and angular velocities concatenated, with shape (num_instances, num_objects, 6). - """ + """Linear and angular velocities with shape ``[env, object, 6]``.""" return torch.cat((self.lin_vel, self.ang_vel), dim=-1) class RigidObjectGroup(BatchEntity): - """RigidObjectGroup represents a batch of rigid bodies in the simulation.""" + """A two-dimensional view over rigid objects owned by DexSim Spawn.""" def __init__( self, cfg: RigidObjectGroupCfg, - entities: List[List[MeshObject]] = None, + entities: Sequence[Sequence[SpawnedObject]] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: self.body_type = cfg.body_type + self._declared_num_objects = len(cfg.rigid_objects) - self._world = dexsim.default_world() - self._ps = self._world.get_physics_scene() - - self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() - self._all_obj_indices = torch.arange( - len(entities[0]), dtype=torch.int32 - ).tolist() - - # data for managing body data (only for dynamic and kinematic bodies) on GPU. - self._data = RigidBodyGroupData(entities=entities, ps=self._ps, device=device) + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared RigidObjectGroup requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities: list[list[SpawnedObject]] = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._all_obj_indices = list(range(self._declared_num_objects)) + return - body_cfgs = list(cfg.rigid_objects.values()) - for instance in entities: - for i, body in enumerate(instance): - body.set_body_scale(*body_cfgs[i].body_scale) - body.set_physical_attr(body_cfgs[i].attrs.attr()) + rows = [list(instance) for instance in entities] + if not rows or any( + len(instance) != self._declared_num_objects for instance in rows + ): + raise ValueError( + "RigidObjectGroup Spawn handles must have shape " + "[num_instances, num_objects]." + ) + if spawn_result is None: + raise ValueError( + "RigidObjectGroup entities must be owned by a SpawnResult." + ) - if device.type == "cuda": - self._world.update(0.001) + self._declared_num_instances = len(rows) + self._spawn_result = spawn_result + self._all_indices = list(range(len(rows))) + self._all_obj_indices = list(range(self._declared_num_objects)) + flat_entities = [entity for instance in rows for entity in instance] + batch = spawn_result.create_rigid_body_batch(flat_entities) + body_view = SpawnRigidBodyView(spawn_result, batch, device) + self._data = RigidBodyGroupData( + body_view, + num_instances=len(rows), + num_objects=self._declared_num_objects, + device=device, + ) - super().__init__(cfg, entities, device) + super().__init__(cfg, rows, device, auto_reset=False) + self.reset() - # set default collision filter - self._set_default_collision_filter() + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for Spawn materialization.""" + return self._spawn_result is None and not self._entities - # reserve flag for collision visible node existence - n_instances = len(self._entities[0]) - self._has_collision_visible_node_list = [False] * n_instances + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to a SpawnResult.""" + return self._spawn_result is not None - def __str__(self) -> str: - parent_str = super().__str__() - return ( - parent_str - + f" | body type: {self.body_type} | num_objects: {self.num_objects}" - ) + @property + def num_instances(self) -> int: + return len(self._entities) if self._entities else self._declared_num_instances @property def num_objects(self) -> int: - """Get the number of objects in each rigid body instance. - - Returns: - int: The number of objects in each rigid body instance. - """ - return self._data.num_objects + return self._declared_num_objects @property def body_data(self) -> RigidBodyGroupData: - """Get the rigid body data manager for this rigid object. - - Returns: - RigidBodyGroupData: The rigid body data manager. - """ + if self._data is None: + raise RuntimeError( + f"RigidObjectGroup {self.uid!r} is not bound; call SimulationManager.prepare()." + ) return self._data @property def body_state(self) -> torch.Tensor: - """Get the body state of the rigid object. - - The body state of a rigid object is represented as a tensor with the following format: - [x, y, z, qw, qx, qy, qz, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z] - - If the rigid object is static, linear and angular velocities will be zero. - - Returns: - torch.Tensor: The body state of the rigid object with shape (num_instances, num_objects, 13), - where N is the number of instances. - """ + """Pose and velocity with shape ``[env, object, 13]``.""" return torch.cat( (self.body_data.pose, self.body_data.lin_vel, self.body_data.ang_vel), dim=-1, @@ -266,46 +185,91 @@ def body_state(self) -> torch.Tensor: @property def is_non_dynamic(self) -> bool: - """Check if the rigid object is non-dynamic (static or kinematic). - - Returns: - bool: True if the rigid object is non-dynamic, False otherwise. - """ return self.body_type in ("static", "kinematic") - def _set_default_collision_filter(self) -> None: - collision_filter_data = torch.zeros( - size=(self.num_instances, 4), dtype=torch.int32 - ) - for i in range(self.num_instances): - collision_filter_data[i, 0] = i - collision_filter_data[i, 1] = 1 - self.set_collision_filter(collision_filter_data) - - def set_collision_filter( - self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None + def bind_spawn( + self, + result: SpawnResult, + entities: Sequence[SpawnedObject], ) -> None: - """set collision filter data for the rigid object group. + """Bind the declaration facade to env-major Spawn handles in place.""" + if self.is_spawn_bound: + raise RuntimeError(f"RigidObjectGroup {self.uid!r} is already Spawn-bound.") + expected = self.num_instances * self.num_objects + if len(entities) != expected: + raise ValueError( + f"RigidObjectGroup {self.uid!r} expected {expected} Spawn handles, " + f"got {len(entities)}." + ) + rows = [ + entities[start : start + self.num_objects] + for start in range(0, expected, self.num_objects) + ] + bound = RigidObjectGroup( + self.cfg, + rows, + self.device, + spawn_result=result, + ) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) - Args: - filter_data (torch.Tensor): [N, 4] of int. - First element of each object is arena id. - If 2nd element is 0, the object will collision with all other objects in world. - 3rd and 4th elements are not used currently. + def __str__(self) -> str: + if self.is_declared: + return ( + f"{self.__class__}: declared {self.num_instances}x{self.num_objects} " + f"Spawn objects | uid: {self.uid} | device: {self.device}" + ) + return ( + super().__str__() + + f" | body type: {self.body_type} | num_objects: {self.num_objects}" + ) - env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. Defaults to None. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids + def _selected_indices( + self, + env_ids: Sequence[int] | torch.Tensor | None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> tuple[list[int], list[int], torch.Tensor]: + env = ( + self._all_indices + if env_ids is None + else torch.as_tensor(env_ids).reshape(-1).cpu().tolist() + ) + objects = ( + self._all_obj_indices + if obj_ids is None + else torch.as_tensor(obj_ids).reshape(-1).cpu().tolist() + ) + if any(index < 0 or index >= self.num_instances for index in env): + raise IndexError("RigidObjectGroup environment index is out of range.") + if any(index < 0 or index >= self.num_objects for index in objects): + raise IndexError("RigidObjectGroup object index is out of range.") + rows = torch.as_tensor( + [ + env_id * self.num_objects + obj_id + for env_id in env + for obj_id in objects + ], + dtype=torch.long, + device=self.device, + ) + return env, objects, rows - if len(local_env_ids) != len(filter_data): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." + def set_collision_filter( + self, + filter_data: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set one PhysX collision filter value for every member in each env.""" + env, _, _ = self._selected_indices(env_ids) + values = np.asarray(filter_data.detach().cpu(), dtype=np.uint32).reshape(-1, 4) + if len(values) != len(env): + raise ValueError( + f"Expected {len(env)} collision filters, got {len(values)}." ) - - filter_data_np = filter_data.cpu().numpy().astype(np.uint32) - for i, env_idx in enumerate(local_env_ids): - for entity in self._entities[env_idx]: - entity.get_physical_body().set_collision_filter_data(filter_data_np[i]) + for row, env_id in enumerate(env): + for entity in self._entities[env_id]: + entity.get_physical_body().set_collision_filter_data(values[row]) def set_local_pose( self, @@ -313,96 +277,43 @@ def set_local_pose( env_ids: Sequence[int] | None = None, obj_ids: Sequence[int] | None = None, ) -> None: - """Set local pose of the rigid object group. - - Args: - pose (torch.Tensor): The local pose of the rigid object group with shape (num_instances, num_objects, 7) or - (num_instances, num_objects, 4, 4). - env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. - obj_ids (Sequence[int] | None, optional): Object indices within the group. If None, all objects are set. Defaults to None. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - local_obj_ids = self._all_obj_indices if obj_ids is None else obj_ids - - if len(local_env_ids) != len(pose): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." + """Set Group poses in ``xyz+wxyz`` or homogeneous-matrix form.""" + env, objects, rows = self._selected_indices(env_ids, obj_ids) + expected_prefix = (len(env), len(objects)) + pose = pose.to(device=self.device, dtype=torch.float32) + if tuple(pose.shape) == (*expected_prefix, 7): + flat = pose.reshape(-1, 7) + target = torch.cat( + (flat[:, :3], convert_quat(flat[:, 3:7], to="xyzw")), dim=-1 ) - - if self.device.type == "cpu": - pose = pose.cpu() - if pose.dim() == 3 and pose.shape[2] == 7: - reshape_pose = pose.reshape(-1, 7) - pose_matrix = ( - torch.eye(4).unsqueeze(0).repeat(reshape_pose.shape[0], 1, 1) - ) - pose_matrix[:, :3, 3] = reshape_pose[:, :3] - pose_matrix[:, :3, :3] = matrix_from_quat(reshape_pose[:, 3:7]) - pose = pose_matrix.reshape(-1, len(local_obj_ids), 4, 4) - elif pose.dim() == 4 and pose.shape[2:] == (4, 4): - pass - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (num_instances, num_objects, 7) or (num_instances, num_objects, 4, 4)." - ) - - for i, env_idx in enumerate(local_env_ids): - for j, obj_idx in enumerate(local_obj_ids): - self._entities[env_idx][obj_idx].set_local_pose(pose[i, j]) - - else: - if pose.dim() == 3 and pose.shape[2] == 7: - xyz = pose[..., :3].reshape(-1, 3) - quat = pose[..., 3:7].reshape(-1, 4) - quat = convert_quat(quat, to="xyzw") - elif pose.dim() == 4 and pose.shape[2:] == (4, 4): - xyz = pose[..., :3, 3].reshape(-1, 3) - mat = pose[..., :3, :3].reshape(-1, 3, 3) - quat = quat_from_matrix(mat) - quat = convert_quat(quat, to="xyzw") - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - # we should keep `pose_` life cycle to the end of the function. - pose = torch.cat((quat, xyz), dim=-1) - indices = self.body_data.gpu_indices[local_env_ids][ - :, local_obj_ids - ].flatten() - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=pose.clone(), - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.POSE, + elif tuple(pose.shape) == (*expected_prefix, 4, 4): + flat = pose.reshape(-1, 4, 4) + target = torch.cat( + ( + flat[:, :3, 3], + convert_quat(quat_from_matrix(flat[:, :3, :3]), to="xyzw"), + ), + dim=-1, ) - self._world.sync_poses_gpu_to_cpu( - rigid_pose=CudaArray(pose), rigid_gpu_indices=CudaArray(indices) + else: + raise ValueError( + f"Expected pose shape {(*expected_prefix, 7)} or " + f"{(*expected_prefix, 4, 4)}, got {tuple(pose.shape)}." ) + self.body_data.body_view.apply_pose(target, rows) def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: - """Get local pose of the rigid object group. - - Args: - to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qw, qx, qy, qz). Defaults to False. - - Returns: - torch.Tensor: The local pose of the rigid object with shape (num_instances, num_objects, 7) or (num_instances, num_objects, 4, 4) depending on `to_matrix`. - """ + """Return all Group poses as ``xyz+wxyz`` or homogeneous matrices.""" pose = self.body_data.pose - if to_matrix: - pose = pose.reshape(-1, 7) - xyz = pose[:, :3] - mat = matrix_from_quat(pose[:, 3:7]) - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze(0) - .repeat(self.num_instances * self.num_objects, 1, 1) - ) - pose[:, :3, 3] = xyz - pose[:, :3, :3] = mat - pose = pose.reshape(self.num_instances, self.num_objects, 4, 4) - return pose + if not to_matrix: + return pose + flat = pose.reshape(-1, 7) + result = torch.eye(4, dtype=torch.float32, device=self.device).repeat( + len(flat), 1, 1 + ) + result[:, :3, 3] = flat[:, :3] + result[:, :3, :3] = matrix_from_quat(flat[:, 3:7]) + return result.reshape(self.num_instances, self.num_objects, 4, 4) def get_object_vertices( self, @@ -410,34 +321,19 @@ def get_object_vertices( env_ids: Sequence[int] | None = None, scale: bool = False, ) -> torch.Tensor: - """Get one constituent object's vertices across selected environments. - - Args: - object_id: Constituent object index within the group. - env_ids: Environment indices. If ``None``, returns all instances. - scale: Whether to apply each object's body scale. - - Returns: - Vertices with shape ``(N, num_vertices, 3)``. - """ - if not 0 <= object_id < self.num_objects: - raise IndexError( - f"object_id {object_id} is outside [0, {self.num_objects - 1}]." - ) - ids = self._all_indices if env_ids is None else env_ids + """Return one member's render vertices across selected environments.""" + env, objects, _ = self._selected_indices(env_ids, [object_id]) + object_id = objects[0] vertices = np.asarray( - [ - get_combined_vertices(self._entities[env_id][object_id]) - for env_id in ids - ], + [get_combined_vertices(self._entities[index][object_id]) for index in env], dtype=np.float32, ) if scale: scales = np.asarray( - [self._entities[env_id][object_id].get_body_scale() for env_id in ids], + [self._entities[index][object_id].get_body_scale() for index in env], dtype=np.float32, ) - vertices = vertices * scales[:, None, :] + vertices *= scales[:, None, :] return torch.as_tensor(vertices, dtype=torch.float32, device=self.device) def get_object_triangles( @@ -445,35 +341,17 @@ def get_object_triangles( object_id: int, env_ids: Sequence[int] | None = None, ) -> torch.Tensor: - """Get one constituent object's triangle indices. - - Args: - object_id: Constituent object index within the group. - env_ids: Environment indices. If ``None``, returns all instances. - - Returns: - Triangle indices with shape ``(N, num_triangles, 3)``. - """ - if not 0 <= object_id < self.num_objects: - raise IndexError( - f"object_id {object_id} is outside [0, {self.num_objects - 1}]." - ) - ids = self._all_indices if env_ids is None else env_ids + """Return one member's render triangles across selected environments.""" + env, objects, _ = self._selected_indices(env_ids, [object_id]) + object_id = objects[0] triangles = np.asarray( - [ - get_combined_triangles(self._entities[env_id][object_id]) - for env_id in ids - ], + [get_combined_triangles(self._entities[index][object_id]) for index in env], dtype=np.int32, ) return torch.as_tensor(triangles, dtype=torch.int32, device=self.device) def get_user_ids(self) -> torch.Tensor: - """Get the user ids of the rigid body group. - - Returns: - torch.Tensor: A tensor of shape (num_envs, num_objects) representing the user ids of the rigid body group. - """ + """Return render user ids with shape ``[env, object]``.""" return torch.as_tensor( [ [entity.get_user_id() for entity in instance] @@ -484,164 +362,78 @@ def get_user_ids(self) -> torch.Tensor: ) def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: - """Clear the dynamics of the rigid bodies by resetting velocities and applying zero forces and torques. - - Args: - env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used. - """ + """Clear velocity and one-step wrench buffers for selected envs.""" if self.is_non_dynamic: return - - local_env_ids = self._all_indices if env_ids is None else env_ids - - if self.device.type == "cpu": - for env_idx in local_env_ids: - for entity in self._entities[env_idx]: - entity.clear_dynamics() - else: - # Apply zero force and torque to the rigid bodies. - zeros = torch.zeros( - (len(local_env_ids) * self.num_objects, 3), - dtype=torch.float32, - device=self.device, - ) - indices = self.body_data.gpu_indices[local_env_ids].flatten() - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.FORCE, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.TORQUE, - ) + _, _, rows = self._selected_indices(env_ids) + zeros = torch.zeros((len(rows), 3), dtype=torch.float32, device=self.device) + view = self.body_data.body_view + view.apply_linear_velocity(zeros, rows) + view.apply_angular_velocity(zeros, rows) + view.apply_force(zeros, rows) + view.apply_torque(zeros, rows) def set_visual_material( - self, mat: VisualMaterial, env_ids: Sequence[int] | None = None + self, + mat: VisualMaterial, + env_ids: Sequence[int] | None = None, ) -> None: - """Set visual material for the rigid object group. - - Note: - For each entity in the rigid object group, a unique material instance will be created and shared - among all objects in that entity. - - Args: - mat (VisualMaterial): The material to set. - env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - - for i, env_idx in enumerate(local_env_ids): - mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}") - for j, entity in enumerate(self._entities[env_idx]): - entity.set_material(mat_inst.mat) - - # Note: The rigid object group is not supported to change the visual material once created. - # If needed, we should create a visual material dict to store the material instances, and - # implement a get_visual_material method to retrieve the material instances. + """Assign one material instance to all members in each selected env.""" + env, _, _ = self._selected_indices(env_ids) + for env_id in env: + material = mat.create_instance(f"{mat.uid}_{self.uid}_{env_id}") + for entity in self._entities[env_id]: + entity.set_material(material.mat) def reset(self, env_ids: Sequence[int] | None = None) -> None: - local_env_ids = self._all_indices if env_ids is None else env_ids - num_instances = len(local_env_ids) - - self.cfg: RigidObjectGroupCfg - body_cfgs = list(self.cfg.rigid_objects.values()) - - init_pos = [] - init_rot = [] - for cfg in body_cfgs: - init_pos.append(cfg.init_pos) - init_rot.append(cfg.init_rot) - - # (num_objects, 3) - pos = torch.as_tensor(init_pos, dtype=torch.float32, device=self.device) - rot = ( - torch.as_tensor(init_rot, dtype=torch.float32, device=self.device) - * torch.pi - / 180.0 - ) - # Convert pos and rot to shape (num_instances, num_objects, dim) - pos = pos.unsqueeze_(0).repeat(num_instances, 1, 1) - rot = rot.unsqueeze_(0).repeat(num_instances, 1, 1) - - mat = matrix_from_euler(rot.reshape(-1, 3), "XYZ") - # Init pose with shape (num_instances, num_objects, 4, 4) - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze_(0) - .repeat(num_instances * self.num_objects, 1, 1) - ) - pose[:, :3, 3] = pos.reshape(-1, 3) - pose[:, :3, :3] = mat - pose = pose.reshape(num_instances, self.num_objects, 4, 4) - self.set_local_pose(pose, env_ids=local_env_ids) - - self.clear_dynamics(env_ids=local_env_ids) + env, _, _ = self._selected_indices(env_ids) + member_poses = [] + for cfg in self.cfg.rigid_objects.values(): + if cfg.init_local_pose is not None: + member_poses.append( + torch.as_tensor( + cfg.init_local_pose, + dtype=torch.float32, + device=self.device, + ).reshape(4, 4) + ) + continue + pose = torch.eye(4, dtype=torch.float32, device=self.device) + pose[:3, 3] = torch.as_tensor( + cfg.init_pos, dtype=torch.float32, device=self.device + ) + rotation = torch.as_tensor( + cfg.init_rot, dtype=torch.float32, device=self.device + ) + pose[:3, :3] = matrix_from_euler( + (rotation * torch.pi / 180.0).reshape(1, 3), "XYZ" + )[0] + member_poses.append(pose) + pose = torch.stack(member_poses).repeat(len(env), 1, 1) + self.set_local_pose(pose.reshape(len(env), self.num_objects, 4, 4), env_ids=env) + self.clear_dynamics(env_ids=env) def set_physical_visible( self, visible: bool = True, rgba: Sequence[float] | None = None, - ): - """set collion render visibility - - Args: - visible (bool, optional): is collision body visible. Defaults to True. - rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None. - """ - rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7) - if len(rgba) != 4: - logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.") - - # create collision visible node if not exist - if visible: - for i, env_idx in enumerate(self._all_indices): - for intance_id, entity in enumerate(self._entities[env_idx]): - if not self._has_collision_visible_node_list[intance_id]: - entity.create_physical_visible_node( - np.array( - [ - rgba[0], - rgba[1], - rgba[2], - rgba[3], - ] - ) - ) - self._has_collision_visible_node_list[intance_id] = True - - # create collision visible node if not exist - for i, env_idx in enumerate(self._all_indices): - for entity in self._entities[env_idx]: - entity.set_physical_visible(visible) + ) -> None: + """Set collision-geometry visibility for every Group member.""" + color = np.asarray( + (0.8, 0.2, 0.2, 0.7) if rgba is None else rgba, + dtype=np.float32, + ) + if color.shape != (4,): + raise ValueError("Collision visualization color must contain four values.") + for instance in self._entities: + for entity in instance: + self._spawn_result.set_physical_visible(entity, color, visible) def set_visible(self, visible: bool = True) -> None: - """Set the visibility of the rigid object group. - - Args: - visible (bool, optional): Whether the rigid object group is visible. Defaults to True. - """ - for i, env_idx in enumerate(self._all_indices): - for entity in self._entities[env_idx]: + """Set render visibility for every Group member.""" + for instance in self._entities: + for entity in instance: entity.set_visible(visible) def destroy(self) -> None: - env = self._world.get_env() - arenas = env.get_all_arenas() - if len(arenas) == 0: - arenas = [env] - for i, instance in enumerate(self._entities): - for entity in instance: - arenas[i].remove_actor(entity) + """Leave topology destruction to SimulationManager and SpawnResult.""" diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 7b8a1340e..e3a50d9cc 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -19,7 +19,7 @@ import torch import numpy as np -from typing import Dict, List, Literal, Sequence, Tuple +from typing import TYPE_CHECKING, Dict, List, Literal, Sequence, Tuple from dataclasses import dataclass, field from tensordict import TensorDict @@ -39,6 +39,9 @@ ) from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedArticulation + @dataclass class ControlGroup: @@ -71,11 +74,14 @@ class Robot(Articulation): def __init__( self, cfg: RobotCfg, - entities: List[_Articulation], + entities: List[_Articulation | SpawnedArticulation] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: - self._entities = entities + self._entities = [] if entities is None else entities self.cfg = cfg # Initialize joint ids for control parts. @@ -91,12 +97,18 @@ def __init__( # cache I/O unless a task actually requests workspace sampling. self._workspaces: Dict[str, RobotWorkspace] = {} - if self.cfg.control_parts: + if entities is not None and self.cfg.control_parts: self._init_control_parts(self.cfg.control_parts) - super().__init__(cfg, entities, device) + super().__init__( + cfg, + entities, + device, + spawn_result=spawn_result, + declared_num_instances=declared_num_instances, + ) - if self.cfg.solver_cfg: + if entities is not None and self.cfg.solver_cfg: self.init_solver(self.cfg.solver_cfg) def __str__(self) -> str: diff --git a/embodichain/lab/sim/objects/soft_object.py b/embodichain/lab/sim/objects/soft_object.py index 9fbc1f2d1..8fccb56a9 100644 --- a/embodichain/lab/sim/objects/soft_object.py +++ b/embodichain/lab/sim/objects/soft_object.py @@ -19,10 +19,11 @@ import torch import dexsim import numpy as np +from copy import deepcopy from functools import cached_property from dataclasses import dataclass -from typing import List, Sequence, Union +from typing import Any, List, Sequence, TYPE_CHECKING, Union from dexsim.models import MeshObject from dexsim.engine import PhysicsScene, SoftBody @@ -47,6 +48,9 @@ ) from embodichain.utils.math import xyz_quat_to_4x4_matrix +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult + __all__ = ["SoftBodyData", "SoftObject", "SoftObjectCfg"] @@ -198,18 +202,48 @@ class SoftObject(BatchEntity): def __init__( self, cfg: SoftObjectCfg, - entities: List[MeshObject] = None, + entities: Sequence[Any] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: - self._world: dexsim.World = dexsim.default_world() - from embodichain.lab.sim.sim_manager import get_physics_scene + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared SoftObject requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._world = None + self._ps = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._visual_material = [None] * declared_num_instances + self.is_shared_visual_material = False + return + + entities = list(entities) + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result + if spawn_result is None: + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene - self._ps = get_physics_scene() + self._ps = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = self._world.get_physics_scene() self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() self._data = SoftBodyData(entities=entities, ps=self._ps, device=device) - self._world.update(0.001) + if spawn_result is None: + self._world.update(0.001) self._visual_material: List[VisualMaterialInst | None] = [None] * len(entities) self.is_shared_visual_material = False @@ -221,6 +255,44 @@ def __init__( # set default collision filter self._set_default_collision_filter() + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized SpawnResult.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its SpawnResult binding.""" + return self._spawn_result is None and len(self._entities) == 0 + + @property + def num_instances(self) -> int: + return len(self._entities) if self._entities else self._declared_num_instances + + def bind_spawn(self, result: SpawnResult, entities: Sequence[Any]) -> None: + """Bind a declared facade to finalized soft-body handles in place.""" + if len(entities) != self._declared_num_instances: + raise ValueError( + f"SoftObject {self.uid!r} expected {self._declared_num_instances} " + f"Spawn handles, got {len(entities)}." + ) + bound = SoftObject( + self.cfg, + entities, + self.device, + spawn_result=result, + ) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) + + def __str__(self) -> str: + if self.is_declared: + return ( + f"{self.__class__}: declared {self.num_instances} Spawn soft " + f"objects | uid: {self.uid} | device: {self.device}" + ) + return super().__str__() + def _initialize_existing_visual_material(self) -> None: """Wrap asset-parsed materials during soft-object construction. @@ -379,28 +451,38 @@ def set_local_pose( arena_offsets = sim.arena_offsets for i, env_idx in enumerate(local_env_ids): # TODO: soft body cannot directly set by `set_local_pose` currently. - rest_collision_vertices = self.body_data.rest_collision_vertices[i] - rest_sim_vertices = self.body_data.rest_sim_vertices[i] + soft_body: SoftBody = self._entities[env_idx].get_physical_body() + rest_collision_vertices = self.body_data.rest_collision_vertices[env_idx] + rest_sim_vertices = self.body_data.rest_sim_vertices[env_idx] + initial_transform = torch.as_tensor( + soft_body.get_initial_transform(), + dtype=torch.float32, + device=self.device, + ) + initial_rotation = initial_transform[:3, :3] + initial_translation = initial_transform[:3, 3] + rest_collision_vertices_local = ( + rest_collision_vertices - initial_translation + ) @ initial_rotation + rest_sim_vertices_local = ( + rest_sim_vertices - initial_translation + ) @ initial_rotation rotation = pose4x4[i][:3, :3] translation = pose4x4[i][:3, 3] - # apply transformation to local rest vertices and back - rest_collision_vertices_local = rest_collision_vertices - arena_offsets[i] transformed_collision_vertices = ( rest_collision_vertices_local @ rotation.T + translation ) transformed_collision_vertices = ( - transformed_collision_vertices + arena_offsets[i] + transformed_collision_vertices + arena_offsets[env_idx] ) - rest_sim_vertices_local = rest_sim_vertices - arena_offsets[i] transformed_sim_vertices = ( rest_sim_vertices_local @ rotation.T + translation ) - transformed_sim_vertices = transformed_sim_vertices + arena_offsets[i] + transformed_sim_vertices = transformed_sim_vertices + arena_offsets[env_idx] # apply vertices to soft body - soft_body: SoftBody = self._entities[env_idx].get_physical_body() collision_position_buffer = soft_body.get_position_inv_mass_buffer() sim_position_buffer = soft_body.get_sim_position_inv_mass_buffer() sim_velocity_buffer = soft_body.get_sim_velocity_buffer() @@ -528,6 +610,8 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: self.set_local_pose(pose, env_ids=local_env_ids) def destroy(self) -> None: + if self.is_spawn_bound: + return # TODO: not tested yet env = self._world.get_env() arenas = env.get_all_arenas() diff --git a/embodichain/lab/sim/physics/base.py b/embodichain/lab/sim/physics/base.py index fd6ffab78..3dab9969f 100644 --- a/embodichain/lab/sim/physics/base.py +++ b/embodichain/lab/sim/physics/base.py @@ -13,14 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Swappable physics-backend abstraction for :class:`SimulationManager`. +"""Spawn-aware physics-backend abstraction for :class:`SimulationManager`. This module defines the contract that every physics backend (DexSim default, Newton/Warp, ...) satisfies. The owning :class:`SimulationManager` holds a single :class:`PhysicsBackend` instance as ``self.physics`` and -delegates the backend-specific lifecycle, scene access, world-config -activation and capability queries to it, instead of branching on a backend -name string throughout the manager. +delegates backend-specific world configuration, compatibility scene access, +and capability queries to it. Scene topology and runtime readiness are owned +by DexSim's ``SceneBuilder`` and ``SpawnResult``. The design deliberately mirrors IsaacLab's split of an orchestrator (``SimulationContext``) from a swappable physics manager (``PhysicsManager``), @@ -92,68 +92,25 @@ def configure_world( def activate(self, sim_config: "SimulationManagerCfg") -> None: """Perform backend setup immediately after the dexsim World is created. - This is the counterpart of the backend split that used to live in - ``SimulationManager.__init__`` (default ``set_physics_config`` vs - ``get_newton_manager``). + Default configures the native PhysX globals. Newton is already + registered from ``WorldConfig.newton_cfg`` and therefore has no + additional activation work. """ - # ------------------------------------------------------------------ # - # Lifecycle - # ------------------------------------------------------------------ # - @abstractmethod - def ensure_initialized(self) -> None: - """Ensure the backend runtime is ready before a physics step. - - Called at the top of :meth:`SimulationManager.update`. For the default - backend this lazy-initializes GPU physics; for Newton it finalizes the - scene (rebuilding if the scene was mutated). Idempotent. - """ - - @abstractmethod - def invalidate(self) -> None: - """Mark the backend scene as needing re-initialization. - - Called after any scene mutation (adding/removing assets) so that the - next :meth:`ensure_initialized` rebuilds as needed. A no-op for - backends without a dirty/finalize lifecycle. - """ - - @abstractmethod - def prepare(self) -> None: - """Force the backend into a ready-to-step state. - - This unifies what the legacy code exposed as two separate operations - - "GPU physics init" on the default backend and "Newton finalize" - into a - single backend-agnostic entry point. It is idempotent: a backend that is - already ready is a no-op, and after :meth:`invalidate` the next call - re-prepares (re-initializes GPU physics / re-finalizes the Newton scene) - as needed. - - Called both lazily by :meth:`ensure_initialized` before each step and - directly by the public :meth:`SimulationManager.init_gpu_physics` and - :meth:`SimulationManager.finalize_newton_physics` entry points (both of - which delegate here). - """ - - @property - @abstractmethod - def is_initialized(self) -> bool: - """Whether the backend runtime has been initialized/finalized.""" - # ------------------------------------------------------------------ # # Scene access # ------------------------------------------------------------------ # @abstractmethod def get_scene(self): - """Return the active physics scene object (default DexSim or Newton).""" + """Return a backend compatibility scene, or raise if none exists.""" @property def newton_manager(self): - """The DexSim Newton manager, or ``None`` if not the Newton backend. + """Return ``None`` because Spawn does not use ``NewtonManager``. - Returns: - The :class:`dexsim.engine.newton_physics.NewtonManager` for the - Newton backend, otherwise ``None``. + The Newton backend overrides this property with an actionable error so + callers do not accidentally mix the removed manager ownership domain + with the World-owned Spawn backend. """ return None diff --git a/embodichain/lab/sim/physics/default.py b/embodichain/lab/sim/physics/default.py index 4cbdde6d9..1401ced33 100644 --- a/embodichain/lab/sim/physics/default.py +++ b/embodichain/lab/sim/physics/default.py @@ -22,27 +22,20 @@ import dexsim from embodichain.lab.sim.cfg import PhysicsCfg -from embodichain.utils import logger from .base import PhysicsBackend if TYPE_CHECKING: - import dexsim as _dexsim # noqa: F401 - from embodichain.lab.sim.cfg import SimulationManagerCfg __all__ = ["DefaultPhysicsBackend"] class DefaultPhysicsBackend(PhysicsBackend): - """The legacy DexSim default physics backend (GPU or CPU).""" + """DexSim's default PhysX backend (GPU or CPU).""" name = "default" - def __init__(self, manager) -> None: - super().__init__(manager) - self._is_initialized_gpu_physics = False - # -- construction / world-config activation ------------------------- # def configure_world(self, world_config, sim_config: "SimulationManagerCfg") -> None: cfg = sim_config.physics_cfg @@ -59,53 +52,10 @@ def activate(self, sim_config: "SimulationManagerCfg") -> None: dexsim.set_physics_config(**cfg.to_dexsim_args()) dexsim.set_physics_gpu_memory_config(**cfg.gpu_memory.to_dict()) - # -- lifecycle ------------------------------------------------------ # - def invalidate(self) -> None: - # The default backend has no dirty/finalize lifecycle. - pass - - @property - def is_initialized(self) -> bool: - return self._is_initialized_gpu_physics - - def prepare(self) -> None: - """Initialize GPU physics for the default backend. - - Implements the unified :meth:`PhysicsBackend.prepare` contract. For the - default backend "becoming ready to step" is initializing GPU physics; on - CPU there is nothing to initialize so this is a no-op. - """ - if not self._manager.is_use_gpu_physics: - logger.log_warning( - "The simulation device is not cuda, cannot initialize GPU physics." - ) - return - - if self._is_initialized_gpu_physics: - return - - for art in self._manager._articulations.values(): - art.reallocate_body_data() - for robot in self._manager._robots.values(): - robot.reallocate_body_data() - - # Re-establish rigid object positions after articulation resets, ensuring - # no articulation kinematics step has inadvertently corrupted the broadphase - # state for rigid bodies. - for rigid_obj in self._manager._rigid_objects.values(): - rigid_obj.reset() - - self._is_initialized_gpu_physics = True - - def ensure_initialized(self) -> None: - if self._manager.is_use_gpu_physics and not self._is_initialized_gpu_physics: - logger.log_warning( - "Using GPU physics, but not initialized yet. Forcing initialization." - ) - self.prepare() - # -- scene ---------------------------------------------------------- # def get_scene(self): + """Return PhysX's compatibility scene after Spawn is prepared.""" + self._manager.prepare() return self._manager._world.get_physics_scene() # -- capabilities --------------------------------------------------- # diff --git a/embodichain/lab/sim/physics/newton.py b/embodichain/lab/sim/physics/newton.py index 86c976396..1d0f79c98 100644 --- a/embodichain/lab/sim/physics/newton.py +++ b/embodichain/lab/sim/physics/newton.py @@ -13,26 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- -"""Newton (Warp) physics backend. - -Wraps DexSim's Newton module (``dexsim.engine.newton_physics``), which itself -runs NVIDIA Newton solvers (MuJoCo-Warp / XPBD / Featherstone / VBD / -semi-implicit) on Warp. The backend owns the lazy finalize/invalidate state -machine that rebuilds the Newton model whenever the scene is mutated. -""" +"""World-owned Newton (Warp) physics backend configuration.""" from __future__ import annotations import importlib from typing import TYPE_CHECKING -from embodichain.utils import logger - from .base import PhysicsBackend if TYPE_CHECKING: - from dexsim.engine.newton_physics import NewtonManager - from embodichain.lab.sim.cfg import SimulationManagerCfg __all__ = ["NewtonPhysicsBackend"] @@ -43,11 +33,6 @@ class NewtonPhysicsBackend(PhysicsBackend): name = "newton" - def __init__(self, manager) -> None: - super().__init__(manager) - self._newton_manager: "NewtonManager | None" = None - self._is_finalized = False - # -- construction / world-config activation ------------------------- # def configure_world(self, world_config, sim_config: "SimulationManagerCfg") -> None: importlib.import_module("dexsim.engine.newton_physics") @@ -58,102 +43,30 @@ def configure_world(self, world_config, sim_config: "SimulationManagerCfg") -> N ) def activate(self, sim_config: "SimulationManagerCfg") -> None: - from dexsim.engine.newton_physics import get_newton_manager - - self._newton_manager = get_newton_manager(self._manager._world) - - # -- lifecycle ------------------------------------------------------ # - def invalidate(self) -> None: - """Mark the Newton scene as needing re-finalization after a mutation.""" - self._is_finalized = False + del sim_config + # WorldConfig.newton_cfg registers the World-owned NewtonBackend. + # SceneBuilder.finalize() completes its model; no second manager-level + # activation or rebuild domain participates. @property - def is_initialized(self) -> bool: - return self._is_finalized - - @property - def newton_manager(self) -> "NewtonManager | None": - if self._newton_manager is None: - from dexsim.engine.newton_physics import get_newton_manager - - self._newton_manager = get_newton_manager(self._manager._world) - return self._newton_manager - - def _lifecycle_state(self) -> str: - """Return the Newton manager lifecycle state name, or empty string.""" - mgr = self.newton_manager - return getattr(getattr(mgr, "lifecycle_state", None), "name", "") - - def _reset_entities_after_finalize(self) -> None: - """Apply deferred initial resets once Newton runtime data is ready.""" - for rigid_obj in self._manager._rigid_objects.values(): - rigid_obj.reset() - for articulation in self._manager._articulations.values(): - articulation.reset() - for robot in self._manager._robots.values(): - robot.reset() - # Rigid object groups are not supported on the Newton backend yet. - - def prepare(self) -> None: - """Finalize the Newton scene if it has not been finalized yet. - - Implements the unified :meth:`PhysicsBackend.prepare` contract: this is - both the "finalize" entry point (public - :meth:`SimulationManager.finalize_newton_physics`) and the "GPU init" - entry point (:meth:`SimulationManager.init_gpu_physics`) for the Newton - backend, since Newton's notion of becoming ready to step is finalizing - the model. - """ - if self._is_finalized and self._lifecycle_state() == "READY": - return - - mgr = self.newton_manager - state = self._lifecycle_state() - - if state != "READY": - from dexsim.engine.newton_physics.rebuild import ( - ensure_simulation_prepared_lazy, - rebuild_newton_from_scene, - ) - - safe_to_continue, _ = ensure_simulation_prepared_lazy( - mgr, - self._manager._world, - rebuild_from_scene=rebuild_newton_from_scene, - warn=True, - ) - if not safe_to_continue: - logger.log_error( - "Failed to finalize Newton physics: model is not ready to build " - f"(lifecycle state {state!r})." - ) - return - - state = self._lifecycle_state() - if state != "READY": - logger.log_error( - "Failed to finalize Newton physics: lifecycle state is " - f"{state!r} after simulation preparation." - ) - - self._is_finalized = True - self._reset_entities_after_finalize() - - def ensure_initialized(self) -> None: - self.prepare() + def newton_manager(self): + """Reject access to the removed, independently owned Newton manager.""" + raise RuntimeError( + "NewtonManager is not part of Spawn scene ownership. Use " + "SimulationManager.spawn_result and its Spawned*/Batch APIs." + ) # -- scene ---------------------------------------------------------- # def get_scene(self): - return self.newton_manager.scene + raise RuntimeError( + "Newton Spawn scenes do not expose a PhysicsScene. Use " + "SimulationManager.spawn_result and its Spawned*/Batch APIs." + ) # -- capabilities --------------------------------------------------- # @property def supports_robot(self) -> bool: - # Robots are URDF articulations; the Newton ``load_urdf`` patch builds a - # NewtonArticulation, and the shared spawn path (add_robot invalidate + - # _reset_entities_after_finalize) handles the Newton lifecycle. Requires - # the dexsim fix to ``NewtonArticulation._joint_metas_from_ids`` so that - # explicit joint_ids use active-joint indexing (matching get_dof()). + # Robots are SpawnedArticulations in the World-owned Newton model. return True @property diff --git a/embodichain/lab/sim/sensors/base_sensor.py b/embodichain/lab/sim/sensors/base_sensor.py index 3fb932f0d..b364c2e09 100644 --- a/embodichain/lab/sim/sensors/base_sensor.py +++ b/embodichain/lab/sim/sensors/base_sensor.py @@ -171,10 +171,18 @@ class BaseSensor(BatchEntity): SUPPORTED_DATA_TYPES = [] def __init__( - self, config: SensorCfg, device: torch.device = torch.device("cpu") + self, + config: SensorCfg, + device: torch.device = torch.device("cpu"), + *, + num_instances: int | None = None, ) -> None: - - num_envs = get_dexsim_arena_num() + num_envs = ( + get_dexsim_arena_num() if num_instances is None else int(num_instances) + ) + if num_envs <= 0: + raise ValueError("A sensor requires at least one simulation instance.") + self._num_instances = num_envs self._data_buffer: TensorDict[str, torch.Tensor] = TensorDict( {}, batch_size=[num_envs], device=device ) @@ -186,7 +194,7 @@ def __init__( @cached_property def num_instances(self) -> int: - return get_dexsim_arena_num() + return self._num_instances @abstractmethod def _build_sensor_from_config( diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py index cc9a7aa44..b118bbeee 100644 --- a/embodichain/lab/sim/sensors/camera.py +++ b/embodichain/lab/sim/sensors/camera.py @@ -21,7 +21,7 @@ import dexsim.render as dr from functools import cached_property -from typing import List, Literal, Sequence, Tuple +from typing import Callable, List, Literal, Sequence, Tuple from embodichain.lab.sim.sensors import BaseSensor, SensorCfg from embodichain.utils.math import matrix_from_quat, quat_from_matrix, look_at_to_pose @@ -134,27 +134,41 @@ class Camera(BaseSensor): SUPPORTED_DATA_TYPES = ["color", "depth", "mask", "normal", "position"] def __init__( - self, config: CameraCfg, device: torch.device = torch.device("cpu") + self, + config: CameraCfg, + device: torch.device = torch.device("cpu"), + *, + world: dexsim.World | None = None, + arenas: Sequence[dexsim.environment.Arena] | None = None, + parent_node_resolver: Callable[[str], Sequence[object]] | None = None, + defer_parent_attachment: bool = False, ) -> None: - super().__init__(config, device) + if world is None or arenas is None: + raise ValueError( + "Camera render resources must be supplied explicitly; construct " + "cameras through SimulationManager.add_sensor()." + ) + self._world = world + self._arenas = list(arenas) + if len(self._arenas) == 0: + raise ValueError("Camera requires at least one materialized Arena.") + self._parent_node_resolver = parent_node_resolver + self._camera_names: list[tuple[dexsim.environment.Arena, str]] = [] + self._is_destroyed = False + super().__init__(config, device, num_instances=len(self._arenas)) + if config.extrinsics.parent is not None and not defer_parent_attachment: + self.attach_to_parent() def _build_sensor_from_config( self, config: CameraCfg, device: torch.device ) -> None: - self._world = dexsim.default_world() - env = self._world.get_env() - arenas = env.get_all_arenas() - if len(arenas) == 0: - arenas = [env] - num_instances = len(arenas) - self._frame_buffer = self._world.create_camera_group( - [config.width, config.height], num_instances, True + [config.width, config.height], self.num_instances, True ) view_attrib = config.get_view_attrib() - for i, arena in enumerate(arenas): - view_name = f"{self.uid}_view{i + 1}" + for i, arena in enumerate(self._arenas): + view_name = f"{config.uid}_view{i + 1}" view = arena.create_camera( view_name, config.width, @@ -167,6 +181,7 @@ def _build_sensor_from_config( view.set_near(config.near) view.set_far(config.far) self._entities[i] = view + self._camera_names.append((arena, view_name)) # Define a mapping of data types to their respective shapes and dtypes buffer_specs = { @@ -202,8 +217,6 @@ def _build_sensor_from_config( ) self.cfg: CameraCfg = config - if self.cfg.extrinsics.parent is not None: - self._attach_to_entity() @cached_property def group_id(self) -> int: @@ -270,21 +283,30 @@ def update(self, **kwargs) -> None: def _attach_to_entity(self) -> None: """Attach the sensor to the parent entity in each environment.""" - env = self._world.get_env() - for i, entity in enumerate(self._entities): - - parent = None - if i == 0: - parent = env.find_node(f"{self.cfg.extrinsics.parent}") - else: - parent = env.find_node(f"{self.cfg.extrinsics.parent}.{i-1}") - if parent is None: - logger.log_error( - f"Failed to find parent entity {self.cfg.extrinsics.parent} for sensor {self.cfg.uid}." - ) - + if self._parent_node_resolver is None: + raise RuntimeError( + f"Camera {self.cfg.uid!r} has parent " + f"{self.cfg.extrinsics.parent!r}, but no Spawn parent resolver " + "was supplied." + ) + parents = list(self._parent_node_resolver(self.cfg.extrinsics.parent)) + if len(parents) != self.num_instances: + raise RuntimeError( + f"Camera parent resolver returned {len(parents)} nodes for " + f"{self.num_instances} camera instances." + ) + for entity, parent in zip(self._entities, parents): entity.attach_node(parent) + def attach_to_parent(self) -> None: + """Resolve and attach a deferred parent after Spawn materialization.""" + if self.cfg.extrinsics.parent is None: + return + self._attach_to_entity() + # Extrinsics are expressed in the parent frame. Reapply them after + # reparenting because the camera was initially reset in Arena space. + self.reset() + def set_local_pose( self, pose: torch.Tensor, env_ids: Sequence[int] | None = None ) -> None: @@ -346,14 +368,10 @@ def get_arena_pose(self, to_matrix: bool = False) -> torch.Tensor: Returns: A tensor representing the pose of the sensor in the arena frame. """ - from embodichain.lab.sim.utility import get_dexsim_arenas - - arenas = get_dexsim_arenas() - poses = [] for i, entity in enumerate(self._entities): pose = entity.get_world_pose() - pose[:2, 3] -= arenas[i].get_root_node().get_local_pose()[:2, 3] + pose[:2, 3] -= self._arenas[i].get_root_node().get_local_pose()[:2, 3] poses.append(torch.as_tensor(pose, dtype=torch.float32)) poses = torch.stack(poses, dim=0).to(self.device) @@ -363,6 +381,28 @@ def get_arena_pose(self, to_matrix: bool = False) -> torch.Tensor: return torch.cat((xyz, quat), dim=-1) return poses + def destroy(self) -> None: + """Remove render cameras before releasing their World-owned group.""" + if self._is_destroyed: + return + self._is_destroyed = True + for arena, camera_name in self._camera_names: + try: + arena.remove_camera(camera_name) + except Exception as error: + logger.log_warning( + f"Failed to remove camera {camera_name!r}: {error!r}" + ) + self._entities = [] + self._camera_names = [] + # DexSim currently has no public remove_camera_group API. The group is + # World-owned; dropping this borrowed facade after removing all views + # is the narrowest safe lifetime boundary available to EmbodiChain. + self._frame_buffer = None + self._parent_node_resolver = None + self._arenas = [] + self._world = None + def look_at( self, eye: torch.Tensor, diff --git a/embodichain/lab/sim/sensors/stereo.py b/embodichain/lab/sim/sensors/stereo.py index 999bedca9..2df992e0d 100644 --- a/embodichain/lab/sim/sensors/stereo.py +++ b/embodichain/lab/sim/sensors/stereo.py @@ -21,7 +21,7 @@ import numpy as np import dexsim.render as dr -from typing import Dict, Tuple, List, Sequence +from typing import Callable, Dict, Tuple, List, Sequence from dexsim.utility import inv_transform from embodichain.lab.sim.sensors import Camera, CameraCfg @@ -155,8 +155,20 @@ def __init__( self, config: StereoCameraCfg, device: torch.device = torch.device("cpu"), + *, + world: dexsim.World | None = None, + arenas: Sequence[dexsim.environment.Arena] | None = None, + parent_node_resolver: Callable[[str], Sequence[object]] | None = None, + defer_parent_attachment: bool = False, ) -> None: - super().__init__(config, device) + super().__init__( + config, + device, + world=world, + arenas=arenas, + parent_node_resolver=parent_node_resolver, + defer_parent_attachment=defer_parent_attachment, + ) # check valid config if self.cfg.enable_disparity and not self.cfg.enable_depth: @@ -165,21 +177,14 @@ def __init__( def _build_sensor_from_config( self, config: StereoCameraCfg, device: torch.device ) -> None: - self._world = dexsim.default_world() - env = self._world.get_env() - arenas = env.get_all_arenas() - if len(arenas) == 0: - arenas = [env] - num_instances = len(arenas) - self._frame_buffer = self._world.create_camera_group( - [config.width, config.height], num_instances * 2, True + [config.width, config.height], self.num_instances * 2, True ) view_attrib = config.get_view_attrib() left_list = [] right_list = [] - for i, arena in enumerate(arenas): - left_view_name = f"{self.uid}_left_view{i + 1}" + for i, arena in enumerate(self._arenas): + left_view_name = f"{config.uid}_left_view{i + 1}" left_view = arena.create_camera( left_view_name, config.width, @@ -192,9 +197,10 @@ def _build_sensor_from_config( left_view.set_near(config.near) left_view.set_far(config.far) left_list.append(left_view) + self._camera_names.append((arena, left_view_name)) - for i, arena in enumerate(arenas): - right_view_name = f"{self.uid}_right_view{i + 1}" + for i, arena in enumerate(self._arenas): + right_view_name = f"{config.uid}_right_view{i + 1}" right_view = arena.create_camera( right_view_name, config.width, @@ -207,8 +213,9 @@ def _build_sensor_from_config( right_view.set_near(config.near) right_view.set_far(config.far) right_list.append(right_view) + self._camera_names.append((arena, right_view_name)) - for i in range(num_instances): + for i in range(self.num_instances): self._entities[i] = PairCameraView( left_list[i], right_list[i], config.left_to_right.cpu().numpy() ) @@ -277,8 +284,6 @@ def _build_sensor_from_config( ][:, :, config.width :, :] self.cfg: CameraCfg = config - if self.cfg.extrinsics.parent is not None: - self._attach_to_entity() def update(self, **kwargs) -> None: """Update the sensor data. @@ -343,14 +348,10 @@ def get_left_right_arena_pose(self) -> torch.Tensor: Returns: torch.Tensor: The local pose of the left camera with shape (num_envs, 4, 4). """ - from embodichain.lab.sim.utility import get_dexsim_arenas - - arenas = get_dexsim_arenas() - left_poses = [] right_poses = [] for i, entity in enumerate(self._entities): - arena_pose = arenas[i].get_root_node().get_local_pose() + arena_pose = self._arenas[i].get_root_node().get_local_pose() left_pose = entity._left_view.get_world_pose() left_pose[:2, 3] -= arena_pose[:2, 3] left_poses.append( diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 11e58ce19..944836182 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -27,7 +27,6 @@ import numpy as np import warp as wp -from tqdm import tqdm from pathlib import Path from copy import deepcopy from datetime import datetime @@ -41,19 +40,22 @@ CONVEX_DECOMP_DIR = SIM_CACHE_DIR / "convex_decomposition" REACHABLE_XPOS_DIR = SIM_CACHE_DIR / "robot_reachable_xpos" + +def _is_usd_path(path: object | None) -> bool: + """Return whether a source path is a USD stage.""" + return path is not None and str(path).lower().endswith((".usd", ".usda", ".usdc")) + + from dexsim.types import ( + ActorType, Backend, ThreadMode, - PhysicalAttr, - ActorType, - RigidBodyShape, ) from dexsim.core import TASK_RETURN -from dexsim.engine import Material, PhysicsScene +from dexsim.engine import Material from dexsim.models import MeshObject -from dexsim.render import Light as _Light, LightType, Windows +from dexsim.render import LightType, Windows from dexsim.engine import GizmoController, ObjectManipulator -from dexsim.engine.newton_physics import NewtonManager, NewtonPhysicsScene from embodichain.lab.sim.objects import ( RigidObject, @@ -93,6 +95,17 @@ RigidConstraintCfg, ) from embodichain.lab.sim.physics import make_physics_backend +from embodichain.lab.sim.spawn.descriptors import ( + articulation_desc_from_cfg, + cloth_desc_from_cfg, + rigid_desc_from_cfg, + soft_desc_from_cfg, +) +from embodichain.lab.sim.spawn.usd import ( + articulation_desc_from_usd, + rigid_desc_from_usd, +) +from embodichain.lab.sim.spawn.scene import SpawnScene from embodichain.lab.sim import VisualMaterial, VisualMaterialCfg from embodichain.lab.sim.profiler import Profiler, ProfilerCfg from embodichain.lab.visualization.cfg import VisualizationCfg @@ -100,6 +113,9 @@ from embodichain.utils.math import look_at_to_pose, matrix_from_quat, pose_inv if TYPE_CHECKING: + from dexsim.engine import PhysicsScene + from dexsim.spawn import SpawnResult + from embodichain.lab.visualization import ( RuntimeHealth, RuntimeStats, @@ -176,7 +192,6 @@ def __init__( self.window_camera_pose = ( WindowCameraPoseCfg() if window_camera_pose is None else window_camera_pose ) - if physics_dt is not None: self.physics_cfg.physics_dt = physics_dt runtime_device = device if device is not None else sim_device @@ -473,7 +488,15 @@ def __init__( self._robots: Dict[str, Robot] = dict() self._sensors: Dict[str, BaseSensor] = dict() - self._lights: Dict[str, _Light] = dict() + self._pending_sensor_attachments: list[Camera] = [] + self._lights: Dict[str, Light] = dict() + + self._spawn_scene = SpawnScene( + self._world, + num_envs=sim_config.num_envs, + spacing=(sim_config.arena_space, sim_config.arena_space, 0.0), + ) + self._arenas = list(self._spawn_scene.builder.prepare_arenas()) self._visualization_runtime = None self._visualization_overlays: SceneOverlays | None = None @@ -492,15 +515,15 @@ def __init__( self._init_sim_resources() - self._create_default_plane() + # The render material is authored on the descriptor before Spawn + # materialization. The plane handle does not exist until prepare. + self._spawn_default_plane_visibility = True self.set_default_background() + self._declare_spawn_default_plane() self.set_default_global_lighting() # Set physics to manual update mode by default. self.set_manual_update(True) - self._build_multiple_arenas(sim_config.num_envs) - self.start_visualization() - if sim_config.headless is False: self._window = self._world.get_windows() @@ -598,7 +621,15 @@ def num_envs(self) -> int: Returns: int: number of arenas. """ - return len(self._arenas) if len(self._arenas) > 0 else 1 + return self.sim_config.num_envs + + @property + def spawn_result(self) -> "SpawnResult | None": + """Return the current SpawnResult, or ``None`` before first prepare.""" + spawn_scene = getattr(self, "_spawn_scene", None) + if spawn_scene is None: + return None + return spawn_scene.result @property def is_use_gpu_physics(self) -> bool: @@ -621,8 +652,13 @@ def is_newton_backend(self) -> bool: return self.physics.name == "newton" @property - def newton_manager(self) -> NewtonManager: - """Return the DexSim Newton manager for this world, if active.""" + def newton_manager(self): + """Compatibility accessor for the removed NewtonManager API. + + A non-Newton backend still returns ``None``. The Newton backend raises + an actionable error because Spawn owns its World-level runtime and no + independent NewtonManager exists. + """ if not self.is_newton_backend: logger.log_warning("Newton backend is not active.") return None @@ -712,6 +748,8 @@ def start_visualization(self) -> VisualizationRuntime | None: """Start the configured live visualizer and publish the current scene.""" if self.sim_config.visualization.backend == "none": return None + if getattr(self, "_spawn_scene", None) is not None: + self.prepare() if getattr(self, "is_window_opened", False): raise RuntimeError( "Cannot start the Viser backend while the native DexSim window " @@ -897,14 +935,21 @@ def _init_sim_resources(self) -> None: self._default_resources = SimResources() - def _invalidate_newton_physics(self) -> None: - """Mark the active backend scene as needing re-initialization. - - Delegates to the active :class:`PhysicsBackend`; a no-op for backends - without a dirty/finalize lifecycle. Called after every scene mutation - (adding assets, creating the default plane). - """ - self.physics.invalidate() + def prepare(self) -> None: + """Materialize physical declarations, then resolve sensor parents.""" + scene = self._spawn_scene + result = scene.result + if result is None or result.needs_rebuild or scene.builder.has_pending_changes: + result = scene.commit() + result.prepare_runtime() + self._env = result.get_arena("default") + self._arenas = [result.get_arena(name) for name in scene.arena_names] + self.__dict__.pop("arena_offsets", None) + scene.bind() + + for sensor in self._pending_sensor_attachments: + sensor.attach_to_parent() + self._pending_sensor_attachments.clear() def enable_physics(self, enable: bool) -> None: """Enable or disable physics simulation. @@ -932,23 +977,20 @@ def set_manual_update(self, enable: bool) -> None: self._world.set_manual_update(enable) def init_gpu_physics(self) -> None: - """Initialize the GPU physics simulation. + """Prepare the Spawn-owned physics runtime. - Delegates to the active backend's unified :meth:`PhysicsBackend.prepare` - (for the default backend this performs the real GPU initialization; for - the Newton backend it finalizes the scene). + This backwards-compatible alias now has the same backend-neutral + behavior as :meth:`prepare`. """ - self.physics.prepare() + self.prepare() def finalize_newton_physics(self) -> None: - """Finalize the Newton scene if it has not been finalized yet. + """Prepare the Spawn-owned physics runtime. - Delegates to the active backend's unified :meth:`PhysicsBackend.prepare` - (for the Newton backend this (re-)finalizes the scene and applies - deferred entity resets; for the default backend it initializes GPU - physics). + This backwards-compatible alias now has the same backend-neutral + behavior as :meth:`prepare`. """ - self.physics.prepare() + self.prepare() def create_differentiable_stepper(self): """Create a single-step differentiable physics primitive (Newton-only). @@ -1019,19 +1061,7 @@ def update(self, physics_dt: float | None = None, step: int = 1) -> None: """ with self.profiler.section("sim_update", is_root=True): with self.profiler.section("gpu_physics_check"): - if hasattr(self, "physics"): - # Lazy GPU initialization for the default backend and scene - # finalization for the Newton backend share one contract. - self.physics.ensure_initialized() - elif self.is_use_gpu_physics and not self._is_initialized_gpu_physics: - # Compatibility for lightweight manager probes that bypass - # ``SimulationManager.__init__``. - logger.log_warning( - "Using GPU physics, but not initialized yet. " - "Forcing initialization." - ) - with self.profiler.section("gpu_physics_init"): - self.init_gpu_physics() + self.prepare() if self.is_physics_manually_update: with self.profiler.section("manual_update"): @@ -1087,8 +1117,12 @@ def get_env(self, arena_index: int = -1) -> dexsim.environment.Arena: def get_world(self) -> dexsim.World: return self._world - def get_physics_scene(self) -> PhysicsScene | NewtonPhysicsScene: - """Get the physics scene of the simulation.""" + def get_physics_scene(self) -> "PhysicsScene": + """Return PhysX's compatibility scene after Spawn preparation. + + Newton has no ``PhysicsScene`` facade and raises with guidance to use + :attr:`spawn_result` instead. + """ return self.physics.get_scene() def can_open_native_window(self) -> bool: @@ -1149,32 +1183,6 @@ def close_window(self) -> None: self._window_camera_pose_input_control = None self.is_window_opened = False - def _build_multiple_arenas(self, num: int, space: float | None = None) -> None: - """Build multiple arenas in a grid pattern. - - This interface is used for vectorized simulation. - - Args: - num (int): number of arenas to build. - space (float | None, optional): The distance between each arena. Defaults to the arena_space in sim_config. - """ - - if space is None: - space = self.sim_config.arena_space - - if num <= 0: - logger.log_warning("Number of arenas must be greater than 0.") - return - - scene_grid_length = int(np.ceil(np.sqrt(num))) - - for i in range(num): - arena = self._env.add_arena(f"arena_{i}") - - id_x, id_y = i % scene_grid_length, i // scene_grid_length - arena.set_root_node_position([id_x * space, id_y * space, 0]) - self._arenas.append(arena) - def set_indirect_lighting(self, name: str) -> None: """Set indirect lighting. @@ -1204,16 +1212,61 @@ def set_emission_light( if intensity is not None: self._env.set_env_light_intensity(intensity) - def _create_default_plane(self): - default_length = 1000 - repeat_uv_size = int(default_length / 2) - self._default_plane = self._env.create_plane( - 0, default_length, repeat_uv_size, repeat_uv_size + def _declare_spawn_default_plane(self) -> None: + """Declare the global ground in the World's Spawn scene.""" + + from dexsim.spawn import ( + CollisionApproximation, + CollisionDesc, + DexsimCollisionDesc, + GeometryDesc, + NewtonCollisionDesc, + ObjectDesc, + RenderDesc, + RigidBodyPhysicsDesc, + ) + + default_length = 1000.0 + geometry = GeometryDesc.plane(default_length) + collision = CollisionDesc.from_geometry( + geometry, + approximation=CollisionApproximation.NONE, + ) + collision.dexsim = DexsimCollisionDesc( + dynamic_friction=0.5, + static_friction=0.5, + ) + collision.newton = NewtonCollisionDesc(mu=0.5) + collision.render_source_index = 0 + descriptor = ObjectDesc( + name="default_plane", + renders=[ + RenderDesc.from_geometry( + geometry, + material=self._spawn_default_plane_material, + ) + ], + collisions=[collision], + physics=RigidBodyPhysicsDesc.static(), + per_env=False, + ) + + def bind_default_plane(_result, handles) -> None: + self._default_plane = handles[0] + self._default_plane.get_render_body().repeat_uv( + np.asarray( + [default_length / 2.0, default_length / 2.0], + dtype=np.float32, + ) + ) + self._default_plane.set_visible(self._spawn_default_plane_visibility) + + self._spawn_scene.declare( + "rigid_object", + "default_plane", + descriptor, + on_bind=bind_default_plane, ) - self._default_plane.set_name("default_plane") - attr = PhysicalAttr(dynamic_friction=0.5, static_friction=0.5) - self._default_plane.add_rigidbody(ActorType.STATIC, RigidBodyShape.PLANE, attr) - self._invalidate_newton_physics() def set_default_global_lighting(self) -> None: """Set default global lighting for the scene. @@ -1230,7 +1283,6 @@ def set_default_background(self) -> None: """Set default background.""" mat_name = "plane_mat" - mat = None mat_path = self._default_resources.get_material_path("PlaneDark") color_texture = os.path.join(mat_path, "PlaneDark_2K_Color.jpg") roughness_texture = os.path.join(mat_path, "PlaneDark_2K_Roughness.jpg") @@ -1243,7 +1295,11 @@ def set_default_background(self) -> None: ) ) - self._default_plane.set_material(mat.get_instance("plane_mat").mat) + material = mat.get_instance("plane_mat").mat + # Consumed by _declare_spawn_default_plane(). Keeping the native + # material in the descriptor preserves the VisualMaterial registry + # used by visual randomization without forcing finalization. + self._spawn_default_plane_material = material self._visual_materials[mat_name] = mat def set_ground_plane_visibility(self, visible: bool) -> None: @@ -1252,10 +1308,10 @@ def set_ground_plane_visibility(self, visible: bool) -> None: Args: visible (bool): _description_ """ - if visible: - self._default_plane.set_visible(True) - else: - self._default_plane.set_visible(False) + self._spawn_default_plane_visibility = bool(visible) + if not hasattr(self, "_default_plane"): + return + self._default_plane.set_visible(bool(visible)) def set_texture_cache( self, key: str, texture: Union[torch.Tensor, List[torch.Tensor]] @@ -1289,7 +1345,15 @@ def get_texture_cache( def get_asset( self, uid: str - ) -> Light | BaseSensor | Robot | RigidObject | Articulation | None: + ) -> ( + Light + | BaseSensor + | Robot + | RigidObject + | RigidObjectGroup + | Articulation + | None + ): """Get an asset by its UID. The asset can be a light, sensor, robot, rigid object or articulation. @@ -1320,7 +1384,6 @@ def get_asset( logger.log_warning(f"Asset {uid} not found.") return None - # Light type string → dexsim LightType enum mapping _LIGHT_TYPE_MAP: dict[str, LightType] = { "point": LightType.POINT, "sun": LightType.SUN, @@ -1329,8 +1392,6 @@ def get_asset( "rect": LightType.RECT, "mesh": LightType.MESH, } - - # Light types that are created as a single global scene light (not per-environment). _GLOBAL_LIGHT_TYPES: tuple[str, ...] = ("sun", "direction") def add_light(self, cfg: LightCfg) -> Light: @@ -1355,7 +1416,7 @@ def add_light(self, cfg: LightCfg) -> Light: Light: The created light instance. Raises: - RuntimeError: If ``cfg.light_type`` is not one of the supported types. + ValueError: If ``cfg.light_type`` is not supported. """ if cfg.uid is None: uid = "light" @@ -1366,45 +1427,41 @@ def add_light(self, cfg: LightCfg) -> Light: if uid in self._lights: logger.log_error(f"Light {uid} already exists.") - light_type_str = cfg.light_type - light_type = self._LIGHT_TYPE_MAP.get(light_type_str) + light_type = self._LIGHT_TYPE_MAP.get(cfg.light_type) if light_type is None: - supported = ", ".join(self._LIGHT_TYPE_MAP.keys()) - logger.log_error( - f"Unsupported light type: '{light_type_str}'. " + supported = ", ".join(self._LIGHT_TYPE_MAP) + raise ValueError( + f"Unsupported light type {cfg.light_type!r}. " f"Supported types: {supported}." ) - # Validation warnings for type-specific constraints - if light_type_str == "mesh" and not cfg.mesh_path: + if cfg.light_type == "mesh" and not cfg.mesh_path: logger.log_warning( f"Mesh light '{uid}' has no mesh_path set. " f"Use set_mesh() to assign a MeshObject." ) - if light_type_str == "rect" and (cfg.rect_width <= 0 or cfg.rect_height <= 0): + if cfg.light_type == "rect" and (cfg.rect_width <= 0 or cfg.rect_height <= 0): logger.log_warning( f"Rect light '{uid}' has zero or negative dimensions " f"(width={cfg.rect_width}, height={cfg.rect_height})." ) if cfg.light_type in self._GLOBAL_LIGHT_TYPES: - # Global scene light: create a single instance on the root - # environment. Infinite-distance lights (sun, direction) are - # physically scene-global and should not be duplicated per arena. - light = self._env.create_light(uid, light_type) - batch_lights = Light(cfg=cfg, entities=[light]) + batch_lights = Light( + cfg=cfg, + entities=[self._env.create_light(uid, light_type)], + ) else: - # Per-environment batched light: one instance per arena. - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - light_list = [] - for i, env in enumerate(env_list): - light_name = f"{uid}_{i}" - light = env.create_light(light_name, light_type) - light_list.append(light) - batch_lights = Light(cfg=cfg, entities=light_list) + batch_lights = Light( + cfg=cfg, + entities=[ + arena.create_light(f"{uid}_{index}", light_type) + for index, arena in enumerate(self._arenas) + ], + ) self._lights[uid] = batch_lights - + self.notify_visualization_topology_changed() return batch_lights def get_light(self, uid: str) -> Light | None: @@ -1429,6 +1486,141 @@ def get_light_uid_list(self) -> List[str]: """ return list(self._lights.keys()) + def add_usd( + self, + name: str, + file_path: str, + *, + pose: np.ndarray | None = None, + robot_cfgs: dict[str, RobotCfg] | None = None, + ) -> dict[str, RigidObject | Articulation | Robot]: + """Declare the supported entities in a USD scene. + + The returned facades are keyed by their USD prim paths. They remain in + declared state until :meth:`prepare` finalizes the shared Spawn scene, + then bind in place to the resulting DexSim handles. + + USD does not identify which articulations should expose EmbodiChain's + robot interface. Pass those explicitly through ``robot_cfgs``; all + other articulation descriptions become :class:`Articulation` objects. + + Args: + name: Name passed to DexSim's USD scene parser. + file_path: USD, USDA, or USDC file path. + pose: Optional scene-root transform. + robot_cfgs: Robot configurations keyed by USD prim path. These + provide robot-side metadata while physics remains authored by + the USD scene. + + Returns: + Supported EmbodiChain facades keyed by USD prim path. + + Raises: + RuntimeError: If called after the Spawn scene was finalized. + """ + if self.spawn_result is not None: + raise RuntimeError( + "add_usd() must be called before SimulationManager.prepare()." + ) + + from dexsim.spawn import ArticulationDesc, MeshObjectDesc + + descriptors = self._spawn_scene.builder.add_usd( + name, + file_path, + pose=pose, + per_env=True, + ) + assets: dict[str, RigidObject | Articulation | Robot] = {} + robot_cfgs = robot_cfgs or {} + + for descriptor in descriptors: + source_path = ( + descriptor.usd.prim_path + if descriptor.usd is not None and descriptor.usd.prim_path + else descriptor.name + ) + + if type(descriptor) is MeshObjectDesc: + body_type = "static" + if descriptor.physics is not None: + body_type = { + ActorType.DYNAMIC: "dynamic", + ActorType.KINEMATIC: "kinematic", + ActorType.STATIC: "static", + }[descriptor.physics.actor_type] + cfg = RigidObjectCfg( + uid=descriptor.name, + init_local_pose=descriptor.pose.copy(), + body_type=body_type, + body_scale=tuple(float(value) for value in descriptor.body_scale), + use_usd_properties=True, + ) + facade = RigidObject( + cfg=cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) + + def bind_rigid(result, handles, target=facade) -> None: + if target.is_declared: + target.bind_spawn(result, handles) + + self._spawn_scene.track( + "rigid_object", + descriptor.name, + descriptor, + on_bind=bind_rigid, + ) + self._rigid_objects[descriptor.name] = facade + assets[source_path] = facade + continue + + if isinstance(descriptor, ArticulationDesc): + robot_cfg = robot_cfgs.get(source_path) + facade_type: type[Articulation] = ( + Robot if robot_cfg is not None else Articulation + ) + cfg = ( + deepcopy(robot_cfg) + if robot_cfg is not None + else ArticulationCfg(uid=descriptor.name) + ) + cfg.uid = descriptor.name + cfg.fpath = file_path + cfg.init_local_pose = descriptor.pose.copy() + cfg.use_usd_properties = True + cfg.fix_base = bool(descriptor.fixed_base) + cfg.disable_self_collision = not descriptor.enable_self_collision + cfg.body_scale = tuple(float(value) for value in descriptor.body_scale) + cfg.build_pk_chain = False + facade = facade_type( + cfg=cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) + + def bind_articulation(result, handles, target=facade) -> None: + if target.is_declared: + target.bind_spawn(result, handles) + + self._spawn_scene.track( + "articulation", + descriptor.name, + descriptor, + on_bind=bind_articulation, + ) + registry = ( + self._robots if robot_cfg is not None else self._articulations + ) + registry[descriptor.name] = facade + assets[source_path] = facade + + self.notify_visualization_topology_changed() + return assets + def add_rigid_object( self, cfg: RigidObjectCfg, @@ -1441,37 +1633,44 @@ def add_rigid_object( Returns: RigidObject: The added rigid object instance handle. """ - from embodichain.lab.sim.utility.sim_utils import ( - load_mesh_objects_from_cfg, - ) - uid = cfg.uid if uid is None: - logger.log_error("Rigid object uid must be specified.") + raise ValueError("Rigid object uid must be specified.") if uid in self._rigid_objects: - logger.log_error(f"Rigid object {uid} already exists.") - - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = load_mesh_objects_from_cfg( - cfg=cfg, - env_list=env_list, - cache_dir=self._convex_decomp_dir, - ) + raise ValueError(f"Rigid object {uid!r} already exists.") + source_path = getattr(cfg.shape, "fpath", None) + if _is_usd_path(source_path): + descriptor, materials = rigid_desc_from_usd(cfg, per_env=True) + else: + descriptor, materials = rigid_desc_from_cfg(cfg, per_env=True) + self._spawn_scene.builder.materials.update(materials) rigid_obj = RigidObject( cfg=cfg, - entities=obj_list, + entities=None, device=self.device, + declared_num_instances=self.sim_config.num_envs, ) - if cfg.shape.visual_material: - mat = self.create_visual_material(cfg.shape.visual_material) - rigid_obj.set_visual_material(mat, update_default=True) + def bind_rigid_object(result, handles) -> None: + if rigid_obj.is_declared: + rigid_obj.bind_spawn(result, handles) + was_materialized = self.spawn_result is not None + self._spawn_scene.declare( + "rigid_object", + uid, + descriptor, + on_bind=bind_rigid_object, + ) self._rigid_objects[uid] = rigid_obj - self._invalidate_newton_physics() self.notify_visualization_topology_changed() + # Preserve the legacy immediate-availability behavior for runtime + # additions. Initial environment construction still batches all + # declarations into one finalize at BaseEnv's prepare boundary. + if was_materialized: + self.prepare() return rigid_obj def add_soft_object(self, cfg: SoftObjectCfg) -> SoftObject: @@ -1484,33 +1683,46 @@ def add_soft_object(self, cfg: SoftObjectCfg) -> SoftObject: SoftObject: The added soft object instance handle. """ if not self.physics.supports_soft_bodies: - logger.log_error( - f"Soft object support is not enabled for the " - f"{self.physics.name} backend yet.", - error_type=NotImplementedError, + raise NotImplementedError( + f"The {self.physics.name} backend does not support soft bodies." + ) + if self.device.type != "cuda": + raise NotImplementedError("SoftObject currently requires a CUDA device.") + if self.spawn_result is not None: + raise NotImplementedError( + "DexSim Spawn does not yet support adding a soft body after finalize." ) - - if not self.is_use_gpu_physics: - logger.log_error("Soft object requires GPU physics to be enabled.") - - from embodichain.lab.sim.utility import ( - load_soft_object_from_cfg, - ) - uid = cfg.uid if uid is None: - logger.log_error("Soft object uid must be specified.") + raise ValueError("Soft object uid must be specified.") + if uid in self._soft_objects: + raise ValueError(f"Soft object {uid!r} already exists.") - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = load_soft_object_from_cfg( - cfg=cfg, - env_list=env_list, + descriptor, materials = soft_desc_from_cfg(cfg, per_env=True) + self._spawn_scene.builder.materials.update(materials) + soft_object = SoftObject( + cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, ) - soft_obj = SoftObject(cfg=cfg, entities=obj_list, device=self.device) - self._soft_objects[uid] = soft_obj + def bind_soft_object(result, handles) -> None: + if soft_object.is_declared: + if cfg.shape.compute_uv: + for handle in handles: + handle.compute_uv_mapping() + soft_object.bind_spawn(result, handles) + + self._spawn_scene.declare( + "soft_object", + uid, + descriptor, + on_bind=bind_soft_object, + ) + self._soft_objects[uid] = soft_object self.notify_visualization_topology_changed() - return soft_obj + return soft_object def add_cloth_object(self, cfg: ClothObjectCfg) -> ClothObject: """Add a cloth object to the scene. @@ -1522,33 +1734,46 @@ def add_cloth_object(self, cfg: ClothObjectCfg) -> ClothObject: ClothObject: The added cloth object instance handle. """ if not self.physics.supports_cloth: - logger.log_error( - f"Cloth object support is not enabled for the " - f"{self.physics.name} backend yet.", - error_type=NotImplementedError, + raise NotImplementedError( + f"The {self.physics.name} backend does not support cloth bodies." + ) + if self.device.type != "cuda": + raise NotImplementedError("ClothObject currently requires a CUDA device.") + if self.spawn_result is not None: + raise NotImplementedError( + "DexSim Spawn does not yet support adding cloth after finalize." ) - - if not self.is_use_gpu_physics: - logger.log_error("Cloth object requires GPU physics to be enabled.") - - from embodichain.lab.sim.utility import ( - load_cloth_object_from_cfg, - ) - uid = cfg.uid if uid is None: - logger.log_error("Cloth object uid must be specified.") + raise ValueError("Cloth object uid must be specified.") + if uid in self._cloth_objects: + raise ValueError(f"Cloth object {uid!r} already exists.") - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = load_cloth_object_from_cfg( - cfg=cfg, - env_list=env_list, + descriptor, materials = cloth_desc_from_cfg(cfg, per_env=True) + self._spawn_scene.builder.materials.update(materials) + cloth_object = ClothObject( + cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, ) - cloth_obj = ClothObject(cfg=cfg, entities=obj_list, device=self.device) - self._cloth_objects[uid] = cloth_obj + def bind_cloth_object(result, handles) -> None: + if cloth_object.is_declared: + if cfg.shape.compute_uv: + for handle in handles: + handle.compute_uv_mapping() + cloth_object.bind_spawn(result, handles) + + self._spawn_scene.declare( + "cloth_object", + uid, + descriptor, + on_bind=bind_cloth_object, + ) + self._cloth_objects[uid] = cloth_object self.notify_visualization_topology_changed() - return cloth_obj + return cloth_object def get_rigid_object(self, uid: str) -> RigidObject | None: """Get a rigid object by its unique ID. @@ -1607,20 +1832,7 @@ def _broadcast_frame( env_ids: Sequence[int], name: str, ) -> list[np.ndarray]: - """Broadcast a local-frame spec to one matrix per target env. - - Args: - frame: None -> identity; (4,4) -> repeated; (N,4,4) -> indexed per env. - num_envs: Total number of arenas (used to validate (N,4,4)). - env_ids: Target env indices to produce frames for. - name: Constraint name (for error messages). - - Returns: - A list of (4,4) numpy arrays, one per env in env_ids. - - Raises: - RuntimeError: If an (N,4,4) frame's N != num_envs, or shape is invalid. - """ + """Broadcast a local constraint frame to the selected environments.""" if frame is None: identity = np.eye(4, dtype=np.float32) return [identity for _ in env_ids] @@ -1670,15 +1882,11 @@ def create_rigid_constraint( cfg: RigidConstraintCfg, env_ids: Sequence[int] | torch.Tensor | None = None, ) -> RigidConstraint: - """Create a fixed constraint between two RigidObjects. + """Create a fixed constraint between two rigid objects. - Binds ``rigid_object_a``'s entity[i] to ``rigid_object_b``'s entity[i] - within arena[i], for each env in ``env_ids``. Local frames default to - welding the objects at their *current* relative pose: - ``local_frame_a`` defaults to identity (object A's origin) and - ``local_frame_b`` defaults to ``inv(pose_B) @ pose_A`` (computed per env), - so the offset is preserved rather than the two origins being pulled - together. Pass explicit frames to define a specific joint frame. + Constraints are native Default/PhysX resources owned by each Arena. + Spawn owns the two actors; this method only borrows their native actor + handles while creating the constraint. Args: cfg: The constraint configuration. @@ -1686,20 +1894,18 @@ def create_rigid_constraint( the :class:`EventManager`) or a sequence of ints. None -> all arenas. Returns: - The created :class:`RigidConstraint`. - - Raises: - RuntimeError: If either object is missing, the name is already in use, - a frame shape is invalid, or dexsim fails to create a handle. + The created constraint batch. """ - # validate constraint type (only fixed supported in v1) + if hasattr(self, "physics") and not self.is_default_backend: + raise NotImplementedError( + "Rigid constraints are currently supported only by the Default/PhysX " + "backend." + ) if cfg.constraint_type != "fixed": logger.log_error( f"Constraint '{cfg.name}' has unsupported type " - f"'{cfg.constraint_type}'. Only 'fixed' is supported in v1." + f"'{cfg.constraint_type}'. Only 'fixed' is supported." ) - - # resolve objects if cfg.rigid_object_a_uid not in self._rigid_objects: logger.log_error( f"RigidObject '{cfg.rigid_object_a_uid}' not found for constraint " @@ -1710,16 +1916,16 @@ def create_rigid_constraint( f"RigidObject '{cfg.rigid_object_b_uid}' not found for constraint " f"'{cfg.name}'. Available: {list(self._rigid_objects.keys())}." ) - rigid_object_a = self._rigid_objects[cfg.rigid_object_a_uid] - rigid_object_b = self._rigid_objects[cfg.rigid_object_b_uid] - - # validate duplicate name if cfg.name in self._constraints: logger.log_error( f"Constraint '{cfg.name}' already exists. Remove it before recreating." ) - # validate object entity counts match num_envs + rigid_object_a = self._rigid_objects[cfg.rigid_object_a_uid] + rigid_object_b = self._rigid_objects[cfg.rigid_object_b_uid] + if hasattr(self, "_spawn_scene"): + self.prepare() + num_envs = self.num_envs if rigid_object_a.num_instances != num_envs: logger.log_error( @@ -1732,50 +1938,52 @@ def create_rigid_constraint( f"{rigid_object_b.num_instances} instances but num_envs is {num_envs}." ) - # resolve target env_ids (accepts None / tensor / sequence) target_env_ids = self._normalize_env_ids(env_ids, num_envs) - - # broadcast local frames. - # local_frame_a defaults to identity (object A's origin). - # local_frame_b defaults to the current relative pose of A w.r.t. B - # (inv(pose_B) @ pose_A), so that with both frames left as None the - # constraint welds the objects at their *current* relative pose instead - # of pulling their origins together. frames_a = self._broadcast_frame( cfg.local_frame_a, num_envs, target_env_ids, cfg.name ) if cfg.local_frame_b is None: pose_a = rigid_object_a.get_local_pose(to_matrix=True) pose_b = rigid_object_b.get_local_pose(to_matrix=True) - frame_b = torch.bmm(pose_inv(pose_b), pose_a) # (N, 4, 4) - frame_b = frame_b.cpu().numpy().astype(np.float32) + frame_b = ( + torch.bmm(pose_inv(pose_b), pose_a).cpu().numpy().astype(np.float32) + ) frames_b = [frame_b[i] for i in target_env_ids] else: frames_b = self._broadcast_frame( cfg.local_frame_b, num_envs, target_env_ids, cfg.name ) - # pre-size handles list with None, fill target envs handles: list = [None] * num_envs try: - for idx, env_id in enumerate(target_env_ids): + for index, env_id in enumerate(target_env_ids): + actor_a = rigid_object_a._entities[env_id] + actor_b = rigid_object_b._entities[env_id] + if getattr(rigid_object_a, "is_spawn_bound", False) is True: + actor_a = actor_a.native + if getattr(rigid_object_b, "is_spawn_bound", False) is True: + actor_b = actor_b.native + if actor_a is None or actor_b is None: + logger.log_error( + f"Constraint '{cfg.name}' references a released Spawn actor " + f"in environment {env_id}." + ) + arena = self.get_env(env_id) - name_i = cfg.name if num_envs <= 1 else f"{cfg.name}_{env_id}" + name = cfg.name if num_envs <= 1 else f"{cfg.name}_{env_id}" handle = arena.create_fixed_constraint( - name_i, - rigid_object_a._entities[env_id], - rigid_object_b._entities[env_id], - frames_a[idx], - frames_b[idx], + name, + actor_a, + actor_b, + frames_a[index], + frames_b[index], ) if handle is None: logger.log_error( - f"Failed to create constraint '{name_i}' in arena {env_id}." + f"Failed to create constraint '{name}' in arena {env_id}." ) handles[env_id] = handle except Exception: - # Ensure partially created per-arena constraints are removed if a later - # arena fails, so create/remove semantics stay consistent. RigidConstraint( cfg=cfg, constraint_handles=handles, @@ -1871,53 +2079,70 @@ def add_rigid_object_group(self, cfg: RigidObjectGroupCfg) -> RigidObjectGroup: Args: cfg (RigidObjectGroupCfg): Configuration for the rigid object group. + + Returns: + The stable Group facade. During initial scene construction it is + bound to Spawn handles by :meth:`prepare`. """ if not self.physics.supports_rigid_object_group: - logger.log_error( - f"Rigid object group support is not enabled for the " - f"{self.physics.name} backend yet.", - error_type=NotImplementedError, + raise NotImplementedError( + f"The {self.physics.name} backend does not support rigid object groups." ) - - from embodichain.lab.sim.utility.sim_utils import ( - load_mesh_objects_from_cfg, - ) - uid = cfg.uid if uid is None: - logger.log_error("Rigid object group uid must be specified.") + raise ValueError("Rigid object group uid must be specified.") if uid in self._rigid_object_groups: - logger.log_error(f"Rigid object group {uid} already exists.") - + raise ValueError(f"Rigid object group {uid!r} already exists.") if cfg.body_type == "static": - logger.log_error("Rigid object group cannot be static.") - - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - - obj_group_list = [] - for key, rigid_cfg in tqdm( - cfg.rigid_objects.items(), desc="Loading rigid objects" - ): - obj_list = load_mesh_objects_from_cfg( - cfg=rigid_cfg, - env_list=env_list, - cache_dir=self._convex_decomp_dir, - ) - obj_group_list.append(obj_list) + raise ValueError("Rigid object group cannot be static.") + if not cfg.rigid_objects: + raise ValueError("Rigid object group must contain at least one object.") + + actor_type = { + "dynamic": ActorType.DYNAMIC, + "kinematic": ActorType.KINEMATIC, + }[cfg.body_type] + descriptors = [] + for index, member in enumerate(cfg.rigid_objects.values()): + member_cfg = deepcopy(member) + member_cfg.uid = f"{uid}__member_{index}" + member_cfg.body_type = cfg.body_type + source_path = getattr(member_cfg.shape, "fpath", None) + if _is_usd_path(source_path): + descriptor, materials = rigid_desc_from_usd(member_cfg, per_env=True) + else: + descriptor, materials = rigid_desc_from_cfg(member_cfg, per_env=True) + if descriptor.physics is None: + raise ValueError( + f"Rigid object group member {index} has no rigid-body physics." + ) + descriptor.physics.actor_type = actor_type + self._spawn_scene.builder.materials.update(materials) + descriptors.append(descriptor) - # Convert [a1, a2, ...], [b1, b2, ...] to [(a1, b1, ...), (a2, b2, ...), ...] - obj_group_list = list(zip(*obj_group_list)) - rigid_obj_group = RigidObjectGroup( - cfg=cfg, - entities=obj_group_list, + group = RigidObjectGroup( + cfg, + entities=None, device=self.device, + declared_num_instances=self.sim_config.num_envs, ) - self._rigid_object_groups[uid] = rigid_obj_group - self._invalidate_newton_physics() - self.notify_visualization_topology_changed() + def bind_group(result, handles) -> None: + if group.is_declared: + group.bind_spawn(result, handles) - return rigid_obj_group + was_materialized = self.spawn_result is not None + self._spawn_scene.declare( + "rigid_object_group", + uid, + tuple(descriptors), + on_bind=bind_group, + ) + self._rigid_object_groups[uid] = group + self.notify_visualization_topology_changed() + if was_materialized: + self.prepare() + return group def get_rigid_object_group(self, uid: str) -> RigidObjectGroup | None: """Get a rigid object group by its unique ID. @@ -1988,39 +2213,21 @@ def add_articulation( """ uid = cfg.uid if uid is None: + if cfg.fpath is None: + raise ValueError( + "Articulation configuration must provide fpath when uid " + "is not specified." + ) uid = os.path.splitext(os.path.basename(cfg.fpath))[0] cfg.uid = uid if uid in self._articulations: - logger.log_error(f"Articulation {uid} already exists.") - - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = [] - - is_usd = cfg.fpath.endswith((".usd", ".usda", ".usdc")) - if is_usd: - from embodichain.lab.sim.utility.sim_utils import ( - spawn_usd_articulation_entities, - ) - - obj_list = spawn_usd_articulation_entities( - cfg, env_list, cache_dir=self._convex_decomp_dir - ) - else: - # non-usd file does not support this option, will be forced set False to avoid potential issues. - cfg.use_usd_properties = False - - from embodichain.lab.sim.utility.sim_utils import ( - spawn_articulation_entities, - ) - - obj_list = spawn_articulation_entities(cfg, env_list) - - articulation = Articulation(cfg=cfg, entities=obj_list, device=self.device) + raise ValueError(f"Articulation {uid!r} already exists.") + was_materialized = self.spawn_result is not None + articulation = self._declare_spawn_articulation(cfg, Articulation) self._articulations[uid] = articulation - self._invalidate_newton_physics() - self.notify_visualization_topology_changed() - + if was_materialized: + self.prepare() return articulation def get_articulation(self, uid: str) -> Articulation | None: @@ -2084,33 +2291,73 @@ def add_robot(self, cfg: RobotCfg) -> Robot | None: logger.log_error(f"Robot {uid} already exists.") return self._robots[uid] - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = [] + was_materialized = self.spawn_result is not None + robot = self._declare_spawn_articulation(cfg, Robot) + self._robots[uid] = robot + if was_materialized: + self.prepare() + return robot - is_usd = cfg.fpath.endswith((".usd", ".usda", ".usdc")) - if is_usd: - from embodichain.lab.sim.utility.sim_utils import ( - spawn_usd_articulation_entities, - ) + def _declare_spawn_articulation( + self, + cfg: ArticulationCfg, + facade_type: type[Articulation], + ) -> Articulation: + """Declare an articulation facade and bind it after Spawn finalize. - obj_list = spawn_usd_articulation_entities(cfg, env_list) + DexSim remains the sole articulation source loader. The facade is + intentionally metadata-empty during scene declaration; once the + adapter has loaded the source exactly once, the bind callback creates + its batch view from the resolved link/joint metadata and applies the + supported live values directly from its EmbodiChain config. + """ + if _is_usd_path(cfg.fpath): + descriptor, materials = articulation_desc_from_usd( + cfg, + per_env=True, + ) + self._spawn_scene.builder.materials.update(materials) else: - # non-usd file does not support this option, will be forced set False to avoid potential issues. - cfg.use_usd_properties = False - - from embodichain.lab.sim.utility.sim_utils import ( - spawn_articulation_entities, + descriptor = articulation_desc_from_cfg(cfg, per_env=True) + if self.is_newton_backend and cfg.qpos_limits is not None: + # Reject before mutating SceneBuilder. Applying this after bind + # would immediately make Newton's immutable model stale. + raise NotImplementedError( + "Newton articulation qpos_limits are not yet supported by the " + "metadata-after-finalize binding path. TODO: add a retained-desc " + "configuration phase that runs before Newton model finalize." ) + if cfg.uid is None: + cfg.uid = descriptor.name - obj_list = spawn_articulation_entities(cfg, env_list) + facade = facade_type( + cfg=cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) - robot = Robot(cfg=cfg, entities=obj_list, device=self.device) + def bind_articulation(result, handles) -> None: + if facade.is_declared: + facade.bind_spawn(result, handles) - self._robots[uid] = robot - self._invalidate_newton_physics() + self._spawn_scene.declare( + "articulation", + descriptor.name, + descriptor, + on_bind=bind_articulation, + ) self.notify_visualization_topology_changed() + return facade - return robot + @staticmethod + def _raise_spawn_feature_todo(feature: str, required_api: str) -> None: + """Reject topology that is not owned by the active Spawn scene.""" + raise NotImplementedError( + f"Spawn scene construction does not integrate {feature} yet. " + f"TODO: route it through {required_api}; falling back to direct " + "Arena construction would create a second topology owner." + ) def get_robot(self, uid: str) -> Robot | None: """Get a Robot by its unique ID. @@ -2388,7 +2635,12 @@ def set_gizmo_visibility( gizmo.set_visible(visible) def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: - """General interface to add a sensor to the scene and returns a handle. + """Create a sensor on the pre-created simulation Arenas. + + Cameras keep EmbodiChain's native CameraGroup implementation. A camera + attached to an articulation link is created immediately and attached + after the physical Spawn scene is prepared. ContactSensor still + requires the Default/PhysX scene and therefore prepares physics first. Args: sensor_cfg (SensorCfg): configuration for the sensor. @@ -2397,28 +2649,123 @@ def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: BaseSensor: The added sensor instance handle. """ sensor_type = sensor_cfg.sensor_type - if sensor_type not in self.SUPPORTED_SENSOR_TYPES: - logger.log_warning(f"Unsupported sensor type: {sensor_type}") - return None + uid = sensor_cfg.uid + if uid is None: + uid = f"{sensor_type.lower()}_{len(self._sensors)}" + sensor_cfg.uid = uid + if uid in self._sensors: + raise ValueError(f"Sensor {uid!r} already exists.") - sensor_uid = sensor_cfg.uid - if sensor_uid is None: - sensor_uid = f"{sensor_type.lower()}_{len(self._sensors)}" - sensor_cfg.uid = sensor_uid + sensor_factory = self.SUPPORTED_SENSOR_TYPES.get(sensor_type) + if sensor_factory is None: + raise ValueError( + f"Unsupported sensor type {sensor_type!r}. Supported types: " + f"{sorted(self.SUPPORTED_SENSOR_TYPES)}." + ) + if sensor_type == "ContactSensor" and self.is_newton_backend: + raise NotImplementedError( + "ContactSensor currently requires the Default/PhysX PhysicsScene. " + "Newton needs a public backend-neutral contact query API in DexSim." + ) - if sensor_uid in self._sensors: - logger.log_warning(f"Sensor {sensor_uid} already exists.") - return None + if isinstance(sensor_factory, type) and issubclass(sensor_factory, Camera): + if len(self._arenas) != self.num_envs: + raise RuntimeError( + "Camera creation requires all Spawn Arenas to be " + f"prepared ({len(self._arenas)} of {self.num_envs} ready)." + ) + sensor = sensor_factory( + sensor_cfg, + self.device, + world=self._world, + arenas=self._arenas, + parent_node_resolver=self._resolve_spawn_sensor_parent_nodes, + defer_parent_attachment=True, + ) + if sensor_cfg.extrinsics.parent is not None: + scene = self._spawn_scene + if ( + scene.result is not None + and not scene.result.needs_rebuild + and not scene.builder.has_pending_changes + ): + sensor.attach_to_parent() + else: + self._pending_sensor_attachments.append(sensor) + else: + # ContactSensor and custom native sensors require a prepared + # physics scene; cameras only depend on the pre-created Arenas. + self.prepare() + # Preserve custom test/plugin factories whose two-argument + # constructor predates the manager-owned render context. + sensor = sensor_factory(sensor_cfg, self.device) + + self._sensors[uid] = sensor + self.notify_visualization_topology_changed() + return sensor - sensor = self.SUPPORTED_SENSOR_TYPES[sensor_type](sensor_cfg, self.device) + def _resolve_spawn_sensor_parent_nodes(self, parent: str) -> list[object]: + """Resolve one canonical articulation link to a render node per Arena. - self._sensors[sensor_uid] = sensor - if isinstance(sensor, Camera): - self.notify_visualization_topology_changed() + A plain link name remains compatible with existing CameraCfg values. + When more than one robot/articulation owns that link, callers can use + ``"/"`` to disambiguate without introducing + backend clone suffixes. + """ + assets: dict[str, Articulation] = { + **self._articulations, + **self._robots, + } + asset_uid: str | None = None + link_name = parent + if "/" in parent: + candidate_uid, candidate_link = parent.split("/", maxsplit=1) + if candidate_uid in assets: + asset_uid = candidate_uid + link_name = candidate_link + + matches: list[tuple[str, list[object]]] = [] + for uid, asset in assets.items(): + if asset_uid is not None and uid != asset_uid: + continue + if not getattr(asset, "is_spawn_bound", False): + continue + handles = list(getattr(asset, "_entities", ())) + if len(handles) != self.num_envs: + continue + if link_name not in handles[0].get_link_names(): + continue - # Check if the sensor needs to change the parent frame. + nodes: list[object] = [] + for handle in handles: + if link_name not in handle.get_link_names(): + raise RuntimeError( + f"Articulation {uid!r} has heterogeneous link topology; " + f"link {link_name!r} is missing in one Arena." + ) + render_body = handle.get_render_body(link_name) + if render_body is None: + raise RuntimeError( + f"Articulation {uid!r} link {link_name!r} has no public " + "render node for camera attachment." + ) + nodes.append(render_body.render_node()) + matches.append((uid, nodes)) - return sensor + if len(matches) == 1: + return matches[0][1] + if len(matches) > 1: + owners = ", ".join(uid for uid, _ in matches) + raise ValueError( + f"Camera parent link {link_name!r} is ambiguous across assets " + f"[{owners}]; use '/{link_name}'." + ) + scope = f" on asset {asset_uid!r}" if asset_uid is not None else "" + raise ValueError( + f"Camera parent link {link_name!r} was not found{scope} in any " + "Spawn-bound Robot or Articulation. Attachment to arbitrary render " + "nodes is not yet supported by the Spawn-only bridge." + ) def get_sensor(self, uid: str) -> BaseSensor | None: """Get a sensor by its UID. @@ -2445,53 +2792,41 @@ def get_sensor_uid_list(self) -> List[str]: def remove_asset(self, uid: str) -> bool: """Remove an asset by its UID. - The asset can be a light, sensor, robot, rigid object or articulation. - - Note: - Currently, lights and sensors are not supported to be removed. + Native render lights are not removed by this method. Sensors and + Spawn-owned physical assets are supported. Args: uid (str): The UID of the asset. Returns: bool: True if the asset is removed successfully, otherwise False. """ - if uid in self._rigid_objects: - obj = self._rigid_objects.pop(uid) - obj.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._soft_objects: - obj = self._soft_objects.pop(uid) - obj.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._cloth_objects: - obj = self._cloth_objects.pop(uid) - obj.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._rigid_object_groups: - group = self._rigid_object_groups.pop(uid) - group.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._articulations: - art = self._articulations.pop(uid) - art.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._robots: - robot = self._robots.pop(uid) - robot.destroy() + if uid in self._sensors: + sensor = self._sensors.pop(uid) + if sensor in self._pending_sensor_attachments: + self._pending_sensor_attachments.remove(sensor) + destroy = getattr(sensor, "destroy", None) + if callable(destroy): + destroy() self.notify_visualization_topology_changed() return True - return False + scene = self._spawn_scene + if uid not in scene: + return False + if uid == "default_plane": + raise ValueError("The Spawn-owned default plane cannot be removed.") + + was_materialized = scene.result is not None + scene.remove(uid) + if was_materialized: + self.prepare() + + self._rigid_objects.pop(uid, None) + self._rigid_object_groups.pop(uid, None) + self._articulations.pop(uid, None) + self._robots.pop(uid, None) + self.notify_visualization_topology_changed() + return True def draw_marker( self, @@ -3339,6 +3674,42 @@ def _deferred_destroy(self) -> None: import sys, gc + # Render-only cameras may be attached to Spawn articulation link + # nodes. Remove their Arena views before closing SpawnResult, which + # releases those parent nodes, and before World.quit releases their + # CameraGroups. + for sensor in list(getattr(self, "_sensors", {}).values()): + try: + sensor.destroy() + except Exception as error: + logger.log_warning( + f"Failed to destroy sensor {getattr(sensor, 'uid', None)!r}: " + f"{error!r}" + ) + + if self._spawn_scene is not None: + # Release result-scoped batches/facades before closing the + # SpawnResult and, finally, the World that owns native resources. + for registry_name in ( + "_rigid_objects", + "_rigid_object_groups", + "_soft_objects", + "_cloth_objects", + "_articulations", + "_robots", + ): + for asset in getattr(self, registry_name, {}).values(): + if hasattr(asset, "_data"): + asset._data = None + if hasattr(asset, "_spawn_result"): + asset._spawn_result = None + if hasattr(asset, "_entities"): + asset._entities = [] + try: + self._spawn_scene.close() + finally: + self._spawn_scene = None + self.clean_materials() if self._env: diff --git a/embodichain/lab/sim/spawn/__init__.py b/embodichain/lab/sim/spawn/__init__.py new file mode 100644 index 000000000..ac930632e --- /dev/null +++ b/embodichain/lab/sim/spawn/__init__.py @@ -0,0 +1,36 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Translate EmbodiChain asset configs into DexSim Spawn descriptors.""" + +from __future__ import annotations + +from .descriptors import ( + articulation_desc_from_cfg, + cloth_desc_from_cfg, + rigid_desc_from_cfg, + soft_desc_from_cfg, +) +from .usd import articulation_desc_from_usd, rigid_desc_from_usd + +__all__ = [ + "articulation_desc_from_cfg", + "articulation_desc_from_usd", + "cloth_desc_from_cfg", + "rigid_desc_from_cfg", + "rigid_desc_from_usd", + "soft_desc_from_cfg", +] diff --git a/embodichain/lab/sim/spawn/descriptors.py b/embodichain/lab/sim/spawn/descriptors.py new file mode 100644 index 000000000..709ba976e --- /dev/null +++ b/embodichain/lab/sim/spawn/descriptors.py @@ -0,0 +1,504 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Translate EmbodiChain asset configurations into DexSim Spawn descriptors. + +This module is deliberately independent of the active physics backend. It +translates one EmbodiChain configuration into a canonical descriptor carrying +both the common physics values and the optional backend extension blocks. The +selected :mod:`dexsim.spawn` adapter remains the only component that chooses +between PhysX and Newton. + +Articulation joint and link names are resolved by the normal DexSim adapter +finalization, not by a second source parser in EmbodiChain. Configuration that +depends on those names is applied directly from the EmbodiChain config after +the facade binds to the finalized result. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import MISSING, fields +import math +import os +from typing import TYPE_CHECKING + +import numpy as np +from dexsim.spawn import ( + ArticulationDesc, + ClothObjectDesc, + CollisionApproximation, + CollisionDesc, + DexsimCollisionDesc, + DexsimPhysicsDesc, + GeometryDesc, + MaterialDesc, + NewtonCollisionDesc, + NewtonJointDesc, + ObjectDesc, + RenderDesc, + RigidBodyPhysicsDesc, + SoftObjectDesc, +) +from dexsim.types import ActorType + +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + ClothObjectCfg, + RigidBodyAttributesCfg, + RigidObjectCfg, + SoftObjectCfg, +) +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, SphereCfg +from embodichain.utils import logger + +if TYPE_CHECKING: + from embodichain.lab.sim.material import VisualMaterialCfg + +__all__ = [ + "articulation_desc_from_cfg", + "cloth_desc_from_cfg", + "rigid_desc_from_cfg", + "soft_desc_from_cfg", +] + + +def rigid_desc_from_cfg( + cfg: RigidObjectCfg, + *, + per_env: bool = True, +) -> tuple[ObjectDesc, dict[str, MaterialDesc]]: + """Translate a rigid-object config into a DexSim Spawn descriptor.""" + uid = _required_uid(cfg.uid, "Rigid object") + if isinstance(cfg.shape, MeshCfg) and _is_usd_path(cfg.shape.fpath): + raise NotImplementedError( + "USD files describe typed scenes; use rigid_desc_from_usd() to " + "select the sole rigid object." + ) + + geometry, approximation, max_hulls = _compile_geometry(cfg) + material_ref, material_entry = _compile_visual_material( + uid, cfg.shape.visual_material + ) + collision = CollisionDesc.from_geometry( + geometry, + approximation=approximation, + ) + collision.enable_collision = bool(cfg.attrs.enable_collision) + collision.decomp_max_hulls = max_hulls + collision.dexsim = _compile_dexsim_collision(cfg.attrs) + collision.newton = _compile_newton_collision( + cfg.attrs, + sdf_resolution=( + _resolved_mesh_collision_settings(cfg)[2] + if isinstance(cfg.shape, MeshCfg) + else 0 + ), + ) + collision.render_source_index = 0 + + descriptor = ObjectDesc( + name=uid, + pose=_pose_from_cfg(cfg), + renders=[RenderDesc.from_geometry(geometry, material_ref=material_ref)], + collisions=[collision], + physics=_compile_rigid_physics(cfg.attrs, cfg.body_type), + per_env=per_env, + body_scale=_vector3(cfg.body_scale, field_name="body_scale"), + ) + materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} + return descriptor, materials + + +def soft_desc_from_cfg( + cfg: SoftObjectCfg, + *, + per_env: bool = True, +) -> tuple[SoftObjectDesc, dict[str, MaterialDesc]]: + """Translate a soft-object config into a DexSim Spawn descriptor.""" + uid = _required_uid(cfg.uid, "Soft object") + if _is_missing(cfg.shape.fpath) or not str(cfg.shape.fpath).strip(): + raise ValueError("SoftObjectCfg.shape.fpath must be a non-empty path.") + geometry = GeometryDesc.mesh(file_path=str(cfg.shape.fpath), segment_name=uid) + material_ref, material_entry = _compile_visual_material( + uid, cfg.shape.visual_material + ) + descriptor = SoftObjectDesc( + name=uid, + pose=_pose_from_cfg(cfg), + renders=[RenderDesc.from_geometry(geometry, material_ref=material_ref)], + voxel_config=cfg.voxel_attr.attr(), + body_attr=cfg.physical_attr.attr(), + per_env=per_env, + ) + materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} + return descriptor, materials + + +def cloth_desc_from_cfg( + cfg: ClothObjectCfg, + *, + per_env: bool = True, +) -> tuple[ClothObjectDesc, dict[str, MaterialDesc]]: + """Translate a cloth-object config into a DexSim Spawn descriptor.""" + uid = _required_uid(cfg.uid, "Cloth object") + if _is_missing(cfg.shape.fpath) or not str(cfg.shape.fpath).strip(): + raise ValueError("ClothObjectCfg.shape.fpath must be a non-empty path.") + geometry = GeometryDesc.mesh(file_path=str(cfg.shape.fpath), segment_name=uid) + material_ref, material_entry = _compile_visual_material( + uid, cfg.shape.visual_material + ) + descriptor = ClothObjectDesc( + name=uid, + pose=_pose_from_cfg(cfg), + renders=[RenderDesc.from_geometry(geometry, material_ref=material_ref)], + body_attr=cfg.physical_attr.attr(), + per_env=per_env, + ) + materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} + return descriptor, materials + + +def articulation_desc_from_cfg( + cfg: ArticulationCfg, + *, + per_env: bool = True, + source_path: str | None = None, +) -> ArticulationDesc: + """Translate an articulation config into a DexSim Spawn descriptor.""" + path = source_path if source_path is not None else cfg.fpath + if path is None or not str(path).strip(): + raise ValueError( + "No articulation source path is available. Assemble the robot URDF " + "before converting its configuration." + ) + if _is_usd_path(path): + raise NotImplementedError( + "USD files describe typed scenes; use articulation_desc_from_usd() " + "to select the sole articulation." + ) + if cfg.use_usd_properties: + logger.log_warning( + "ArticulationCfg.use_usd_properties only applies to USD sources and " + "is ignored for URDF articulations." + ) + if cfg.min_position_iters != 4 or cfg.min_velocity_iters != 1: + logger.log_warning( + "Per-articulation solver iteration counts are not exposed by the " + "backend-neutral Spawn facade and were not applied." + ) + + target_mode = {"force": 3, "none": 0}.get(cfg.drive_pros.drive_type) + return ArticulationDesc( + name=_articulation_uid(cfg.uid, str(path)), + pose=_pose_from_cfg(cfg), + path=str(path), + urdf_path=str(path), + fixed_base=bool(cfg.fix_base), + enable_self_collision=not bool(cfg.disable_self_collision), + urdf_fix_root_link=bool(cfg.fix_base), + per_env=per_env, + body_scale=_vector3(cfg.body_scale, field_name="body_scale"), + newton_drive=( + None if target_mode is None else NewtonJointDesc(target_mode=target_mode) + ), + newton_collision=_compile_newton_collision(cfg.attrs), + ) + + +def _compile_rigid_physics( + attrs: RigidBodyAttributesCfg, + body_type: str, +) -> RigidBodyPhysicsDesc: + actor_types = { + "dynamic": ActorType.DYNAMIC, + "kinematic": ActorType.KINEMATIC, + "static": ActorType.STATIC, + } + try: + actor_type = actor_types[body_type] + except KeyError as exc: + raise ValueError( + f"Unsupported rigid body_type {body_type!r}; expected one of " + f"{tuple(actor_types)}." + ) from exc + + if attrs.mass is not None and attrs.mass < 0: + raise ValueError("Rigid-body mass cannot be negative.") + if attrs.mass == 0 and (attrs.density is None or attrs.density <= 0): + raise ValueError("Rigid-body density must be positive when mass is zero.") + + mass = float(attrs.mass) if attrs.mass is not None and attrs.mass > 0 else None + density = ( + float(attrs.density) + if mass is None and attrs.density is not None and attrs.density > 0 + else None + ) + return RigidBodyPhysicsDesc( + actor_type=actor_type, + mass=mass, + density=density, + dexsim=DexsimPhysicsDesc( + linear_damping=float(attrs.linear_damping), + angular_damping=float(attrs.angular_damping), + max_linear_velocity=float(attrs.max_linear_velocity), + max_angular_velocity=float(attrs.max_angular_velocity), + max_depenetration_velocity=float(attrs.max_depenetration_velocity), + enable_ccd=bool(attrs.enable_ccd), + min_position_iters=int(attrs.min_position_iters), + min_velocity_iters=int(attrs.min_velocity_iters), + sleep_threshold=float(attrs.sleep_threshold), + ), + ) + + +def _compile_dexsim_collision( + attrs: RigidBodyAttributesCfg, +) -> DexsimCollisionDesc: + return DexsimCollisionDesc( + dynamic_friction=float(attrs.dynamic_friction), + static_friction=float(attrs.static_friction), + restitution=float(attrs.restitution), + contact_offset=float(attrs.contact_offset), + rest_offset=float(attrs.rest_offset), + ) + + +def _compile_newton_collision( + attrs: RigidBodyAttributesCfg, + *, + sdf_resolution: int = 0, +) -> NewtonCollisionDesc: + # ``None`` means "leave the backend default untouched". Initializing every + # field avoids accidentally authoring NewtonCollisionDesc's convenience + # defaults when the EmbodiChain Newton sub-config did not set them. + values = {field.name: None for field in fields(NewtonCollisionDesc)} + if attrs.newton is not None: + for name in values: + if hasattr(attrs.newton, name): + values[name] = getattr(attrs.newton, name) + if "mu" in values: + values["mu"] = float(attrs.dynamic_friction) + if "restitution" in values: + values["restitution"] = float(attrs.restitution) + if sdf_resolution > 0: + if "force_sdf" in values: + values["force_sdf"] = True + if values["sdf_max_resolution"] is None: + values["sdf_max_resolution"] = int(sdf_resolution) + return NewtonCollisionDesc(**values) + + +def _compile_geometry( + cfg: RigidObjectCfg, +) -> tuple[GeometryDesc, CollisionApproximation, int]: + shape = cfg.shape + if isinstance(shape, MeshCfg): + if _is_missing(shape.fpath) or not str(shape.fpath).strip(): + raise ValueError("MeshCfg.fpath must be a non-empty path.") + max_hulls, acd_method, sdf_resolution = _resolved_mesh_collision_settings(cfg) + if sdf_resolution > 0: + approximation = CollisionApproximation.SDF + elif max_hulls > 1: + approximation = CollisionApproximation.CONVEX_DECOMPOSITION + else: + approximation = CollisionApproximation.CONVEX_HULL + + option = shape.load_option + if any( + ( + option.rebuild_normals, + option.rebuild_tangent, + option.rebuild_3rdnormal, + option.rebuild_3rdtangent, + option.smooth != -1.0, + ) + ): + logger.log_warning( + "Mesh LoadOption is not represented by ObjectDesc; the Spawn " + "adapter will use its default mesh loading policy." + ) + if shape.compute_uv: + logger.log_warning( + "Mesh UV projection is not represented by GeometryDesc and was " + "not applied." + ) + if max_hulls > 1 and str(acd_method).lower() != "coacd": + logger.log_warning( + f"Spawn preserves max_convex_hull_num={max_hulls}, but does not " + f"expose the requested ACD method {acd_method!r}." + ) + if sdf_resolution > 0: + logger.log_warning( + "CollisionApproximation.SDF is preserved and Newton receives " + "sdf_max_resolution, but the PhysX descriptor does not expose " + "its cooking resolution." + ) + return ( + GeometryDesc.mesh( + file_path=str(shape.fpath), segment_name=cfg.uid or "mesh" + ), + approximation, + max(1, max_hulls), + ) + + if isinstance(shape, CubeCfg): + size = tuple(float(value) for value in shape.size) + if len(size) != 3 or any(value <= 0 for value in size): + raise ValueError("CubeCfg.size must contain three positive values.") + return GeometryDesc.cube(size), CollisionApproximation.NONE, 1 + + if isinstance(shape, SphereCfg): + if shape.radius <= 0: + raise ValueError("SphereCfg.radius must be positive.") + return ( + GeometryDesc.sphere(float(shape.radius)), + CollisionApproximation.NONE, + 1, + ) + + raise NotImplementedError( + f"RigidObjectCfg shape {type(shape).__name__!r} is not supported by " + "the Spawn converter; supported shapes are MeshCfg, CubeCfg, and SphereCfg." + ) + + +def _compile_visual_material( + object_uid: str, + cfg: VisualMaterialCfg | None, +) -> tuple[str | None, tuple[str, MaterialDesc] | None]: + if cfg is None: + return None, None + key = str(cfg.uid or f"{object_uid}_material") + base_color = tuple(float(value) for value in cfg.base_color) + if len(base_color) != 4: + raise ValueError("VisualMaterialCfg.base_color must be RGBA.") + emissive_rgb = tuple( + float(value) * float(cfg.emissive_intensity) for value in cfg.emissive + ) + if len(emissive_rgb) != 3: + raise ValueError("VisualMaterialCfg.emissive must be RGB.") + desc = MaterialDesc( + name=key, + base_color=base_color, + base_color_map=cfg.base_color_texture, + normal_map=cfg.normal_texture, + emissive=(*emissive_rgb, 1.0), + roughness=float(cfg.roughness), + roughness_map=cfg.roughness_texture, + metallic=float(cfg.metallic), + metallic_map=cfg.metallic_texture, + ao_map=cfg.ao_texture, + ior=float(cfg.ior), + ) + return key, (key, desc) + + +def _resolved_mesh_collision_settings( + cfg: RigidObjectCfg, +) -> tuple[int, str, int]: + if not isinstance(cfg.shape, MeshCfg): + return 1, "coacd", 0 + + def first_value(values: Sequence[object], default: object) -> object: + for value in values: + if not _is_missing(value): + return value + return default + + max_hulls = int( + first_value((cfg.max_convex_hull_num, cfg.shape.max_convex_hull_num), 1) + ) + acd_method = str(first_value((cfg.acd_method, cfg.shape.acd_method), "coacd")) + sdf_resolution = int(first_value((cfg.sdf_resolution, cfg.shape.sdf_resolution), 0)) + if max_hulls < 1: + raise ValueError("max_convex_hull_num must be at least 1.") + if sdf_resolution < 0: + raise ValueError("sdf_resolution cannot be negative.") + return max_hulls, acd_method, sdf_resolution + + +def _pose_from_cfg(cfg: object) -> np.ndarray: + local_pose = getattr(cfg, "init_local_pose", None) + if local_pose is not None: + pose = np.asarray(local_pose, dtype=np.float32).reshape(4, 4).copy() + else: + position = _vector3(getattr(cfg, "init_pos"), field_name="init_pos") + rotation_deg = _vector3(getattr(cfg, "init_rot"), field_name="init_rot") + rx, ry, rz = np.deg2rad(rotation_deg) + cx, sx = math.cos(rx), math.sin(rx) + cy, sy = math.cos(ry), math.sin(ry) + cz, sz = math.cos(rz), math.sin(rz) + rot_x = np.array( + ((1.0, 0.0, 0.0), (0.0, cx, -sx), (0.0, sx, cx)), + dtype=np.float32, + ) + rot_y = np.array( + ((cy, 0.0, sy), (0.0, 1.0, 0.0), (-sy, 0.0, cy)), + dtype=np.float32, + ) + rot_z = np.array( + ((cz, -sz, 0.0), (sz, cz, 0.0), (0.0, 0.0, 1.0)), + dtype=np.float32, + ) + pose = np.eye(4, dtype=np.float32) + # Match EmbodiChain's shared matrix_from_euler(..., "XYZ") contract + # used by the legacy RigidObject reset path. + pose[:3, :3] = rot_x @ rot_y @ rot_z + pose[:3, 3] = position + + if not np.isfinite(pose).all(): + raise ValueError("init_local_pose must contain finite values.") + if not np.allclose(pose[3], (0.0, 0.0, 0.0, 1.0), atol=1e-6): + raise ValueError("init_local_pose must be a homogeneous 4x4 transform.") + return pose + + +def _vector3(value: object, *, field_name: str) -> np.ndarray: + result = np.asarray(value, dtype=np.float32).reshape(-1) + if result.size != 3 or not np.isfinite(result).all(): + raise ValueError(f"{field_name} must contain three finite values.") + if field_name == "body_scale" and np.any(result <= 0): + raise ValueError("body_scale values must be positive.") + return result.copy() + + +def _required_uid(value: str | None, label: str) -> str: + if value is None or not str(value).strip(): + raise ValueError(f"{label} uid must be specified before Spawn conversion.") + uid = str(value) + if "/" in uid: + raise ValueError(f"{label} uid cannot contain '/': {uid!r}.") + return uid + + +def _articulation_uid(value: str | None, path: str | None) -> str: + if value is not None and str(value).strip(): + return _required_uid(str(value), "Articulation") + if path is None or not str(path).strip(): + raise ValueError( + "Articulation uid is required when its source path is unresolved." + ) + inferred = os.path.splitext(os.path.basename(str(path)))[0] + return _required_uid(inferred, "Articulation") + + +def _is_usd_path(path: object) -> bool: + return str(path).lower().endswith((".usd", ".usda", ".usdc")) + + +def _is_missing(value: object) -> bool: + # ``@configclass`` deepcopy can create a distinct _MISSING_TYPE instance. + return value is MISSING or isinstance(value, type(MISSING)) diff --git a/embodichain/lab/sim/spawn/scene.py b/embodichain/lab/sim/spawn/scene.py new file mode 100644 index 000000000..e4a42b9ca --- /dev/null +++ b/embodichain/lab/sim/spawn/scene.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Thin EmbodiChain coordination around DexSim Spawn.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Literal + +__all__ = ["SpawnScene"] + +AssetBindCallback = Callable[[Any, tuple[Any, ...]], None] +_AssetKind = Literal[ + "rigid_object", + "rigid_object_group", + "articulation", + "soft_object", + "cloth_object", +] + + +@dataclass(slots=True) +class _AssetDeclaration: + kind: _AssetKind + descriptor: Any + on_bind: AssetBindCallback | None + + +class SpawnScene: + """Map EmbodiChain asset declarations onto one DexSim Spawn scene. + + DexSim's ``SceneBuilder`` and ``SpawnResult`` own lifecycle state and + revisions. This class only remembers how logical asset ids map to Spawn + paths and how the resulting handles bind back into EmbodiChain facades. + """ + + def __init__( + self, + world: Any, + *, + num_envs: int, + spacing: tuple[float, float, float] = (0.0, 0.0, 0.0), + ) -> None: + from dexsim.spawn import SceneBuilder + + self.builder = SceneBuilder(world) + self.builder.replicate( + count=num_envs, + spacing=spacing, + name_format="arena_{i}", + ) + self.result: Any | None = None + self._assets: dict[str, _AssetDeclaration] = {} + + @property + def arena_names(self) -> tuple[str, ...]: + """Names of the replicated per-environment Arenas.""" + return tuple(self.builder.replicate_plan.env_names()) + + def __contains__(self, uid: str) -> bool: + return uid in self._assets + + def declare( + self, + kind: _AssetKind, + uid: str, + descriptor: Any, + *, + on_bind: AssetBindCallback | None = None, + ) -> None: + """Add a descriptor to the Builder and remember its facade binding.""" + if uid in self._assets: + raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") + declaration = _AssetDeclaration( + kind=kind, + descriptor=descriptor, + on_bind=on_bind, + ) + + if kind == "rigid_object_group": + declaration.descriptor = tuple( + self.builder.add_object(member) for member in descriptor + ) + else: + add_name = { + "rigid_object": "add_object", + "articulation": "add_articulation", + "soft_object": "add_soft_object", + "cloth_object": "add_cloth_object", + }[kind] + declaration.descriptor = getattr(self.builder, add_name)(descriptor) + self._assets[uid] = declaration + + def track( + self, + kind: _AssetKind, + uid: str, + descriptor: Any, + *, + on_bind: AssetBindCallback | None = None, + ) -> None: + """Track a descriptor that was already added to ``SceneBuilder``.""" + if uid in self._assets: + raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") + self._assets[uid] = _AssetDeclaration(kind, descriptor, on_bind) + + def remove(self, uid: str) -> None: + """Remove a declared asset from its DexSim owner.""" + declaration = self._assets[uid] + if declaration.kind in {"soft_object", "cloth_object"}: + raise NotImplementedError( + "DexSim Spawn does not yet expose pending removal for " + f"{declaration.kind.replace('_', ' ')}." + ) + if declaration.kind == "rigid_object_group": + for member in declaration.descriptor: + self.builder.remove_object(member.name) + else: + remove_name = { + "rigid_object": "remove_object", + "articulation": "remove_articulation", + }[declaration.kind] + removed = getattr(self.builder, remove_name)(declaration.descriptor.name) + if removed is None: + raise KeyError(f"Spawn asset is absent from SceneBuilder: {uid!r}.") + del self._assets[uid] + + def commit(self) -> Any: + """Finalize once or let ``SpawnResult`` consume pending changes.""" + if self.result is None: + self.result = self.builder.finalize() + elif self.builder.has_pending_changes or self.result.needs_rebuild: + self.result = self.result.rebuild(self.builder) + return self.result + + def bind(self) -> None: + """Resolve current Spawn handles and bind every declared facade.""" + if self.result is None: + raise RuntimeError("Spawn scene must be materialized before binding.") + + for declaration in self._assets.values(): + if declaration.on_bind is None: + continue + paths = self._paths(declaration) + handles = tuple(self.result.handles[path] for path in paths) + declaration.on_bind(self.result, handles) + + def close(self) -> None: + """Release Spawn resources and facade callback references.""" + if self.result is not None: + self.result.close() + self.result = None + self._assets.clear() + + def _paths(self, declaration: _AssetDeclaration) -> tuple[str, ...]: + if declaration.kind == "rigid_object_group": + return tuple( + f"{arena}/{member.name}" + for arena in self.arena_names + for member in declaration.descriptor + ) + name = declaration.descriptor.name + if not declaration.descriptor.per_env: + return (name,) + return tuple(f"{arena}/{name}" for arena in self.arena_names) diff --git a/embodichain/lab/sim/spawn/usd.py b/embodichain/lab/sim/spawn/usd.py new file mode 100644 index 000000000..5422b17cf --- /dev/null +++ b/embodichain/lab/sim/spawn/usd.py @@ -0,0 +1,166 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Compatibility translation for EmbodiChain's singleton USD APIs.""" + +from __future__ import annotations + +import os +from dataclasses import replace + +from dexsim.spawn import ( + ArticulationDesc, + MaterialDesc, + NewtonJointDesc, + ObjectDesc, + RenderDesc, +) +from dexsim.types import ActorType + +from embodichain.lab.sim.cfg import ArticulationCfg, RigidObjectCfg +from embodichain.lab.sim.spawn.descriptors import ( + _compile_dexsim_collision, + _compile_newton_collision, + _compile_rigid_physics, + _compile_visual_material, + _pose_from_cfg, + _required_uid, + _vector3, +) + +__all__ = ["articulation_desc_from_usd", "rigid_desc_from_usd"] + + +def rigid_desc_from_usd( + cfg: RigidObjectCfg, + *, + per_env: bool = True, +) -> tuple[ObjectDesc, dict[str, MaterialDesc]]: + """Select the sole rigid object in a USD stage.""" + uid = _required_uid(cfg.uid, "Rigid object") + path = getattr(cfg.shape, "fpath", None) + scene, desc = _parse_singleton(path, "mesh_objects", "rigid object") + + desc.name = uid + desc.pose = _pose_from_cfg(cfg) + desc.per_env = per_env + materials = _namespace_materials(desc.renders, scene.materials, uid) + + if cfg.use_usd_properties: + if desc.physics is None: + raise ValueError(f"USD rigid object {path!r} has no physics.") + cfg.body_type = { + ActorType.DYNAMIC: "dynamic", + ActorType.KINEMATIC: "kinematic", + ActorType.STATIC: "static", + }[desc.physics.actor_type] + cfg.body_scale = tuple(float(value) for value in desc.body_scale) + return desc, materials + + desc.physics = _compile_rigid_physics(cfg.attrs, cfg.body_type) + desc.body_scale = _vector3(cfg.body_scale, field_name="body_scale") + for collision in desc.collisions: + collision.enable_collision = bool(cfg.attrs.enable_collision) + collision.dexsim = _compile_dexsim_collision(cfg.attrs) + collision.newton = _compile_newton_collision(cfg.attrs) + + material_ref, material_entry = _compile_visual_material( + uid, + cfg.shape.visual_material, + ) + if material_entry is not None: + materials = {material_entry[0]: material_entry[1]} + for render in desc.renders: + render.material = None + render.material_ref = material_ref + return desc, materials + + +def articulation_desc_from_usd( + cfg: ArticulationCfg, + *, + per_env: bool = True, + source_path: str | None = None, +) -> tuple[ArticulationDesc, dict[str, MaterialDesc]]: + """Select the sole articulation in a USD stage.""" + path = source_path or cfg.fpath + scene, desc = _parse_singleton(path, "articulations", "articulation") + uid = _required_uid( + cfg.uid or os.path.splitext(os.path.basename(str(path)))[0], + "Articulation", + ) + cfg.uid = uid + desc.name = uid + desc.pose = _pose_from_cfg(cfg) + desc.per_env = per_env + renders = [visual for link in desc.links for visual in link.visuals] + materials = _namespace_materials(renders, scene.materials, uid) + + if cfg.use_usd_properties: + cfg.fix_base = bool(desc.fixed_base) + cfg.disable_self_collision = not desc.enable_self_collision + cfg.body_scale = tuple(float(value) for value in desc.body_scale) + else: + desc.fixed_base = bool(cfg.fix_base) + desc.enable_self_collision = not bool(cfg.disable_self_collision) + desc.body_scale = _vector3(cfg.body_scale, field_name="body_scale") + target_mode = {"force": 3, "none": 0}.get(cfg.drive_pros.drive_type) + if target_mode is not None: + for joint in desc.joints: + joint.newton = ( + NewtonJointDesc(target_mode=target_mode) + if joint.newton is None + else replace(joint.newton, target_mode=target_mode) + ) + return desc, materials + + +def _parse_singleton(path: object, collection: str, label: str): + if path is None: + raise ValueError(f"A USD path is required for the {label}.") + + from dexsim.kit.usd import parse_usd + + scene = parse_usd(str(path)) + candidates = getattr(scene, collection) + if len(candidates) != 1: + found = [ + (item.name, None if item.usd is None else item.usd.prim_path) + for item in candidates + ] + raise ValueError( + f"Expected exactly one {label} in USD file {path!r}, found " + f"{len(candidates)}: {found}." + ) + return scene, candidates[0] + + +def _namespace_materials( + renders: list[RenderDesc], + materials: dict[str, MaterialDesc], + uid: str, +) -> dict[str, MaterialDesc]: + selected = {} + for render in renders: + if render.material_ref is None: + continue + source_ref = render.material_ref + material = materials[source_ref] + render.material_ref = f"{uid}::{source_ref}" + selected[render.material_ref] = replace( + material, + name=f"{uid}::{material.name}", + ) + return selected diff --git a/examples/sim/demo/grasp_cup_to_caffe.py b/examples/sim/demo/grasp_cup_to_caffe.py index 803214723..9c46bad6d 100644 --- a/examples/sim/demo/grasp_cup_to_caffe.py +++ b/examples/sim/demo/grasp_cup_to_caffe.py @@ -428,6 +428,7 @@ def main(): caffe = create_caffe(sim) cup = create_cup(sim) + sim.prepare() sim.update(step=1) # apply random perturbation diff --git a/examples/sim/demo/pick_up_cloth.py b/examples/sim/demo/pick_up_cloth.py index b049dbcb2..f4e60a8c8 100644 --- a/examples/sim/demo/pick_up_cloth.py +++ b/examples/sim/demo/pick_up_cloth.py @@ -272,7 +272,7 @@ def main(): robot = create_robot(sim) cloth = create_cloth(sim) padding_box = create_padding_box(sim) - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() sim.update(step=10) # Let the cloth settle before interaction diff --git a/examples/sim/demo/press_softbody.py b/examples/sim/demo/press_softbody.py index 214ca4b23..017235276 100644 --- a/examples/sim/demo/press_softbody.py +++ b/examples/sim/demo/press_softbody.py @@ -190,7 +190,7 @@ def main(): robot = create_robot(sim) soft_cow = create_soft_cow(sim) - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() diff --git a/examples/sim/demo/scoop_ice.py b/examples/sim/demo/scoop_ice.py index 1cca5c58c..4a524f056 100644 --- a/examples/sim/demo/scoop_ice.py +++ b/examples/sim/demo/scoop_ice.py @@ -309,6 +309,7 @@ def create_ice_cubes(sim: SimulationManager): material_type="BSDF", ) ) + sim.prepare() ice_cubes.set_visual_material(mat=ice_mat) return ice_cubes diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index 832f818b0..a690c7189 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -103,6 +103,7 @@ def main(): # Add camera to simulation camera = sim.add_sensor(sensor_cfg=camera_cfg) + sim.prepare() # Wait for initialization time.sleep(0.2) diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index 8fefc7ceb..600a61c5e 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -93,6 +93,7 @@ def main(): init_pos=[0.3, 0.0, 1.0], ) ) + sim.prepare() native_window_opened = False if not args.headless: @@ -128,9 +129,6 @@ def main(): def run_simulation(sim: SimulationManager): """Run the simulation loop.""" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 gizmo_enabled = True try: diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 2750a8a80..604b5c001 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -104,6 +104,7 @@ def main(): init_qpos=[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0, 0.0, 0.0], ) robot = sim.add_robot(cfg=robot_cfg) + sim.prepare() # Set initial joint positions initial_qpos = torch.tensor( diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index a2cca4a48..fb1943553 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -126,12 +126,6 @@ def main(): device="cpu", ) - left_joint_ids = robot.get_joint_ids("left_arm") - right_joint_ids = robot.get_joint_ids("right_arm") - - robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) - robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) - # Create a rigid object (cube) positioned to the side of the robot cube_cfg = RigidObjectCfg( uid="interactive_cube", @@ -163,6 +157,12 @@ def main(): ), ) camera = sim.add_sensor(sensor_cfg=camera_cfg) + sim.prepare() + + left_joint_ids = robot.get_joint_ids("left_arm") + right_joint_ids = robot.get_joint_ids("right_arm") + robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) + robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) native_window_opened = False if not args.headless: diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py index 2f830a8bc..b0f4ef55c 100644 --- a/examples/sim/gizmo/gizmo_w1.py +++ b/examples/sim/gizmo/gizmo_w1.py @@ -130,6 +130,7 @@ def main(): 0.0000e00, ] robot = sim.add_robot(cfg=cfg) + sim.prepare() # Set initial joint positions for both arms # Left arm: 8 joints (WAIST + 7 LEFT_J), Right arm: 8 joints (WAIST + 7 RIGHT_J) diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py index 908f8619a..3bb800d02 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -457,11 +457,6 @@ def _build_scene( if robot is None: raise RuntimeError(f"Failed to add robot '{robot_type}' to the cuRobo demo.") target_xpos = _resolve_batched_target(target_xpos, robot.num_instances) - if robot_type == "w1": - # Keep the W1-specific IK diagnostic batched so it remains useful when - # checking solver and cuRobo reachability across multiple environments. - is_success, ik_qpos = robot.compute_ik(pose=target_xpos, name=control_part) - print(f"robot compute ik success: {is_success}, ik_qpos: {ik_qpos}") # This object is also exported into the cuRobo collision world below via # CuroboWorldCfg.rigid_objects, so the simulator and planner share geometry @@ -476,6 +471,13 @@ def _build_scene( init_rot=(0.0, 0.0, 0.0), ) ) + sim.prepare() + + if robot_type == "w1": + # Keep the W1-specific IK diagnostic batched so it remains useful when + # checking solver and cuRobo reachability across multiple environments. + is_success, ik_qpos = robot.compute_ik(pose=target_xpos, name=control_part) + print(f"robot compute ik success: {is_success}, ik_qpos: {ik_qpos}") return sim, robot, demo_block, target_xpos, control_part @@ -698,8 +700,6 @@ def main() -> None: effective_gpu_id, visualization_cfg_from_args(args), ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() obstacles = [demo_block] obstacle_poses = _perturb_obstacles( diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/planners/neural_planner.py index 115282753..d234f001f 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/planners/neural_planner.py @@ -221,8 +221,7 @@ def main() -> None: arm_name = "arm" device = robot.device - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() diff --git a/examples/sim/robot/dexforce_w1.py b/examples/sim/robot/dexforce_w1.py index 9a4e78383..51b0afa2d 100644 --- a/examples/sim/robot/dexforce_w1.py +++ b/examples/sim/robot/dexforce_w1.py @@ -70,6 +70,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: ) robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.update(step=1) print("DexforceW1 with a user defined end-effector added to the simulation.") diff --git a/examples/sim/scene/scene_demo.py b/examples/sim/scene/scene_demo.py index 45866ee31..68646b590 100644 --- a/examples/sim/scene/scene_demo.py +++ b/examples/sim/scene/scene_demo.py @@ -78,9 +78,6 @@ def resolve_asset_path(scene_name: str) -> str: def run_simulation(sim: SimulationManager): """Run the simulation loop.""" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - try: while True: time.sleep(0.01) @@ -181,6 +178,8 @@ def main(): logger.log_info(f"Failed to load scene asset: {e}") return + sim.prepare() + logger.log_info(f"Scene '{args.scene}' setup complete!") logger.log_info(f"Running simulation with {args.num_envs} environment(s)") logger.log_info("Press Ctrl+C to stop the simulation") diff --git a/examples/sim/sensors/batch_camera.py b/examples/sim/sensors/batch_camera.py index 97c606adf..0af567c7b 100644 --- a/examples/sim/sensors/batch_camera.py +++ b/examples/sim/sensors/batch_camera.py @@ -60,8 +60,7 @@ def main(args): ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() @@ -123,6 +122,8 @@ def main(args): else: plt.show() + sim.destroy() + if __name__ == "__main__": import argparse diff --git a/examples/sim/sensors/create_contact_sensor.py b/examples/sim/sensors/create_contact_sensor.py index ebcf0b94c..e918e81dc 100644 --- a/examples/sim/sensors/create_contact_sensor.py +++ b/examples/sim/sensors/create_contact_sensor.py @@ -209,6 +209,7 @@ def main(): cube1 = create_cube(sim, "cube1", position=[0.0, 0.0, 0.06]) cube2 = create_cube(sim, "cube2", position=[0.0, 0.0, 0.09]) robot = create_robot(sim, "UR10_PGI", position=[0.5, 0.0, 0.0]) + sim.prepare() print("[INFO]: Scene setup complete!") print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") @@ -230,10 +231,6 @@ def run_simulation(sim: SimulationManager): sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 # contact filter config contact_filter_cfg = ContactSensorCfg() diff --git a/examples/sim/solvers/differential_solver.py b/examples/sim/solvers/differential_solver.py index ec6424844..111cd4c53 100644 --- a/examples/sim/solvers/differential_solver.py +++ b/examples/sim/solvers/differential_solver.py @@ -82,6 +82,7 @@ def main( } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() # Prepare initial joint positions for all environments rad = torch.deg2rad(torch.tensor(45.0)) diff --git a/examples/sim/solvers/neural_ik_solver.py b/examples/sim/solvers/neural_ik_solver.py index 5df974cdb..2fdd6ae43 100644 --- a/examples/sim/solvers/neural_ik_solver.py +++ b/examples/sim/solvers/neural_ik_solver.py @@ -128,6 +128,7 @@ def main() -> None: ) robot: Robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.open_window() diff --git a/examples/sim/solvers/opw_solver.py b/examples/sim/solvers/opw_solver.py index 5890a55e4..56ae124eb 100644 --- a/examples/sim/solvers/opw_solver.py +++ b/examples/sim/solvers/opw_solver.py @@ -89,6 +89,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: # Add robot to simulation robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + sim.prepare() # Left arm control arm_name = "left_arm" diff --git a/examples/sim/solvers/pink_solver.py b/examples/sim/solvers/pink_solver.py index 33a65cfbe..9d0e71b4e 100644 --- a/examples/sim/solvers/pink_solver.py +++ b/examples/sim/solvers/pink_solver.py @@ -76,6 +76,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() # Define a sample target pose as a 1x4x4 homogeneous matrix rad = torch.deg2rad(torch.tensor(45.0)) diff --git a/examples/sim/solvers/pinocchio_solver.py b/examples/sim/solvers/pinocchio_solver.py index fb43138dd..bfc3610a9 100644 --- a/examples/sim/solvers/pinocchio_solver.py +++ b/examples/sim/solvers/pinocchio_solver.py @@ -76,6 +76,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() arm_name = "left_arm" # Set initial joint positions for left arm qpos_seed = torch.tensor( diff --git a/examples/sim/solvers/pytorch_solver.py b/examples/sim/solvers/pytorch_solver.py index bef9750e1..46749573a 100644 --- a/examples/sim/solvers/pytorch_solver.py +++ b/examples/sim/solvers/pytorch_solver.py @@ -82,6 +82,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: # Add robot to simulation robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() # Prepare initial joint positions for all environments arm_name = "left_arm" diff --git a/examples/sim/solvers/srs_solver.py b/examples/sim/solvers/srs_solver.py index ecb6142d8..76693f96c 100644 --- a/examples/sim/solvers/srs_solver.py +++ b/examples/sim/solvers/srs_solver.py @@ -53,6 +53,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: sim.set_manual_update(False) robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) + sim.prepare() arm_name = "left_arm" # Set initial joint positions for left arm qpos_fk_list = [ diff --git a/examples/sim/workspace/analyze_cartesian_workspace.py b/examples/sim/workspace/analyze_cartesian_workspace.py index fb9160067..d514c71e2 100644 --- a/examples/sim/workspace/analyze_cartesian_workspace.py +++ b/examples/sim/workspace/analyze_cartesian_workspace.py @@ -101,6 +101,7 @@ def main() -> None: } ) robot = sim.add_robot(cfg=cfg) + sim.prepare() print("DexforceW1 robot added to the simulation.") left_qpos = torch.tensor( diff --git a/examples/sim/workspace/analyze_joint_workspace.py b/examples/sim/workspace/analyze_joint_workspace.py index 3695bdb79..ba96f7ca2 100644 --- a/examples/sim/workspace/analyze_joint_workspace.py +++ b/examples/sim/workspace/analyze_joint_workspace.py @@ -98,6 +98,7 @@ def main() -> None: } ) robot = sim_manager.add_robot(cfg=cfg) + sim_manager.prepare() print("DexforceW1 robot added to the simulation.") analyzer = WorkspaceAnalyzer( diff --git a/examples/sim/workspace/analyze_plane_workspace.py b/examples/sim/workspace/analyze_plane_workspace.py index 95e381e1e..7fbcccb24 100644 --- a/examples/sim/workspace/analyze_plane_workspace.py +++ b/examples/sim/workspace/analyze_plane_workspace.py @@ -101,6 +101,7 @@ def main() -> None: } ) robot = sim.add_robot(cfg=cfg) + sim.prepare() print("DexforceW1 robot added to the simulation.") left_qpos = torch.tensor( diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index c7429e62a..c1b747e9f 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -242,6 +242,7 @@ def run_assemble_demo( create_support_surface(sim) can = create_assemble_object(sim) cube = create_base_object(sim) + sim.prepare() settle_object(sim, can, step=0) clone_local_pose_from_first_env(can) diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index 2b52be831..442a22292 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -241,6 +241,7 @@ def create_pickment_object( body_scale=preset.body_scale, ) ) + sim.prepare() obj.cfg.init_pos = compute_supported_init_pos(obj, preset) obj.reset() return obj diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index f24f8aba9..2d116a85f 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -535,6 +535,7 @@ def run_coordinated_placement_demo( create_table(sim) bread = create_bread(sim) pan = create_pan(sim) + sim.prepare() settle_object(sim, bread, step=0) settle_object(sim, pan, step=0) bread_pose_batch = clone_local_pose_from_first_env(bread) diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index 446596c86..73ef050ec 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -180,6 +180,7 @@ def run_handover_demo( """Plan and optionally execute a pick-up followed by a handover.""" create_support_surface(sim) obj = create_handover_object(sim) + sim.prepare() settle_object(sim, obj, step=0) clone_local_pose_from_first_env(obj) obj.clear_dynamics() diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index 797e157c4..12ff42096 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -93,6 +93,7 @@ def create_pick_object(sim) -> RigidObject: body_scale=(0.75, 0.75, 1.0), ) ) + sim.prepare() sim.update(step=10) clone_local_pose_from_first_env(obj) obj.clear_dynamics() diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index b5450f6bc..05664573b 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -95,6 +95,7 @@ def create_pick_object(sim) -> RigidObject: init_pos=[*OBJECT_XY, OBJECT_SIZE[2]], ) ) + sim.prepare() sim.update(step=10) clone_local_pose_from_first_env(obj) obj.clear_dynamics() diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index 17a1bad8d..0c3ea02f1 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -93,6 +93,7 @@ def create_pick_object(sim) -> RigidObject: init_pos=[*OBJECT_XY, 0.5 * OBJECT_SIZE[2]], ) ) + sim.prepare() sim.update(step=10) clone_local_pose_from_first_env(obj) obj.clear_dynamics() diff --git a/scripts/tutorials/atomic_action/scenario_utils.py b/scripts/tutorials/atomic_action/scenario_utils.py index 4fb0f8fda..a5d98ff1c 100644 --- a/scripts/tutorials/atomic_action/scenario_utils.py +++ b/scripts/tutorials/atomic_action/scenario_utils.py @@ -426,8 +426,6 @@ def add_support_surface( def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: """Reset, settle, and freeze an object before tutorial planning.""" - if sim.device.type == "cuda": - sim.init_gpu_physics() obj.reset() if step > 0: sim.update(step=step) diff --git a/scripts/tutorials/grasp/grasp_generator.py b/scripts/tutorials/grasp/grasp_generator.py index 328835311..b8f755af1 100644 --- a/scripts/tutorials/grasp/grasp_generator.py +++ b/scripts/tutorials/grasp/grasp_generator.py @@ -226,6 +226,7 @@ def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tenso sim = initialize_simulation(args) robot = create_robot(sim, position=[0.0, 0.0, 0.0]) obj = create_obj(sim) + sim.prepare() # get mug grasp pose grasp_cfg = GraspGeneratorCfg( diff --git a/scripts/tutorials/sim/create_cloth.py b/scripts/tutorials/sim/create_cloth.py index 202b5fd02..1e0639fb9 100644 --- a/scripts/tutorials/sim/create_cloth.py +++ b/scripts/tutorials/sim/create_cloth.py @@ -123,7 +123,7 @@ def main(): mass=0.01, youngs=1e9, poissons=0.4, - thickness=0.04, + thickness=0.004, bending_stiffness=0.01, bending_damping=0.1, dynamic_friction=0.95, @@ -151,6 +151,8 @@ def main(): padding_box = sim.add_rigid_object(cfg=padding_box_cfg) print("[INFO]: Add soft object complete!") + sim.prepare() + # Open window when the scene has been set up if not args.headless: sim.open_window() @@ -170,9 +172,6 @@ def run_simulation(sim: SimulationManager, cloth: ClothObject) -> None: soft_obj: soft object """ - # Initialize GPU physics - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py index b9517c241..682b2c816 100644 --- a/scripts/tutorials/sim/create_rigid_constraint.py +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -101,8 +101,7 @@ def main(): ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() print("[INFO]: Scene setup complete with two cubes (cube_a, cube_b).") diff --git a/scripts/tutorials/sim/create_rigid_object_group.py b/scripts/tutorials/sim/create_rigid_object_group.py index 7399d872c..d6aa22b75 100644 --- a/scripts/tutorials/sim/create_rigid_object_group.py +++ b/scripts/tutorials/sim/create_rigid_object_group.py @@ -107,6 +107,7 @@ def main(): print("[INFO]: Press Ctrl+C to stop the simulation") # Open window when the scene has been set up + sim.prepare() if not args.headless: sim.open_window() @@ -121,10 +122,6 @@ def run_simulation(sim: SimulationManager): sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/create_robot.py b/scripts/tutorials/sim/create_robot.py index e393c7b05..07ef35405 100644 --- a/scripts/tutorials/sim/create_robot.py +++ b/scripts/tutorials/sim/create_robot.py @@ -74,9 +74,9 @@ def main(): # Create robot configuration robot = create_robot(sim) - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + # Materialize the declared scene before accessing robot metadata. + sim.prepare() + print(f"Robot created successfully with {robot.dof} joints") # Open visualization window if not headless if not args.headless: @@ -138,8 +138,6 @@ def create_robot(sim): # Add robot to simulation robot: Robot = sim.add_robot(cfg=cfg) - print(f"Robot created successfully with {robot.dof} joints") - return robot diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index 89e04dd1d..1188e25cc 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -91,7 +91,6 @@ def main() -> None: uid="cube", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="dynamic", - body_scale=[0.5, 0.5, 0.5], attrs=RigidBodyAttributesCfg( mass=0.1, dynamic_friction=0.5, @@ -119,6 +118,9 @@ def main() -> None: ) ) + # Materialize the complete initial scene before exposing it to the viewer. + sim.prepare() + print("[INFO]: Scene setup complete!") print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") print("[INFO]: Press Ctrl+C to stop the simulation") @@ -157,10 +159,6 @@ def run_simulation( max_steps: Optional maximum number of simulation steps to execute. """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/create_sensor.py b/scripts/tutorials/sim/create_sensor.py index a09da16e4..7231f5af8 100644 --- a/scripts/tutorials/sim/create_sensor.py +++ b/scripts/tutorials/sim/create_sensor.py @@ -112,8 +112,6 @@ def main() -> None: # Create robot configuration robot = create_robot(sim) - sensor = create_sensor(sim, args) - # Add a cube to the scene cube_cfg = RigidObjectCfg( uid="cube", @@ -123,9 +121,12 @@ def main() -> None: ) sim.add_rigid_object(cfg=cube_cfg) - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + # Materialize all physical assets before reading robot metadata or + # constructing render-only sensors. + sim.prepare() + print(f"Robot created successfully with {robot.dof} joints") + + sensor = create_sensor(sim, args) # Open visualization window if not headless if not args.headless: @@ -234,8 +235,6 @@ def create_robot(sim): # Add robot to simulation robot: Robot = sim.add_robot(cfg=cfg) - print(f"Robot created successfully with {robot.dof} joints") - return robot diff --git a/scripts/tutorials/sim/create_softbody.py b/scripts/tutorials/sim/create_softbody.py index 38046f397..aab5b4112 100644 --- a/scripts/tutorials/sim/create_softbody.py +++ b/scripts/tutorials/sim/create_softbody.py @@ -93,6 +93,8 @@ def main(): ) print("[INFO]: Add soft object complete!") + sim.prepare() + # Open window when the scene has been set up if not args.headless: sim.open_window() @@ -112,9 +114,6 @@ def run_simulation(sim: SimulationManager, soft_obj: SoftObject) -> None: soft_obj: soft object """ - # Initialize GPU physics - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/export_usd.py b/scripts/tutorials/sim/export_usd.py index cf402baab..98b8fc721 100644 --- a/scripts/tutorials/sim/export_usd.py +++ b/scripts/tutorials/sim/export_usd.py @@ -263,7 +263,9 @@ def main(): caffe = create_caffe(sim) cup = create_cup(sim) - sim.export_usd("w1_coffee_scene.usda") + sim.prepare() + + sim.export_usd("w1_coffee_scene.usd") logger.log_info("Scene exported successfully.") diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index c2171897a..9f850de3d 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -99,6 +99,8 @@ def main(): dtype=torch.float32, device="cpu", ) + + sim.prepare() joint_ids = robot.get_joint_ids("arm") robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) diff --git a/scripts/tutorials/sim/import_usd.py b/scripts/tutorials/sim/import_usd.py index 02d5cc089..abf4859a0 100644 --- a/scripts/tutorials/sim/import_usd.py +++ b/scripts/tutorials/sim/import_usd.py @@ -113,6 +113,7 @@ def main(): ) # Open window when the scene has been set up + sim.prepare() if not args.headless: sim.open_window() @@ -130,10 +131,6 @@ def run_simulation(sim: SimulationManager): sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index fb4b3169f..351b00009 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -238,8 +238,7 @@ def main() -> None: robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict({"uid": "CobotMagic"})) arm_name = "left_arm" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 606c64a8d..394e25184 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -53,6 +53,9 @@ def main(visualization: VisualizationCfg | None = None) -> None: sim.set_manual_update(False) robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) + + sim.prepare() + arm_name = "left_arm" # Set initial joint positions for left arm qpos_fk_list = [ diff --git a/scripts/tutorials/visualization/viser_scene.py b/scripts/tutorials/visualization/viser_scene.py index a3d329932..9350391fc 100644 --- a/scripts/tutorials/visualization/viser_scene.py +++ b/scripts/tutorials/visualization/viser_scene.py @@ -159,8 +159,7 @@ def main() -> None: build_pk_chain=False, ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() visualization_cfg = VisualizationCfg( backend="viser", diff --git a/tests/sim/test_newton_finalize_lifecycle.py b/tests/sim/test_newton_finalize_lifecycle.py deleted file mode 100644 index 3b2adefd8..000000000 --- a/tests/sim/test_newton_finalize_lifecycle.py +++ /dev/null @@ -1,198 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- -"""Unit tests for the Newton physics backend finalize/invalidate lifecycle. - -These tests exercise :class:`NewtonPhysicsBackend` in isolation (no GPU and no -live dexsim world required) by injecting a fake Newton manager and patching the -``ensure_simulation_prepared_lazy`` rebuild entry point. They verify the -backend owns the dirty/finalize state machine that used to live inline in -:class:`SimulationManager`. -""" - -from __future__ import annotations - -from types import SimpleNamespace -from unittest.mock import patch - -from embodichain.lab.sim.physics import NewtonPhysicsBackend - - -class _Resettable: - """Stand-in for a RigidObject/Articulation with a reset() call counter.""" - - def __init__(self) -> None: - self.reset_calls = 0 - - def reset(self) -> None: - self.reset_calls += 1 - - -class _FakeNewtonManager: - """Stand-in for dexsim's NewtonManager exposing only the lifecycle state.""" - - def __init__(self) -> None: - self.lifecycle_state = SimpleNamespace(name="BUILDER") - - -def _make_backend() -> tuple[ - NewtonPhysicsBackend, - _FakeNewtonManager, - _Resettable, - _Resettable, - _Resettable, - _Resettable, -]: - rigid_obj = _Resettable() - rigid_group = _Resettable() # groups must NOT be reset by the Newton backend. - articulation = _Resettable() - robot = _Resettable() # a robot is an articulation and is reset like one. - newton_mgr = _FakeNewtonManager() - - # Minimal owning-SimulationManager stand-in: only the attributes the backend - # touches during finalize / reset are needed. - manager = SimpleNamespace( - _world=object(), - _rigid_objects={"rigid": rigid_obj}, - _rigid_object_groups={"rigid_group": rigid_group}, - _articulations={"art": articulation}, - _robots={"robot": robot}, - ) - - backend = NewtonPhysicsBackend(manager) - # Inject the fake manager so finalize() does not call get_newton_manager. - backend._newton_manager = newton_mgr - return backend, newton_mgr, rigid_obj, rigid_group, articulation, robot - - -def _fake_ensure_prepared_lazy(mgr, world, *, rebuild_from_scene, warn): - """Mimic the real rebuild: bring the Newton model to the READY state.""" - mgr.lifecycle_state.name = "READY" - return True, None - - -@patch( - "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", - new=_fake_ensure_prepared_lazy, -) -def test_finalize_resets_entities_after_ready() -> None: - ( - backend, - newton_mgr, - rigid_obj, - rigid_group, - articulation, - robot, - ) = _make_backend() - - assert not backend.is_initialized - backend.prepare() - - assert newton_mgr.lifecycle_state.name == "READY" - assert backend.is_initialized - assert rigid_obj.reset_calls == 1 - assert articulation.reset_calls == 1 - assert robot.reset_calls == 1 - # Rigid object groups are not supported on the Newton backend: not reset. - assert rigid_group.reset_calls == 0 - - -@patch( - "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", - new=_fake_ensure_prepared_lazy, -) -def test_finalize_does_not_repeat_deferred_reset() -> None: - ( - backend, - _newton_mgr, - rigid_obj, - _rigid_group, - articulation, - robot, - ) = _make_backend() - - backend.prepare() - backend.prepare() - - assert rigid_obj.reset_calls == 1 - assert articulation.reset_calls == 1 - assert robot.reset_calls == 1 - - -@patch( - "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", - new=_fake_ensure_prepared_lazy, -) -def test_invalidation_allows_next_finalize_to_reset_again() -> None: - ( - backend, - _newton_mgr, - rigid_obj, - _rigid_group, - articulation, - robot, - ) = _make_backend() - - backend.prepare() - backend.invalidate() - assert not backend.is_initialized - backend.prepare() - - assert rigid_obj.reset_calls == 2 - assert articulation.reset_calls == 2 - assert robot.reset_calls == 2 - - -@patch( - "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", - new=_fake_ensure_prepared_lazy, -) -def test_finalize_raises_when_rebuild_unsafe() -> None: - backend, _newton_mgr, rigid_obj, _rigid_group, _articulation, _robot = ( - _make_backend() - ) - - # An unsafe rebuild makes finalize() raise (logger.log_error raises by - # default). It must not mark itself initialized nor reset entities. - with patch( - "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", - new=lambda mgr, world, *, rebuild_from_scene, warn: (False, None), - ): - try: - backend.prepare() - except RuntimeError: - pass - else: # pragma: no cover - defensive - raise AssertionError("finalize() should raise on an unsafe rebuild") - - assert not backend.is_initialized - assert rigid_obj.reset_calls == 0 - - -def test_invalidate_is_idempotent_and_only_clears_finalized_flag() -> None: - ( - backend, - _newton_mgr, - _rigid_obj, - _rigid_group, - _articulation, - _robot, - ) = _make_backend() - backend._is_finalized = True - - backend.invalidate() - backend.invalidate() - - assert not backend.is_initialized diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 699acb0e2..afa692f98 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -422,8 +422,9 @@ def test_start_visualization_rejects_open_native_window() -> None: sim.start_visualization() -def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> None: +def test_constructor_only_declares_spawn_scene(monkeypatch) -> None: lifecycle: list[str] = [] + spawn_scene = MagicMock() world = MagicMock() world.get_physics_scene.return_value = MagicMock() world.get_env.return_value = MagicMock() @@ -433,6 +434,11 @@ def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> No ) monkeypatch.setattr(sim_manager_module.wp, "init", lambda: None) monkeypatch.setattr(sim_manager_module.dexsim, "World", lambda _cfg: world) + monkeypatch.setattr( + sim_manager_module, + "SpawnScene", + lambda *_args, **_kwargs: spawn_scene, + ) monkeypatch.setattr( sim_manager_module.dexsim, "set_physics_config", lambda **_kwargs: None ) @@ -454,7 +460,7 @@ def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> No ) monkeypatch.setattr( SimulationManager, - "_create_default_plane", + "_declare_spawn_default_plane", lambda _self: lifecycle.append("plane"), ) monkeypatch.setattr( @@ -468,14 +474,9 @@ def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> No lambda _self: lifecycle.append("lighting"), ) - def build_arenas(sim: SimulationManager, num: int) -> None: - lifecycle.append("arenas") - sim._arenas.extend([object() for _ in range(num)]) - def start_visualization(sim: SimulationManager) -> None: lifecycle.append(f"visualization:{sim.num_envs}") - monkeypatch.setattr(SimulationManager, "_build_multiple_arenas", build_arenas) monkeypatch.setattr( SimulationManager, "start_visualization", @@ -487,22 +488,33 @@ def start_visualization(sim: SimulationManager) -> None: assert lifecycle == [ "resources", - "plane", "background", + "plane", "lighting", - "arenas", - "visualization:3", ] + assert sim._spawn_scene is spawn_scene + assert sim._arenas == [] def test_remove_asset_marks_visualization_topology_dirty() -> None: sim, runtime = _make_visualization_sim_manager() rigid_object = MagicMock() + spawn_scene = MagicMock() + spawn_scene.__contains__.return_value = True + spawn_scene.result = object() + sim._spawn_scene = spawn_scene + sim.prepare = MagicMock() sim._rigid_objects = {"cube": rigid_object} + sim._articulations = {} + sim._robots = {} + sim._lights = {} assert sim.remove_asset("cube") - rigid_object.destroy.assert_called_once_with() + spawn_scene.remove.assert_called_once_with("cube") + sim.prepare.assert_called_once_with() + rigid_object.destroy.assert_not_called() + assert "cube" not in sim._rigid_objects assert sim._visualization_topology_revision == 3 sim.stop_visualization() assert runtime.stopped diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py index 5236b311e..6f68f0984 100644 --- a/tests/sim/test_sim_manager_cfg.py +++ b/tests/sim/test_sim_manager_cfg.py @@ -17,6 +17,7 @@ from __future__ import annotations import torch +import pytest from embodichain.lab.sim import SimulationManagerCfg from embodichain.lab.sim.cfg import NewtonPhysicsCfg, WindowCameraPoseCfg @@ -62,6 +63,14 @@ def test_simulation_manager_cfg_initializes_window_camera_pose() -> None: assert cfg.window_camera_pose == window_camera_pose +def test_simulation_manager_cfg_has_no_scene_construction_switch() -> None: + cfg = SimulationManagerCfg() + + assert "scene_construction" not in cfg.to_dict() + with pytest.raises(TypeError, match="scene_construction"): + SimulationManagerCfg(scene_construction="legacy") + + def test_newton_physics_cfg_uses_device() -> None: cfg = NewtonPhysicsCfg(device="cuda:1")