diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 05cc2434..0577c1c0 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -75,7 +75,7 @@ jobs: - name: Build docs shell: bash run: | - pip install -e . --extra-index-url http://pyp.open3dv.site:2345/simple/ --trusted-host pyp.open3dv.site + pip install -e . --extra-index-url http://pyp.open3dv.site:2345/simple/ --trusted-host pyp.open3dv.site pip install -r docs/requirements.txt python3 docs/scripts/sync_readme.py cd ${GITHUB_WORKSPACE}/docs diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 1a0fa89e..0fac4f88 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -129,9 +129,7 @@ def __init__( self._setup_scene(**kwargs) - # TODO: To be removed. - if self.device.type == "cuda": - self.sim.init_gpu_physics() + self.sim.prepare_physics() if not self.sim_cfg.headless: self.sim.open_window() diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 0b10a725..84bc93fb 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -77,7 +77,7 @@ def to_dexsim_flags(self): @configclass -class PhysicsCfg: +class DefaultPhysicsCfg: gravity: np.ndarray = field(default_factory=lambda: np.array([0, 0, -9.81])) """Gravity vector for the simulation environment.""" @@ -124,6 +124,95 @@ def to_dexsim_args(self) -> Dict[str, Any]: return args +# Backwards-compatible alias for existing task configs. +PhysicsCfg = DefaultPhysicsCfg + + +@configclass +class NewtonPhysicsCfg: + """Configuration for DexSim Newton physics backend.""" + + num_substeps: int = 10 + """Number of Newton solver substeps per EmbodiChain physics step.""" + + device: str | None = None + """Newton device. If None, derived from ``SimulationManagerCfg.sim_device`` and ``gpu_id``.""" + + require_grad: bool = False + """Whether to finalize the Newton model for differentiable simulation.""" + + use_cuda_graph: bool = True + """Whether to use CUDA graph capture for Newton stepping when supported.""" + + debug_mode: bool = False + """Whether to enable Newton debug mode.""" + + solver_type: Literal["mjwarp", "xpbd", "semi_implicit", "featherstone", "vbd"] = ( + "semi_implicit" + ) + """Newton solver preset.""" + + broad_phase: Literal["nxn", "sap", "explicit"] | None = None + """Newton collision broad-phase implementation. If None, DexSim chooses its default.""" + + visualizer_enabled: bool = False + """Whether to enable the Newton visualizer.""" + + def to_dexsim_cfg( + self, + physics_dt: float, + sim_device: str | torch.device, + gpu_id: int, + ): + """Convert this config to ``dexsim.engine.newton_physics.NewtonCfg``.""" + from dexsim.engine.newton_physics import ( + FeatherstoneSolverCfg, + MJWarpSolverCfg, + NewtonCfg, + NewtonCollisionPipelineCfg, + SemiImplicitSolverCfg, + VBDSolverCfg, + XPBDSolverCfg, + ) + + torch_device = ( + torch.device(sim_device) if isinstance(sim_device, str) else sim_device + ) + device = self.device + if device is None: + device = f"cuda:{gpu_id}" if torch_device.type == "cuda" else "cpu" + + solver_cfg_map = { + "mjwarp": MJWarpSolverCfg, + "xpbd": XPBDSolverCfg, + "semi_implicit": SemiImplicitSolverCfg, + "featherstone": FeatherstoneSolverCfg, + "vbd": VBDSolverCfg, + } + solver_cfg = solver_cfg_map[self.solver_type]() + + if self.require_grad and self.solver_type != "semi_implicit": + logger.log_error( + "Newton gradient mode requires solver_type='semi_implicit'." + ) + + cfg = NewtonCfg( + dt=physics_dt, + num_substeps=self.num_substeps, + device=device, + debug_mode=self.debug_mode, + require_grad=self.require_grad, + solver_cfg=solver_cfg, + collision_pipeline_cfg=NewtonCollisionPipelineCfg( + broad_phase=self.broad_phase, + requires_grad=self.require_grad, + ), + ) + cfg.use_cuda_graph = self.use_cuda_graph and not self.require_grad + cfg._visualizer_enabled = self.visualizer_enabled + return cfg + + @configclass class MarkerCfg: """Configuration for visual markers in the simulation. diff --git a/embodichain/lab/sim/objects/backends/__init__.py b/embodichain/lab/sim/objects/backends/__init__.py new file mode 100644 index 00000000..e65f00a5 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/__init__.py @@ -0,0 +1,26 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +from .base import RigidBodyViewBase +from .default import DefaultRigidBodyView +from .newton import NewtonRigidBodyView, is_newton_scene + +__all__ = [ + "RigidBodyViewBase", + "DefaultRigidBodyView", + "NewtonRigidBodyView", + "is_newton_scene", +] diff --git a/embodichain/lab/sim/objects/backends/base.py b/embodichain/lab/sim/objects/backends/base.py new file mode 100644 index 00000000..8a44344f --- /dev/null +++ b/embodichain/lab/sim/objects/backends/base.py @@ -0,0 +1,128 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Sequence + +import torch + +__all__ = ["RigidBodyViewBase"] + + +class RigidBodyViewBase(ABC): + """Abstract interface for physics-backend rigid body data access. + + All pose/velocity/acceleration data uses EmbodiChain convention: + ``(x, y, z, qx, qy, qz, qw)``. + """ + + # -- Lifecycle & State -------------------------------------------------- + + @property + @abstractmethod + def is_ready(self) -> bool: + """Whether the backend simulation is finalized and data can be accessed.""" + ... + + # -- Body ID Management ------------------------------------------------- + + @property + @abstractmethod + def body_ids(self) -> list[int]: + """Backend body IDs for all managed entities.""" + ... + + @property + @abstractmethod + def body_ids_tensor(self) -> torch.Tensor: + """Body IDs as an int32 tensor on ``device``.""" + ... + + @abstractmethod + def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> list[int]: + """Return body IDs for the given entity indices.""" + ... + + # -- Pose --------------------------------------------------------------- + + @abstractmethod + def fetch_pose(self, body_ids: Sequence[int] | None = None) -> torch.Tensor: + """Fetch poses as ``(N, 7)`` tensor in ``(x, y, z, qx, qy, qz, qw)``.""" + ... + + @abstractmethod + def apply_pose(self, pose: torch.Tensor, body_ids: Sequence[int]) -> None: + """Apply poses from ``(N, 7)`` tensor in ``(x, y, z, qx, qy, qz, qw)``.""" + ... + + # -- Velocity ----------------------------------------------------------- + + @abstractmethod + def fetch_linear_velocity( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Fetch linear velocities as ``(N, 3)`` tensor.""" + ... + + @abstractmethod + def fetch_angular_velocity( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Fetch angular velocities as ``(N, 3)`` tensor.""" + ... + + @abstractmethod + def apply_linear_velocity( + self, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + """Set linear velocities from ``(N, 3)`` tensor.""" + ... + + @abstractmethod + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + """Set angular velocities from ``(N, 3)`` tensor.""" + ... + + # -- Acceleration ------------------------------------------------------- + + @abstractmethod + def fetch_linear_acceleration( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Fetch linear accelerations as ``(N, 3)`` tensor.""" + ... + + @abstractmethod + def fetch_angular_acceleration( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Fetch angular accelerations as ``(N, 3)`` tensor.""" + ... + + # -- Force & Torque ----------------------------------------------------- + + @abstractmethod + def apply_force(self, data: torch.Tensor, body_ids: Sequence[int]) -> None: + """Apply external forces ``(N, 3)``. One-shot — consumed on next step.""" + ... + + @abstractmethod + def apply_torque(self, data: torch.Tensor, body_ids: Sequence[int]) -> None: + """Apply external torques ``(N, 3)``. One-shot — consumed on next step.""" + ... diff --git a/embodichain/lab/sim/objects/backends/default.py b/embodichain/lab/sim/objects/backends/default.py new file mode 100644 index 00000000..d9139c4b --- /dev/null +++ b/embodichain/lab/sim/objects/backends/default.py @@ -0,0 +1,283 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +from __future__ import annotations + +from typing import Sequence + +import numpy as np +import torch + +from dexsim.models import MeshObject +from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase +from embodichain.utils.math import convert_quat, matrix_from_quat + +__all__ = ["DefaultRigidBodyView"] + + +class DefaultRigidBodyView(RigidBodyViewBase): + """Default DexSim backend rigid body data adapter. + + Encapsulates both GPU (PhysX) and CPU entity-level data paths. + The default GPU API stores pose as ``(qx, qy, qz, qw, x, y, z)``; this + adapter converts to / from the EmbodiChain convention + ``(x, y, z, qx, qy, qz, qw)`` transparently. + """ + + def __init__( + self, + entities: Sequence[MeshObject], + ps: object, + device: torch.device, + ) -> None: + self.entities = list(entities) + self.ps = ps + self.device = device + self._is_gpu = device.type == "cuda" + + if self._is_gpu: + self._gpu_indices = torch.as_tensor( + [entity.get_gpu_index() for entity in self.entities], + dtype=torch.int32, + device=self.device, + ) + else: + self._gpu_indices = None + + # -- RigidBodyViewBase: lifecycle ---------------------------------------- + + @property + def is_ready(self) -> bool: + return True + + # -- RigidBodyViewBase: body IDs ----------------------------------------- + + @property + def body_ids(self) -> list[int]: + if self._is_gpu: + return self._gpu_indices.tolist() + return list(range(len(self.entities))) + + @property + def body_ids_tensor(self) -> torch.Tensor: + if self._is_gpu: + return self._gpu_indices + return torch.arange(len(self.entities), dtype=torch.int32, device=self.device) + + def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> list[int]: + if isinstance(indices, torch.Tensor): + indices = indices.detach().cpu().tolist() + if self._is_gpu: + return self._gpu_indices[list(int(i) for i in indices)].tolist() + return [int(i) for i in indices] + + # -- RigidBodyViewBase: pose --------------------------------------------- + + def fetch_pose(self, body_ids: Sequence[int] | None = None) -> torch.Tensor: + if self._is_gpu: + indices = self._indices_tensor(body_ids) + out = torch.zeros( + (len(indices), 7), dtype=torch.float32, device=self.device + ) + self.ps.gpu_fetch_rigid_body_data( + data=out, + gpu_indices=indices, + data_type=RigidBodyGPUAPIReadType.POSE, + ) + # Convert (qx, qy, qz, qw, x, y, z) -> (x, y, z, qx, qy, qz, qw) + quat = out[:, :4].clone() + xyz = out[:, 4:7].clone() + out[:, :3] = xyz + out[:, 3:7] = quat + return out + + entities = self._select_entities(body_ids) + xyzs = torch.as_tensor( + np.array([e.get_location() for e in entities]), + dtype=torch.float32, + device=self.device, + ) + quats = torch.as_tensor( + np.array([e.get_rotation_quat() for e in entities]), + dtype=torch.float32, + device=self.device, + ) + return torch.cat((xyzs, quats), dim=-1) + + def apply_pose(self, pose: torch.Tensor, body_ids: Sequence[int]) -> None: + pose = pose.to(dtype=torch.float32) + if self._is_gpu: + # Convert (x, y, z, qx, qy, qz, qw) -> (qx, qy, qz, qw, x, y, z) + xyz = pose[:, :3] + quat = pose[:, 3:7] + gpu_pose = torch.cat((quat, xyz), dim=-1) + indices = self._indices_tensor(body_ids) + torch.cuda.synchronize(self.device) + self.ps.gpu_apply_rigid_body_data( + data=gpu_pose.clone(), + gpu_indices=indices, + data_type=RigidBodyGPUAPIWriteType.POSE, + ) + return + + # CPU: convert (x, y, z, qx, qy, qz, qw) -> 4x4 matrix per entity + indices = list(body_ids) + pose_cpu = pose.cpu() + mat = torch.eye(4, dtype=torch.float32).unsqueeze(0).repeat(len(indices), 1, 1) + mat[:, :3, 3] = pose_cpu[:, :3] + mat[:, :3, :3] = matrix_from_quat(convert_quat(pose_cpu[:, 3:7], to="wxyz")) + for i, idx in enumerate(indices): + self.entities[idx].set_local_pose(mat[i]) + + # -- RigidBodyViewBase: velocity ----------------------------------------- + + def fetch_linear_velocity( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3( + RigidBodyGPUAPIReadType.LINEAR_VELOCITY, + "get_linear_velocity", + body_ids, + ) + + def fetch_angular_velocity( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3( + RigidBodyGPUAPIReadType.ANGULAR_VELOCITY, + "get_angular_velocity", + body_ids, + ) + + def apply_linear_velocity( + self, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, + "set_linear_velocity", + data, + body_ids, + ) + + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, + "set_angular_velocity", + data, + body_ids, + ) + + # -- RigidBodyViewBase: acceleration ------------------------------------- + + def fetch_linear_acceleration( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3( + RigidBodyGPUAPIReadType.LINEAR_ACCELERATION, + "get_linear_acceleration", + body_ids, + ) + + def fetch_angular_acceleration( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3( + RigidBodyGPUAPIReadType.ANGULAR_ACCELERATION, + "get_angular_acceleration", + body_ids, + ) + + # -- RigidBodyViewBase: force & torque ----------------------------------- + + def apply_force(self, data: torch.Tensor, body_ids: Sequence[int]) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.FORCE, + "add_force", + data, + body_ids, + ) + + def apply_torque(self, data: torch.Tensor, body_ids: Sequence[int]) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.TORQUE, + "add_torque", + data, + body_ids, + ) + + # -- Internal helpers ---------------------------------------------------- + + def _indices_tensor(self, body_ids: Sequence[int] | None) -> torch.Tensor: + """Return GPU indices as an int32 tensor on device.""" + if body_ids is None: + return self._gpu_indices + if isinstance(body_ids, torch.Tensor): + return body_ids.to(device=self.device, dtype=torch.int32) + return torch.as_tensor(body_ids, dtype=torch.int32, device=self.device) + + def _select_entities(self, body_ids: Sequence[int] | None) -> list[MeshObject]: + """Select entities by body IDs (entity list indices for CPU).""" + if body_ids is None: + return self.entities + return [self.entities[int(i)] for i in body_ids] + + def _fetch_vec3( + self, + gpu_read_type, + cpu_method: str, + body_ids: Sequence[int] | None, + ) -> torch.Tensor: + """Fetch a vec3 field from GPU or CPU entities.""" + if self._is_gpu: + indices = self._indices_tensor(body_ids) + out = torch.zeros( + (len(indices), 3), dtype=torch.float32, device=self.device + ) + self.ps.gpu_fetch_rigid_body_data( + data=out, gpu_indices=indices, data_type=gpu_read_type + ) + return out + + entities = self._select_entities(body_ids) + return torch.as_tensor( + np.array([getattr(e, cpu_method)() for e in entities]), + dtype=torch.float32, + device=self.device, + ) + + def _apply_vec3( + self, + gpu_write_type, + cpu_method: str, + data: torch.Tensor, + body_ids: Sequence[int], + ) -> None: + """Apply a vec3 field to GPU or CPU entities.""" + data = data.to(dtype=torch.float32) + if self._is_gpu: + indices = self._indices_tensor(body_ids) + torch.cuda.synchronize(self.device) + self.ps.gpu_apply_rigid_body_data( + data=data, gpu_indices=indices, data_type=gpu_write_type + ) + return + + indices = list(body_ids) + data_cpu = data.cpu().numpy() + for i, idx in enumerate(indices): + getattr(self.entities[idx], cpu_method)(data_cpu[i]) diff --git a/embodichain/lab/sim/objects/backends/newton.py b/embodichain/lab/sim/objects/backends/newton.py new file mode 100644 index 00000000..122c44c7 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -0,0 +1,223 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +from __future__ import annotations + +from typing import Sequence + +import numpy as np +import torch +import warp as wp + +from dexsim.models import MeshObject +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase +from embodichain.utils import logger + +__all__ = ["NewtonRigidBodyView", "is_newton_scene"] + +_UINT64_MAX = (1 << 64) - 1 +_INT32_MAX = (1 << 31) - 1 + + +def _normalize_native_handle(handle: int, owner: str) -> int: + value = int(handle) + if value < 0: + value &= _UINT64_MAX + if value > _UINT64_MAX: + logger.log_error(f"{owner} native handle is outside uint64 range: {value}.") + return value + + +def is_newton_scene(scene: object) -> bool: + """Return whether *scene* looks like a DexSim Newton scene view.""" + return ( + scene is not None + and hasattr(scene, "manager") + and hasattr(scene, "gpu_fetch_rigid_body_data") + and hasattr(scene, "gpu_apply_rigid_body_data") + ) + + +class NewtonRigidBodyView(RigidBodyViewBase): + """Adapter around DexSim Newton rigid body scene APIs. + + EmbodiChain public rigid-body pose convention is + ``(x, y, z, qx, qy, qz, qw)``. + DexSim Newton exposes the same pose convention through its unified rigid + data API. + """ + + _DATA_TYPE = None # lazily resolved NewtonRigidDataType + + def __init__( + self, + entities: Sequence[MeshObject], + scene: object, + device: torch.device, + ) -> None: + self.entities = list(entities) + self.scene = scene + self.device = device + self.entity_handles = [ + _normalize_native_handle(entity.get_native_handle(), "MeshObject") + for entity in self.entities + ] + self._body_ids = [self._resolve_body_id(entity) for entity in self.entities] + if any(bid < 0 or bid > _INT32_MAX for bid in self._body_ids): + logger.log_error( + "Newton rigid body view found an entity without a Newton body id." + ) + self._body_ids_tensor = torch.as_tensor( + self._body_ids, dtype=torch.int32, device=self.device + ) + + # -- Lazy enum access --------------------------------------------------- + + @classmethod + def _get_data_type(cls): + """Lazily resolve *NewtonRigidDataType* to avoid eager import.""" + if cls._DATA_TYPE is None: + from dexsim.engine.newton_physics import NewtonRigidDataType + + cls._DATA_TYPE = NewtonRigidDataType + return cls._DATA_TYPE + + # -- RigidBodyViewBase: lifecycle ---------------------------------------- + + @property + def is_ready(self) -> bool: + manager = getattr(self.scene, "manager", None) + return ( + manager is not None + and getattr(getattr(manager, "lifecycle_state", None), "name", "") + == "READY" + ) + + # -- RigidBodyViewBase: body IDs ----------------------------------------- + + @property + def body_ids(self) -> list[int]: + return self._body_ids + + @property + def body_ids_tensor(self) -> torch.Tensor: + return self._body_ids_tensor + + def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> list[int]: + if isinstance(indices, torch.Tensor): + indices = indices.detach().cpu().tolist() + return [self._body_ids[int(index)] for index in indices] + + # -- RigidBodyViewBase: pose --------------------------------------------- + + def fetch_pose(self, body_ids: Sequence[int] | None = None) -> torch.Tensor: + body_ids = self._body_ids if body_ids is None else list(body_ids) + out = self._warp_array((len(body_ids), 7)) + self.scene.gpu_fetch_rigid_body_data(body_ids, self._get_data_type().POSE, out) + return self._to_torch(out) + + def apply_pose(self, pose: torch.Tensor, body_ids: Sequence[int]) -> None: + self._apply_data(body_ids, self._get_data_type().POSE, pose) + + # -- RigidBodyViewBase: velocity ----------------------------------------- + + def fetch_linear_velocity( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3(self._get_data_type().LINEAR_VELOCITY, body_ids) + + def fetch_angular_velocity( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3(self._get_data_type().ANGULAR_VELOCITY, body_ids) + + def apply_linear_velocity( + self, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + self._apply_data(body_ids, self._get_data_type().LINEAR_VELOCITY, data) + + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + self._apply_data(body_ids, self._get_data_type().ANGULAR_VELOCITY, data) + + # -- RigidBodyViewBase: acceleration ------------------------------------- + + def fetch_linear_acceleration( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3(self._get_data_type().LINEAR_ACCELERATION, body_ids) + + def fetch_angular_acceleration( + self, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + return self._fetch_vec3(self._get_data_type().ANGULAR_ACCELERATION, body_ids) + + # -- RigidBodyViewBase: force & torque ----------------------------------- + + def apply_force(self, data: torch.Tensor, body_ids: Sequence[int]) -> None: + self._apply_data(body_ids, self._get_data_type().FORCE, data) + + def apply_torque(self, data: torch.Tensor, body_ids: Sequence[int]) -> None: + self._apply_data(body_ids, self._get_data_type().TORQUE, data) + + # -- Internal helpers ---------------------------------------------------- + + def _resolve_body_id(self, entity: MeshObject) -> int: + manager = getattr(self.scene, "manager", None) + if manager is not None and hasattr(entity, "get_native_handle"): + entity_handle = _normalize_native_handle( + entity.get_native_handle(), "MeshObject" + ) + body_id = getattr(manager, "dexsim2newton_body", {}).get(entity_handle) + if body_id is not None: + return int(body_id) + + if hasattr(entity, "get_gpu_index"): + body_id = int(entity.get_gpu_index()) + if 0 <= body_id <= _INT32_MAX: + return body_id + return -1 + + def _warp_array(self, shape: tuple[int, int]): + """Allocate a Warp float32 array on the simulation device.""" + manager = self.scene.manager + state = getattr(manager, "_state_0", None) + warp_device = state.body_q.device if state is not None else manager._device + return wp.empty(shape, dtype=wp.float32, device=warp_device) + + def _to_torch(self, array) -> torch.Tensor: + """Convert a Warp array to a float32 torch tensor on ``self.device``.""" + if str(array.device).startswith("cuda"): + return wp.to_torch(array).to(device=self.device, dtype=torch.float32) + return torch.as_tensor(array.numpy(), dtype=torch.float32, device=self.device) + + def _fetch_vec3( + self, data_type, body_ids: Sequence[int] | None = None + ) -> torch.Tensor: + body_ids = self._body_ids if body_ids is None else list(body_ids) + out = self._warp_array((len(body_ids), 3)) + self.scene.gpu_fetch_rigid_body_data(body_ids, data_type, out) + return self._to_torch(out) + + def _apply_data( + self, body_ids: Sequence[int], data_type, data: torch.Tensor + ) -> None: + """Apply data to bodies via the unified Newton GPU API.""" + data = data.to(dtype=torch.float32) + state = getattr(self.scene.manager, "_state_0", None) + is_cuda = state is not None and str(state.body_q.device).startswith("cuda") + payload = data if is_cuda else data.detach().cpu().numpy() + self.scene.gpu_apply_rigid_body_data(list(body_ids), data_type, payload) diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index 2202bbec..9d2aa924 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -23,9 +23,14 @@ from functools import cached_property from dexsim.models import MeshObject -from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType -from dexsim.engine import CudaArray, PhysicsScene +from dexsim.engine import PhysicsScene from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg +from embodichain.lab.sim.objects.backends import ( + DefaultRigidBodyView, + NewtonRigidBodyView, + is_newton_scene, +) +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase from embodichain.lab.sim import ( VisualMaterial, VisualMaterialInst, @@ -40,8 +45,8 @@ class RigidBodyData: """Data manager for rigid body with body type of dynamic or kinematic. - Note: - 1. The pose data managed by dexsim is in the format of (qx, qy, qz, qw, x, y, z), but in SimulationManager, we use (x, y, z, qw, qx, qy, qz) format. + All pose/velocity/acceleration data uses EmbodiChain convention: + ``(x, y, z, qx, qy, qz, qw)``. """ def __init__( @@ -59,16 +64,18 @@ def __init__( self.num_instances = len(entities) self.device = device - # get gpu indices for the entities. - self.gpu_indices = ( - torch.as_tensor( - [entity.get_gpu_index() for entity in self.entities], - dtype=torch.int32, - device=self.device, + # Create the appropriate backend view. + if is_newton_scene(ps): + self._body_view: RigidBodyViewBase = NewtonRigidBodyView( + entities=entities, scene=ps, device=device ) - if self.device.type == "cuda" - else None - ) + else: + self._body_view = DefaultRigidBodyView( + entities=entities, ps=ps, device=device + ) + + # Kept for backward compatibility with callers that index gpu_indices directly. + self.gpu_indices = self._body_view.body_ids_tensor # Initialize rigid body data. self._pose = torch.zeros( @@ -86,7 +93,7 @@ def __init__( self._ang_acc = torch.zeros( (self.num_instances, 3), dtype=torch.float32, device=self.device ) - # center of mass pose in format (x, y, z, qw, qx, qy, qz) + # center of mass pose in format (x, y, z, qx, qy, qz, qw) self.default_com_pose = torch.zeros( (self.num_instances, 7), dtype=torch.float32, device=self.device ) @@ -94,67 +101,56 @@ def __init__( (self.num_instances, 7), dtype=torch.float32, device=self.device ) + @property + def is_newton_backend(self) -> bool: + return isinstance(self._body_view, NewtonRigidBodyView) + + @property + def is_newton_ready(self) -> bool: + return self.is_newton_backend and self._body_view.is_ready + + def body_ids_for(self, env_ids: Sequence[int]) -> list[int]: + return self._body_view.select_body_ids(env_ids) + @property def pose(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch pose from CPU entities - xyzs = torch.as_tensor( - np.array([entity.get_location() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - quats = torch.as_tensor( - np.array( - [entity.get_rotation_quat() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, + if self._body_view.is_ready: + self._pose = self._body_view.fetch_pose() + return self._pose + + # Newton backend not yet finalized — use entity API fallback. + for i, entity in enumerate(self.entities): + pos = entity.get_location() + quat = entity.get_rotation_quat() + self._pose[i, :3] = torch.as_tensor( + pos, dtype=torch.float32, device=self.device ) - quats = convert_quat(quats, to="wxyz") - self._pose = torch.cat((xyzs, quats), dim=-1) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._pose, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.POSE, + self._pose[i, 3:7] = torch.as_tensor( + quat, dtype=torch.float32, device=self.device ) - self._pose[:, :4] = convert_quat(self._pose[:, :4], to="wxyz") - self._pose = self._pose[:, [4, 5, 6, 0, 1, 2, 3]] 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( - np.array([entity.get_linear_velocity() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._lin_vel, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.LINEAR_VELOCITY, + if self._body_view.is_ready: + self._lin_vel = self._body_view.fetch_linear_velocity() + return self._lin_vel + + for i, entity in enumerate(self.entities): + self._lin_vel[i] = torch.as_tensor( + entity.get_linear_velocity(), dtype=torch.float32, device=self.device ) 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( - np.array( - [entity.get_angular_velocity() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._ang_vel, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.ANGULAR_VELOCITY, + if self._body_view.is_ready: + self._ang_vel = self._body_view.fetch_angular_velocity() + return self._ang_vel + + for i, entity in enumerate(self.entities): + self._ang_vel[i] = torch.as_tensor( + entity.get_angular_velocity(), dtype=torch.float32, device=self.device ) return self._ang_vel @@ -169,38 +165,30 @@ def vel(self) -> torch.Tensor: @property def lin_acc(self) -> torch.Tensor: - if self.device.type == "cpu": - self._lin_acc = torch.as_tensor( - np.array( - [entity.get_linear_acceleration() for entity in self.entities], - ), + if self._body_view.is_ready: + self._lin_acc = self._body_view.fetch_linear_acceleration() + return self._lin_acc + + for i, entity in enumerate(self.entities): + self._lin_acc[i] = torch.as_tensor( + entity.get_linear_acceleration(), dtype=torch.float32, device=self.device, ) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._lin_acc, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.LINEAR_ACCELERATION, - ) return self._lin_acc @property def ang_acc(self) -> torch.Tensor: - if self.device.type == "cpu": - self._ang_acc = torch.as_tensor( - np.array( - [entity.get_angular_acceleration() for entity in self.entities], - ), + if self._body_view.is_ready: + self._ang_acc = self._body_view.fetch_angular_acceleration() + return self._ang_acc + + for i, entity in enumerate(self.entities): + self._ang_acc[i] = torch.as_tensor( + entity.get_angular_acceleration(), dtype=torch.float32, device=self.device, ) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._ang_acc, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.ANGULAR_ACCELERATION, - ) return self._ang_acc @property @@ -219,13 +207,35 @@ def com_pose(self) -> torch.Tensor: Returns: torch.Tensor: The center of mass pose with shape (N, 7). """ + if self.is_newton_backend: + manager = self._body_view.scene.manager + for i, entity_handle in enumerate(self._body_view.entity_handles): + attr = manager.dexsim_meta.get(entity_handle, {}).get("attr") + if attr is None: + pos = np.zeros(3, dtype=np.float32) + quat = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32) + else: + pos = np.asarray(attr.com_position, dtype=np.float32).copy() + quat = np.asarray(attr.com_quaternion, dtype=np.float32).copy() + self._com_pose[i, :3] = torch.as_tensor( + pos, dtype=torch.float32, device=self.device + ) + self._com_pose[i, 3:7] = torch.as_tensor( + convert_quat(quat, to="xyzw"), + dtype=torch.float32, + device=self.device, + ) + return self._com_pose + for i, entity in enumerate(self.entities): pos, quat = entity.get_physical_body().get_cmass_local_pose() self._com_pose[i, :3] = torch.as_tensor( pos, dtype=torch.float32, device=self.device ) self._com_pose[i, 3:7] = torch.as_tensor( - quat, dtype=torch.float32, device=self.device + convert_quat(np.asarray(quat, dtype=np.float32), to="xyzw"), + dtype=torch.float32, + device=self.device, ) return self._com_pose @@ -265,6 +275,8 @@ def __init__( # Determine if we should use USD properties or cfg properties. if not cfg.use_usd_properties: for entity in entities: + if is_newton_scene(self._ps): + continue entity.set_body_scale(*cfg.body_scale) entity.set_physical_attr(cfg.attrs.attr()) else: @@ -277,7 +289,7 @@ def __init__( first_entity.get_physical_attr().as_dict() ) - if device.type == "cuda": + if device.type == "cuda" and not is_newton_scene(self._ps): self._world.update(0.001) super().__init__(cfg, entities, device) @@ -286,8 +298,8 @@ def __init__( self._set_default_collision_filter() # update default center of mass pose (only for non-static bodies with body data). - if self.body_data is not None: - self.body_data.default_com_pose = self.body_data.com_pose.clone() + if self._data is not None: + self._data.default_com_pose = self._data.com_pose.clone() # TODO: Must be called after setting all attributes. # May be improved in the future. @@ -336,7 +348,7 @@ 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] + [x, y, z, qx, qy, qz, qw, 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. @@ -402,9 +414,11 @@ def set_collision_filter( filter_data_np = filter_data.cpu().numpy().astype(np.uint32) for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].get_physical_body().set_collision_filter_data( - filter_data_np[i] - ) + entity = self._entities[env_idx] + if is_newton_scene(self._ps): + entity.set_collision_filter_data(filter_data_np[i]) + else: + entity.get_physical_body().set_collision_filter_data(filter_data_np[i]) def set_local_pose( self, pose: torch.Tensor, env_ids: Sequence[int] | None = None @@ -422,50 +436,47 @@ def set_local_pose( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." ) - if self.device.type == "cpu" or self.is_static: - pose = pose.cpu() - if pose.dim() == 2 and pose.shape[1] == 7: - pose_matrix = torch.eye(4).unsqueeze(0).repeat(pose.shape[0], 1, 1) - pose_matrix[:, :3, 3] = pose[:, :3] - pose_matrix[:, :3, :3] = matrix_from_quat(pose[:, 3:7]) - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_local_pose(pose_matrix[i]) - elif pose.dim() == 3 and pose.shape[1:] == (4, 4): - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_local_pose(pose[i]) - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - + # Normalize pose to (N, 7) format in (x, y, z, qx, qy, qz, qw). + if pose.dim() == 2 and pose.shape[1] == 7: + target_pose = pose.to(device=self.device, dtype=torch.float32) + elif pose.dim() == 3 and pose.shape[1:] == (4, 4): + xyz = pose[:, :3, 3] + quat = convert_quat(quat_from_matrix(pose[:, :3, :3]), to="xyzw") + target_pose = torch.cat((xyz, quat), dim=-1).to( + device=self.device, dtype=torch.float32 + ) else: - if pose.dim() == 2 and pose.shape[1] == 7: - xyz = pose[:, :3] - quat = convert_quat(pose[:, 3:7], to="xyzw") - elif pose.dim() == 3 and pose.shape[1:] == (4, 4): - xyz = pose[:, :3, 3] - quat = quat_from_matrix(pose[:, :3, :3]) - 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] - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=pose.clone(), - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.POSE, + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." ) + return + + # Use backend view if available and ready. + if ( + self._data is not None + and self._data._body_view.is_ready + and not self.is_static + ): + body_ids = self._data.body_ids_for(local_env_ids) + self._data._body_view.apply_pose(target_pose, body_ids) + return + + # Static bodies and non-ready backends (notably Newton before finalize) + # still accept direct entity pose updates. + target_pose = target_pose.cpu() + pose_matrix = torch.eye(4).unsqueeze(0).repeat(len(local_env_ids), 1, 1) + pose_matrix[:, :3, 3] = target_pose[:, :3] + pose_matrix[:, :3, :3] = matrix_from_quat( + convert_quat(target_pose[:, 3:7], to="wxyz") + ) + for i, env_idx in enumerate(local_env_ids): + self._entities[env_idx].set_local_pose(pose_matrix[i]) def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: """Get local pose of the rigid object. 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. + to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qx, qy, qz, qw). Defaults to False. Returns: torch.Tensor: The local pose of the rigid object with shape (N, 7) or (N, 4, 4) depending on `to_matrix`. @@ -484,7 +495,6 @@ def get_local_pose_cpu( quats = torch.as_tensor( [entity.get_rotation_quat() for entity in entities] ) - quats = convert_quat(quats, to="wxyz") pose = torch.cat((xyzs, quats), dim=-1) return pose @@ -495,7 +505,7 @@ def get_local_pose_cpu( pose = self.body_data.pose if to_matrix: xyz = pose[:, :3] - mat = matrix_from_quat(pose[:, 3:7]) + mat = matrix_from_quat(convert_quat(pose[:, 3:7], to="wxyz")) pose = ( torch.eye(4, dtype=torch.float32, device=self.device) .unsqueeze(0) @@ -550,29 +560,21 @@ def add_force_torque( f"Length of env_ids {len(local_env_ids)} does not match torque length {len(torque)}." ) - if self.device.type == "cpu": + if self._data is not None and self._data._body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + if force is not None: + self._data._body_view.apply_force(force, body_ids) + if torque is not None: + self._data._body_view.apply_torque(torque, body_ids) + elif self._data is not None and self._data.is_newton_backend: + return + else: for i, env_idx in enumerate(local_env_ids): if force is not None: self._entities[env_idx].add_force(force[i].cpu().numpy()) if torque is not None: self._entities[env_idx].add_torque(torque[i].cpu().numpy()) - else: - indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) - if force is not None: - self._ps.gpu_apply_rigid_body_data( - data=force, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.FORCE, - ) - if torque is not None: - self._ps.gpu_apply_rigid_body_data( - data=torque, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.TORQUE, - ) - def set_velocity( self, lin_vel: torch.Tensor | None = None, @@ -608,7 +610,15 @@ def set_velocity( f"Length of env_ids {len(local_env_ids)} does not match ang_vel length {len(ang_vel)}." ) - if self.device.type == "cpu": + if self._data is not None and self._data._body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + if lin_vel is not None: + self._data._body_view.apply_linear_velocity(lin_vel, body_ids) + if ang_vel is not None: + self._data._body_view.apply_angular_velocity(ang_vel, body_ids) + elif self._data is not None and self._data.is_newton_backend: + return + else: for i, env_idx in enumerate(local_env_ids): if lin_vel is not None: self._entities[env_idx].set_linear_velocity( @@ -618,21 +628,6 @@ def set_velocity( self._entities[env_idx].set_angular_velocity( ang_vel[i].cpu().numpy() ) - else: - indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) - if lin_vel is not None: - self._ps.gpu_apply_rigid_body_data( - data=lin_vel, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, - ) - if ang_vel is not None: - self._ps.gpu_apply_rigid_body_data( - data=ang_vel, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, - ) def set_attrs( self, @@ -941,7 +936,7 @@ def set_body_scale( def set_com_pose( self, com_pose: torch.Tensor, env_ids: Sequence[int] | None = None ) -> None: - """Set the center of mass pose of the rigid body. The pose format is (x, y, z, qw, qx, qy, qz). + """Set the center of mass pose of the rigid body. The pose format is (x, y, z, qx, qy, qz, qw). Args: com_pose (torch.Tensor): The center of mass pose to set with shape (N, 7). @@ -963,8 +958,13 @@ def set_com_pose( com_pose = com_pose.cpu().numpy() for i, env_idx in enumerate(local_env_ids): pos = com_pose[i, :3] - quat = com_pose[i, 3:7] - self._entities[env_idx].get_physical_body().set_cmass_local_pose(pos, quat) + quat = convert_quat(com_pose[i, 3:7], to="wxyz") + if self._data is not None and self._data.is_newton_backend: + self._entities[env_idx].set_cmass_local_pose(pos, quat) + else: + self._entities[env_idx].get_physical_body().set_cmass_local_pose( + pos, quat + ) def set_body_type(self, body_type: str) -> None: """Set the body type of the rigid object. @@ -1081,36 +1081,20 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: 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: - self._entities[env_idx].clear_dynamics() - else: - # Apply zero force and torque to the rigid bodies. + if self._data is not None and self._data._body_view.is_ready: zeros = torch.zeros( (len(local_env_ids), 3), dtype=torch.float32, device=self.device ) - indices = self.body_data.gpu_indices[local_env_ids] - 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, - ) + body_ids = self._data.body_ids_for(local_env_ids) + self._data._body_view.apply_linear_velocity(zeros, body_ids) + self._data._body_view.apply_angular_velocity(zeros, body_ids) + self._data._body_view.apply_force(zeros, body_ids) + self._data._body_view.apply_torque(zeros, body_ids) + elif self._data is not None and self._data.is_newton_backend: + return + else: + for env_idx in local_env_ids: + self._entities[env_idx].clear_dynamics() def set_physical_visible( self, @@ -1190,4 +1174,7 @@ def destroy(self) -> None: if len(arenas) == 0: arenas = [env] for i, entity in enumerate(self._entities): - arenas[i].remove_actor(entity) + if is_newton_scene(self._ps): + arenas[i].remove_actor(entity.get_name()) + else: + arenas[i].remove_actor(entity) diff --git a/embodichain/lab/sim/objects/rigid_object_group.py b/embodichain/lab/sim/objects/rigid_object_group.py index e4cca592..dc2007a0 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/embodichain/lab/sim/objects/rigid_object_group.py @@ -22,12 +22,17 @@ from typing import List, Sequence, Union from dexsim.models import MeshObject -from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType -from dexsim.engine import CudaArray, PhysicsScene +from dexsim.engine import PhysicsScene from embodichain.lab.sim.cfg import ( RigidObjectGroupCfg, RigidBodyAttributesCfg, ) +from embodichain.lab.sim.objects.backends import ( + DefaultRigidBodyView, + NewtonRigidBodyView, + is_newton_scene, +) +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase from embodichain.lab.sim import ( BatchEntity, ) @@ -56,19 +61,21 @@ def __init__( self.num_instances = len(entities) self.num_objects = len(entities[0]) self.device = device + self.flat_entities = [entity for instance in entities for entity in instance] - # 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, + # Create the appropriate backend view. + if is_newton_scene(ps): + self._body_view: RigidBodyViewBase = NewtonRigidBodyView( + entities=self.flat_entities, scene=ps, device=device + ) + else: + self._body_view = DefaultRigidBodyView( + entities=self.flat_entities, ps=ps, device=device ) - if self.device.type == "cuda" - else None + + # get gpu indices for the rigid bodies with shape of (num_instances, num_objects) + self.gpu_indices = self._body_view.body_ids_tensor.reshape( + self.num_instances, self.num_objects ) # Initialize rigid body group data tensors. Shape of (num_instances, num_objects, data_dim) @@ -88,79 +95,77 @@ def __init__( device=self.device, ) + @property + def is_newton_backend(self) -> bool: + return isinstance(self._body_view, NewtonRigidBodyView) + + @property + def is_newton_ready(self) -> bool: + return self.is_newton_backend and self._body_view.is_ready + + def body_ids_for( + self, + env_ids: Sequence[int], + obj_ids: Sequence[int] | None = None, + ) -> list[int]: + local_obj_ids = range(self.num_objects) if obj_ids is None else obj_ids + flat_indices = [] + for env_idx in env_ids: + for obj_idx in local_obj_ids: + flat_indices.append(int(env_idx) * self.num_objects + int(obj_idx)) + return self._body_view.select_body_ids(flat_indices) + @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, + if self._body_view.is_ready: + self._pose = self._body_view.fetch_pose().reshape( + self.num_instances, self.num_objects, 7 ) - 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 + # Newton not ready — entity API fallback. + for i, instance in enumerate(self.entities): + for j, entity in enumerate(instance): + self._pose[i, j, :3] = torch.as_tensor( + entity.get_location(), dtype=torch.float32, device=self.device + ) + self._pose[i, j, 3:7] = torch.as_tensor( + entity.get_rotation_quat(), dtype=torch.float32, device=self.device + ) + 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, + if self._body_view.is_ready: + self._lin_vel = self._body_view.fetch_linear_velocity().reshape( + self.num_instances, self.num_objects, 3 ) + return self._lin_vel + + for i, instance in enumerate(self.entities): + for j, entity in enumerate(instance): + self._lin_vel[i, j] = torch.as_tensor( + entity.get_linear_velocity(), + dtype=torch.float32, + device=self.device, + ) 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_linear_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, + if self._body_view.is_ready: + self._ang_vel = self._body_view.fetch_angular_velocity().reshape( + self.num_instances, self.num_objects, 3 ) + return self._ang_vel + + for i, instance in enumerate(self.entities): + for j, entity in enumerate(instance): + self._ang_vel[i, j] = torch.as_tensor( + entity.get_angular_velocity(), + dtype=torch.float32, + device=self.device, + ) return self._ang_vel @property @@ -198,10 +203,12 @@ def __init__( body_cfgs = list(cfg.rigid_objects.values()) for instance in entities: for i, body in enumerate(instance): + if is_newton_scene(self._ps): + continue body.set_body_scale(*body_cfgs[i].body_scale) body.set_physical_attr(body_cfgs[i].attrs.attr()) - if device.type == "cuda": + if device.type == "cuda" and not is_newton_scene(self._ps): self._world.update(0.001) super().__init__(cfg, entities, device) @@ -243,7 +250,7 @@ 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] + [x, y, z, qx, qy, qz, qw, 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. @@ -297,7 +304,12 @@ def set_collision_filter( 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]) + if is_newton_scene(self._ps): + entity.set_collision_filter_data(filter_data_np[i]) + else: + entity.get_physical_body().set_collision_filter_data( + filter_data_np[i] + ) def set_local_pose( self, @@ -321,62 +333,47 @@ def set_local_pose( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." ) - 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, + # Normalize pose to (N*M, 7) format in (x, y, z, qx, qy, qz, qw). + if pose.dim() == 3 and pose.shape[2] == 7: + target_pose = pose.reshape(-1, 7).to( + device=self.device, dtype=torch.float32 + ) + 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 = convert_quat(quat_from_matrix(mat), to="xyzw") + target_pose = torch.cat((xyz, quat), dim=-1).to( + device=self.device, dtype=torch.float32 ) - self._world.sync_poses_gpu_to_cpu( - rigid_pose=CudaArray(pose), rigid_gpu_indices=CudaArray(indices) + else: + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, M, 7) or (N, M, 4, 4)." ) + return + + # Use backend view if ready. + if self._data._body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids, local_obj_ids) + self._data._body_view.apply_pose(target_pose, body_ids) + return + + # Newton not ready — entity API fallback. + target_pose = target_pose.cpu() + pose_matrix = torch.eye(4).unsqueeze(0).repeat(target_pose.shape[0], 1, 1) + pose_matrix[:, :3, 3] = target_pose[:, :3] + pose_matrix[:, :3, :3] = matrix_from_quat( + convert_quat(target_pose[:, 3:7], to="wxyz") + ) + pose_matrix = pose_matrix.reshape(-1, len(local_obj_ids), 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_matrix[i, j]) 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. + to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qx, qy, qz, qw). 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`. @@ -385,7 +382,7 @@ def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: if to_matrix: pose = pose.reshape(-1, 7) xyz = pose[:, :3] - mat = matrix_from_quat(pose[:, 3:7]) + mat = matrix_from_quat(convert_quat(pose[:, 3:7], to="wxyz")) pose = ( torch.eye(4, dtype=torch.float32, device=self.device) .unsqueeze(0) @@ -422,39 +419,23 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: 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. + if self._data._body_view.is_ready: 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, - ) + body_ids = self._data.body_ids_for(local_env_ids) + self._data._body_view.apply_linear_velocity(zeros, body_ids) + self._data._body_view.apply_angular_velocity(zeros, body_ids) + self._data._body_view.apply_force(zeros, body_ids) + self._data._body_view.apply_torque(zeros, body_ids) + elif self._data.is_newton_backend: + return + else: + for env_idx in local_env_ids: + for entity in self._entities[env_idx]: + entity.clear_dynamics() def set_visual_material( self, mat: VisualMaterial, env_ids: Sequence[int] | None = None diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 8f2e257f..31c66210 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -22,6 +22,7 @@ import queue import time import threading +import importlib import dexsim import torch import numpy as np @@ -76,6 +77,8 @@ from embodichain.lab.sim.cfg import ( RenderCfg, PhysicsCfg, + DefaultPhysicsCfg, + NewtonPhysicsCfg, MarkerCfg, GPUMemoryCfg, WindowRecordCfg, @@ -144,14 +147,30 @@ class SimulationManagerCfg: sim_device: Union[str, torch.device] = "cpu" """The device for the physics simulation. Can be 'cpu', 'cuda', or a torch.device object.""" - physics_config: PhysicsCfg = field(default_factory=PhysicsCfg) - """The physics configuration parameters.""" + physics_backend: str = "default" + """Physics backend name. Supported values are 'default' and 'newton'.""" + + default_physics_cfg: DefaultPhysicsCfg = field(default_factory=DefaultPhysicsCfg) + """The existing DexSim default-backend physics configuration parameters.""" + + newton_physics_cfg: NewtonPhysicsCfg = field(default_factory=NewtonPhysicsCfg) + """DexSim Newton backend physics configuration parameters.""" + + physics_config: PhysicsCfg | None = None + """Deprecated alias for ``default_physics_cfg`` kept for existing configs.""" + gpu_memory_config: GPUMemoryCfg = field(default_factory=GPUMemoryCfg) """The GPU memory configuration parameters.""" window_record: WindowRecordCfg = field(default_factory=WindowRecordCfg) """Viewer window recording settings (hotkey, paths, FPS, memory budget).""" + def __post_init__(self): + if self.physics_config is not None: + self.default_physics_cfg = self.physics_config + else: + self.physics_config = self.default_physics_cfg + @dataclass class _WindowRecordState: @@ -230,6 +249,15 @@ def __init__( self.sim_config = sim_config self.device = torch.device("cpu") + self._physics_backend = getattr( + sim_config, "physics_backend", "default" + ).lower() + if self._physics_backend not in ("default", "newton"): + logger.log_error( + f"Unsupported physics backend '{self._physics_backend}'. " + "Supported backends are 'default' and 'newton'." + ) + self._newton_manager = None world_config = self._convert_sim_config(sim_config) @@ -258,8 +286,15 @@ def __init__( self._world.set_delta_time(sim_config.physics_dt) self._world.show_coordinate_axis(False) - dexsim.set_physics_config(**sim_config.physics_config.to_dexsim_args()) - dexsim.set_physics_gpu_memory_config(**sim_config.gpu_memory_config.to_dict()) + if self.is_default_backend: + dexsim.set_physics_config(**sim_config.default_physics_cfg.to_dexsim_args()) + dexsim.set_physics_gpu_memory_config( + **sim_config.gpu_memory_config.to_dict() + ) + else: + from dexsim.engine.newton_physics import get_newton_manager + + self._newton_manager = get_newton_manager(self._world) self._is_initialized_gpu_physics = False self._ps = self._world.get_physics_scene() @@ -368,8 +403,55 @@ def num_envs(self) -> int: @property def is_use_gpu_physics(self) -> bool: - """Check if the physics simulation is using GPU.""" - return self.device.type == "cuda" + """Check if the default backend GPU physics API is active.""" + return self.is_default_gpu_backend + + @property + def physics_backend(self) -> str: + """Return the active physics backend name.""" + return self._physics_backend + + @property + def is_default_backend(self) -> bool: + """Whether the existing DexSim default physics backend is active.""" + return self._physics_backend == "default" + + @property + def is_newton_backend(self) -> bool: + """Whether the DexSim Newton physics backend is active.""" + return self._physics_backend == "newton" + + @property + def is_default_gpu_backend(self) -> bool: + """Whether the default backend is using the DexSim GPU physics API.""" + return self.is_default_backend and self.device.type == "cuda" + + @property + def is_newton_gpu_backend(self) -> bool: + """Whether Newton is configured to run on CUDA.""" + if not self.is_newton_backend: + return False + mgr = self.newton_manager + if mgr is None: + return self.device.type == "cuda" + return str(mgr.cfg.device).startswith("cuda") + + @property + def newton_manager(self): + """Return the DexSim Newton manager for this world, if active.""" + if not self.is_newton_backend: + return None + if self._newton_manager is None: + from dexsim.engine.newton_physics import get_newton_manager + + self._newton_manager = get_newton_manager(self._world) + return self._newton_manager + + @property + def newton_scene(self): + """Return the DexSim Newton scene view, if active.""" + mgr = self.newton_manager + return None if mgr is None else mgr.newton_scene @property def is_physics_manually_update(self) -> bool: @@ -409,8 +491,8 @@ def _convert_sim_config( world_config.backend = Backend.VULKAN world_config.thread_mode = sim_config.thread_mode world_config.cache_path = str(self._material_cache_dir) - world_config.length_tolerance = sim_config.physics_config.length_tolerance - world_config.speed_tolerance = sim_config.physics_config.speed_tolerance + world_config.length_tolerance = sim_config.default_physics_cfg.length_tolerance + world_config.speed_tolerance = sim_config.default_physics_cfg.speed_tolerance world_config.renderer = sim_config.render_cfg.to_dexsim_flags() if sim_config.render_cfg.enable_denoiser is False: @@ -423,9 +505,6 @@ def _convert_sim_config( self.device = sim_config.sim_device if self.device.type == "cuda": - world_config.enable_gpu_sim = True - world_config.direct_gpu_api = True - if self.device.index is not None and sim_config.gpu_id != self.device.index: logger.log_warning( f"Conflict gpu_id {sim_config.gpu_id} and device index {self.device.index}. Using device index." @@ -434,6 +513,19 @@ def _convert_sim_config( self.device = torch.device(f"cuda:{sim_config.gpu_id}") + if self.is_default_backend and self.device.type == "cuda": + world_config.enable_gpu_sim = True + world_config.direct_gpu_api = True + + if self.is_newton_backend: + importlib.import_module("dexsim.engine.newton_physics") + + world_config.newton_cfg = sim_config.newton_physics_cfg.to_dexsim_cfg( + physics_dt=sim_config.physics_dt, + sim_device=self.device, + gpu_id=sim_config.gpu_id, + ) + world_config.gpu_id = sim_config.gpu_id return world_config @@ -465,7 +557,10 @@ def set_manual_update(self, enable: bool) -> None: def init_gpu_physics(self) -> None: """Initialize the GPU physics simulation.""" - if self.device.type != "cuda": + if self.is_newton_backend: + return + + if not self.is_default_gpu_backend: logger.log_warning( "The simulation device is not cuda, cannot initialize GPU physics." ) @@ -483,6 +578,20 @@ def init_gpu_physics(self) -> None: self._is_initialized_gpu_physics = True + def prepare_physics(self) -> None: + """Prepare backend-specific runtime data after scene construction.""" + if self.is_default_gpu_backend: + self.init_gpu_physics() + elif self.is_newton_backend: + self._world.update(0.0) + + def forward_physics(self) -> None: + """Refresh backend physics state without advancing time when supported.""" + if self.is_newton_backend: + mgr = self.newton_manager + if mgr is not None and getattr(mgr.lifecycle_state, "name", "") == "READY": + mgr.forward_kinematics() + def render_camera_group(self, group_ids: list[int]) -> None: """Render all camera group in the simulation. @@ -501,7 +610,7 @@ def update(self, physics_dt: float | None = None, step: int = 10) -> None: physics_dt (float | None, optional): the time step for physics simulation. Defaults to None. step (int, optional): the number of steps to update physics. Defaults to 10. """ - if self.is_use_gpu_physics and not self._is_initialized_gpu_physics: + if self.is_default_gpu_backend and not self._is_initialized_gpu_physics: logger.log_warning( f"Using GPU physics, but not initialized yet. Forcing initialization." ) @@ -624,6 +733,11 @@ def _create_default_plane(self): self._default_plane = self._env.create_plane( 0, default_length, repeat_uv_size, repeat_uv_size ) + if self.is_newton_backend and self.newton_manager is not None: + plane_handle = int(self._default_plane.get_native_handle()) + if plane_handle < 0: + plane_handle &= (1 << 64) - 1 + self.newton_manager.dexsim_meta.pop(plane_handle, None) self._default_plane.set_name("default_plane") plane_collision = self._env.create_cube( default_length, default_length, default_length / 10 @@ -841,6 +955,12 @@ def add_soft_object(self, cfg: SoftObjectCfg) -> SoftObject: Returns: SoftObject: The added soft object instance handle. """ + if self.is_newton_backend: + logger.log_error( + "Soft object support for the Newton backend is not enabled in EmbodiChain yet.", + error_type=NotImplementedError, + ) + if not self.is_use_gpu_physics: logger.log_error("Soft object requires GPU physics to be enabled.") @@ -871,6 +991,12 @@ def add_cloth_object(self, cfg: ClothObjectCfg) -> ClothObject: Returns: ClothObject: The added cloth object instance handle. """ + if self.is_newton_backend: + logger.log_error( + "Cloth object support for the Newton backend is not enabled in EmbodiChain yet.", + error_type=NotImplementedError, + ) + if not self.is_use_gpu_physics: logger.log_error("Cloth object requires GPU physics to be enabled.") @@ -1067,6 +1193,11 @@ def add_articulation( Returns: Articulation: The added articulation instance handle. """ + if self.is_newton_backend: + logger.log_error( + "Newton articulation support is under development in DexSim and is not enabled in EmbodiChain yet.", + error_type=NotImplementedError, + ) uid = cfg.uid if uid is None: @@ -1147,6 +1278,11 @@ def add_robot(self, cfg: RobotCfg) -> Robot | None: Returns: Robot | None: The added robot instance handle, or None if failed. """ + if self.is_newton_backend: + logger.log_error( + "Newton robot support depends on DexSim Newton articulation support and is not enabled in EmbodiChain yet.", + error_type=NotImplementedError, + ) uid = cfg.uid if cfg.fpath is None: diff --git a/embodichain/lab/sim/utility/sim_utils.py b/embodichain/lab/sim/utility/sim_utils.py index 9a3f1eea..398cfe1c 100644 --- a/embodichain/lab/sim/utility/sim_utils.py +++ b/embodichain/lab/sim/utility/sim_utils.py @@ -26,7 +26,6 @@ LoadOption, RigidBodyShape, SDFConfig, - PhysicalAttr, ) from dexsim.engine import Articulation from dexsim.environment import Env, Arena @@ -44,6 +43,18 @@ import numpy as np +def _is_newton_backend_active() -> bool: + """Return whether the current default world uses the Newton physics scene.""" + from embodichain.lab.sim.objects.backends import is_newton_scene + + return is_newton_scene(dexsim.default_world().get_physics_scene()) + + +def _set_body_scale_after_rigidbody(obj: MeshObject, body_scale: tuple | list) -> None: + """Set body scale after rigid body creation for Newton compatibility.""" + obj.set_body_scale(*body_scale) + + def get_dexsim_arenas() -> List[dexsim.environment.Arena]: """Get all arenas in the default dexsim world. @@ -220,6 +231,7 @@ def load_mesh_objects_from_cfg( """ obj_list = [] body_type = cfg.to_dexsim_body_type() + is_newton_backend = _is_newton_backend_active() if isinstance(cfg.shape, MeshCfg): option = LoadOption() @@ -274,19 +286,27 @@ def load_mesh_objects_from_cfg( obj = env.load_actor( fpath, duplicate=True, attach_scene=True, option=option ) + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) sdf_cfg = SDFConfig() sdf_cfg.resolution = cfg.sdf_resolution obj.add_physical_body( body_type, RigidBodyShape.SDF, config=sdf_cfg, - attr=PhysicalAttr(), + attr=cfg.attrs.attr(), ) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) else: obj = env.load_actor( fpath, duplicate=True, attach_scene=True, option=option ) - obj.add_rigidbody(body_type, RigidBodyShape.CONVEX) + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.CONVEX, cfg.attrs.attr()) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) obj.set_name(f"{cfg.uid}_{i}") obj_list.append(obj) @@ -305,7 +325,11 @@ def load_mesh_objects_from_cfg( obj_list = create_cube(env_list, cfg.shape.size, uid=cfg.uid) for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.BOX) + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.BOX, cfg.attrs.attr()) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) elif isinstance(cfg.shape, SphereCfg): from embodichain.lab.sim.utility.sim_utils import create_sphere @@ -314,7 +338,11 @@ def load_mesh_objects_from_cfg( env_list, cfg.shape.radius, cfg.shape.resolution, uid=cfg.uid ) for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.SPHERE) + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.SPHERE, cfg.attrs.attr()) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) else: logger.log_error( f"Unsupported rigid object shape type: {type(cfg.shape)}. Supported types: MeshCfg, CubeCfg, SphereCfg." diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index b8f6c727..a104e831 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -38,6 +38,18 @@ def main(): description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--physics_backend", + choices=["default", "newton"], + default="default", + help="Physics backend to use for the simulation.", + ) + parser.add_argument( + "--max_steps", + type=int, + default=None, + help="Maximum number of simulation steps to run before exiting.", + ) args = parser.parse_args() # Configure the simulation @@ -47,6 +59,7 @@ def main(): headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) sim_device=args.device, + physics_backend=args.physics_backend, render_cfg=RenderCfg( renderer=args.renderer, ), @@ -98,10 +111,10 @@ def main(): sim.open_window() # Run the simulation - run_simulation(sim) + run_simulation(sim, max_steps=args.max_steps) -def run_simulation(sim: SimulationManager): +def run_simulation(sim: SimulationManager, max_steps: int | None = None): """Run the simulation loop. Args: @@ -122,6 +135,9 @@ def run_simulation(sim: SimulationManager): sim.update(step=1) step_count += 1 + if max_steps is not None and step_count >= max_steps: + break + # Print FPS every second if step_count % 100 == 0: current_time = time.time() diff --git a/tests/sim/objects/test_rigid_object.py b/tests/sim/objects/test_rigid_object.py index 5beebe26..c572db0f 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -15,21 +15,20 @@ # ---------------------------------------------------------------------------- import os -import torch + import pytest +import torch from embodichain.lab.sim import ( SimulationManager, SimulationManagerCfg, VisualMaterialCfg, ) -from embodichain.lab.sim.objects import RigidObject -from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg -from embodichain.lab.sim.shapes import MeshCfg from embodichain.data import get_data_path -from dexsim.types import ActorType - from embodichain.lab.sim.cfg import RenderCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import MeshCfg DUCK_PATH = "ToyDuck/toy_duck.glb" TABLE_PATH = "ShopTableSimple/shop_table_simple.ply" @@ -39,13 +38,18 @@ class BaseRigidObjectTest: - """Shared test logic for CPU and CUDA.""" + """Shared rigid object test logic across physics backends.""" - def setup_simulation(self, sim_device): + def setup_simulation(self, physics_backend: str): config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + headless=True, + sim_device="cpu", + num_envs=NUM_ARENAS, + physics_backend=physics_backend, + render_cfg=RenderCfg(renderer="hybrid"), ) self.sim = SimulationManager(config) + self.physics_backend = physics_backend self.sim.enable_physics(False) duck_path = get_data_path(DUCK_PATH) assert os.path.isfile(duck_path) @@ -80,10 +84,8 @@ def setup_simulation(self, sim_device): ), ) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): - self.sim.init_gpu_physics() - self.sim.enable_physics(True) + self.sim.prepare_physics() def test_is_static(self): """Test the is_static() method of duck, table, and chair objects.""" @@ -158,9 +160,10 @@ def test_local_pose_behavior(self): assert all( abs(x) < 1e-5 for x in table_xyz_after ), f"FAIL: Table moved unexpectedly: {table_xyz_after}" - assert torch.allclose( - chair_xyz_after, expected_chair_pos, atol=1e-5 - ), f"FAIL: Chair pose changed unexpectedly: {chair_xyz_after.tolist()}" + if self.physics_backend == "default": + assert torch.allclose( + chair_xyz_after, expected_chair_pos, atol=1e-5 + ), f"FAIL: Chair pose changed unexpectedly: {chair_xyz_after.tolist()}" def test_add_force_torque(self): """Test that add_force applies force correctly to the duck object.""" @@ -404,6 +407,9 @@ def test_physical_attributes(self): assert self.table.is_non_dynamic, "Static table should be is_non_dynamic" assert self.chair.is_non_dynamic, "Kinematic chair should be is_non_dynamic" + if self.physics_backend == "newton": + return + # 3. body_type assert self.duck.body_type == "dynamic" self.duck.set_body_type("kinematic") @@ -590,14 +596,14 @@ def teardown_method(self): gc.collect() -class TestRigidObjectCPU(BaseRigidObjectTest): +class TestRigidObjectDefaultBackend(BaseRigidObjectTest): def setup_method(self): - self.setup_simulation("cpu") + self.setup_simulation("default") -class TestRigidObjectCUDA(BaseRigidObjectTest): +class TestRigidObjectNewtonBackend(BaseRigidObjectTest): def setup_method(self): - self.setup_simulation("cuda") + self.setup_simulation("newton") if __name__ == "__main__":