From 78237317e2eec85cdeffa9c052aabc4f6f58fab7 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Thu, 21 May 2026 17:16:38 +0800 Subject: [PATCH 1/2] Add Newton physics backend support Integrate Newton-aware simulation config, manager, and rigid body adapters.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .github/workflows/main.yml | 2 +- embodichain/lab/gym/envs/base_env.py | 4 +- embodichain/lab/sim/cfg.py | 85 ++++++- .../lab/sim/objects/backends/__init__.py | 27 +++ .../lab/sim/objects/backends/newton.py | 174 ++++++++++++++ embodichain/lab/sim/objects/rigid_object.py | 215 +++++++++++++++--- .../lab/sim/objects/rigid_object_group.py | 161 ++++++++++--- embodichain/lab/sim/sim_manager.py | 162 +++++++++++-- embodichain/lab/sim/utility/sim_utils.py | 13 +- 9 files changed, 763 insertions(+), 80 deletions(-) create mode 100644 embodichain/lab/sim/objects/backends/__init__.py create mode 100644 embodichain/lab/sim/objects/backends/newton.py 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..aa395589 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,89 @@ 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"] = ( + "mjwarp" + ) + """Newton solver preset.""" + + 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..076bad73 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/__init__.py @@ -0,0 +1,27 @@ +# ---------------------------------------------------------------------------- +# 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 .newton import ( + NewtonRigidBodyView, + is_newton_scene, + newton_rigid_data_type, +) + +__all__ = [ + "NewtonRigidBodyView", + "is_newton_scene", + "newton_rigid_data_type", +] diff --git a/embodichain/lab/sim/objects/backends/newton.py b/embodichain/lab/sim/objects/backends/newton.py new file mode 100644 index 00000000..89ac3ae0 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -0,0 +1,174 @@ +# ---------------------------------------------------------------------------- +# 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.utils import logger + +_UINT64_MAX = (1 << 64) - 1 +_INT32_MAX = (1 << 31) - 1 + + +def newton_rigid_data_type(name: str): + from dexsim.engine.newton_physics.newton_physics_scene import NewtonRigidDataType + + return getattr(NewtonRigidDataType, name) + + +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: + """Thin 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. + """ + + 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(body_id < 0 or body_id > _INT32_MAX for body_id 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 + ) + + @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" + ) + + 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] + + 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 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._empty_warp((len(body_ids), 7)) + self.scene.gpu_fetch_rigid_body_data( + body_ids, + newton_rigid_data_type("POSE"), + out, + ) + return self._warp_to_torch(out) + + def apply_pose(self, pose: torch.Tensor, body_ids: Sequence[int]) -> None: + pose = pose.to(dtype=torch.float32) + self.scene.gpu_apply_rigid_body_data( + list(body_ids), + newton_rigid_data_type("POSE"), + self._to_numpy(pose), + ) + + 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._empty_warp((len(body_ids), 3)) + self.scene.gpu_fetch_rigid_body_data(body_ids, data_type, out) + return self._warp_to_torch(out) + + def apply_vec3( + self, data_type, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + self.scene.gpu_apply_rigid_body_data( + list(body_ids), + data_type, + self._to_numpy(data.to(dtype=torch.float32)), + ) + + def apply_force( + self, data_type, data: torch.Tensor, body_ids: Sequence[int] + ) -> None: + self.scene.gpu_apply_rigid_body_data( + list(body_ids), + data_type, + data.to(dtype=torch.float32, device=self.device), + ) + + def _empty_warp(self, shape: tuple[int, int]): + 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 _warp_to_torch(self, array) -> torch.Tensor: + 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 _to_numpy(self, tensor: torch.Tensor) -> np.ndarray: + return tensor.detach().cpu().numpy().astype(np.float32, copy=False) diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index 2202bbec..b4273296 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -26,6 +26,11 @@ from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType from dexsim.engine import CudaArray, PhysicsScene from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg +from embodichain.lab.sim.objects.backends import ( + NewtonRigidBodyView, + is_newton_scene, + newton_rigid_data_type, +) from embodichain.lab.sim import ( VisualMaterial, VisualMaterialInst, @@ -41,7 +46,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. + 1. The default DexSim GPU API stores pose as ``(qx, qy, qz, qw, x, y, z)``. + EmbodiChain and DexSim Newton use ``(x, y, z, qx, qy, qz, qw)``. """ def __init__( @@ -58,16 +64,28 @@ def __init__( self.ps = ps self.num_instances = len(entities) self.device = device + self._newton_view = ( + NewtonRigidBodyView(entities=entities, scene=ps, device=device) + if is_newton_scene(ps) + else None + ) # 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, + self._newton_view.body_ids_tensor + if self.is_newton_backend + else ( + torch.as_tensor( + [entity.get_gpu_index() for entity in self.entities], + dtype=torch.int32, + device=self.device, + ) + if self.device.type == "cuda" + else None ) - if self.device.type == "cuda" - else None + ) + self.newton_body_ids = ( + self._newton_view.body_ids if self.is_newton_backend else None ) # Initialize rigid body data. @@ -86,7 +104,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,8 +112,23 @@ def __init__( (self.num_instances, 7), dtype=torch.float32, device=self.device ) + @property + def is_newton_backend(self) -> bool: + return self._newton_view is not None + + @property + def is_newton_ready(self) -> bool: + return self._newton_view is not None and self._newton_view.is_ready + + def newton_body_ids_for(self, env_ids: Sequence[int]) -> list[int]: + return self._newton_view.select_body_ids(env_ids) + @property def pose(self) -> torch.Tensor: + if self.is_newton_ready: + self._pose = self._newton_view.fetch_pose() + return self._pose + if self.device.type == "cpu": # Fetch pose from CPU entities xyzs = torch.as_tensor( @@ -110,7 +143,6 @@ def pose(self) -> torch.Tensor: 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( @@ -118,12 +150,20 @@ def pose(self) -> torch.Tensor: gpu_indices=self.gpu_indices, data_type=RigidBodyGPUAPIReadType.POSE, ) - self._pose[:, :4] = convert_quat(self._pose[:, :4], to="wxyz") - self._pose = self._pose[:, [4, 5, 6, 0, 1, 2, 3]] + quat = self._pose[:, :4].clone() + xyz = self._pose[:, 4:7].clone() + self._pose[:, :3] = xyz + self._pose[:, 3:7] = quat return self._pose @property def lin_vel(self) -> torch.Tensor: + if self.is_newton_ready: + self._lin_vel = self._newton_view.fetch_vec3( + newton_rigid_data_type("LINEAR_VELOCITY") + ) + return self._lin_vel + if self.device.type == "cpu": # Fetch linear velocity from CPU entities self._lin_vel = torch.as_tensor( @@ -141,6 +181,12 @@ def lin_vel(self) -> torch.Tensor: @property def ang_vel(self) -> torch.Tensor: + if self.is_newton_ready: + self._ang_vel = self._newton_view.fetch_vec3( + newton_rigid_data_type("ANGULAR_VELOCITY") + ) + return self._ang_vel + if self.device.type == "cpu": # Fetch angular velocity from CPU entities self._ang_vel = torch.as_tensor( @@ -169,6 +215,12 @@ def vel(self) -> torch.Tensor: @property def lin_acc(self) -> torch.Tensor: + if self.is_newton_ready: + self._lin_acc = self._newton_view.fetch_vec3( + newton_rigid_data_type("LINEAR_ACCELERATION") + ) + return self._lin_acc + if self.device.type == "cpu": self._lin_acc = torch.as_tensor( np.array( @@ -187,6 +239,12 @@ def lin_acc(self) -> torch.Tensor: @property def ang_acc(self) -> torch.Tensor: + if self.is_newton_ready: + self._ang_acc = self._newton_view.fetch_vec3( + newton_rigid_data_type("ANGULAR_ACCELERATION") + ) + return self._ang_acc + if self.device.type == "cpu": self._ang_acc = torch.as_tensor( np.array( @@ -219,13 +277,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._newton_view.scene.manager + for i, entity_handle in enumerate(self._newton_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 +345,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 +359,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 +368,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 +418,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 +484,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,12 +506,34 @@ 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: + if self._data is not None and self._data.is_newton_ready and not self.is_static: + if pose.dim() == 2 and pose.shape[1] == 7: + newton_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") + newton_pose = torch.cat((xyz, quat), dim=-1) + else: + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." + ) + + body_ids = self._data.newton_body_ids_for(local_env_ids) + self._data._newton_view.apply_pose(newton_pose, body_ids) + return + + if ( + self.device.type == "cpu" + or self.is_static + or (self._data is not None and self._data.is_newton_backend) + ): 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]) + pose_matrix[:, :3, :3] = matrix_from_quat( + convert_quat(pose[:, 3:7], to="wxyz") + ) 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): @@ -441,7 +547,7 @@ def set_local_pose( else: if pose.dim() == 2 and pose.shape[1] == 7: xyz = pose[:, :3] - quat = convert_quat(pose[:, 3:7], to="xyzw") + quat = pose[:, 3:7] elif pose.dim() == 3 and pose.shape[1:] == (4, 4): xyz = pose[:, :3, 3] quat = quat_from_matrix(pose[:, :3, :3]) @@ -465,7 +571,7 @@ 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 +590,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 +600,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,7 +655,19 @@ 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.is_newton_ready: + body_ids = self._data.newton_body_ids_for(local_env_ids) + if force is not None: + self._data._newton_view.apply_force( + newton_rigid_data_type("FORCE"), force, body_ids + ) + if torque is not None: + self._data._newton_view.apply_force( + newton_rigid_data_type("TORQUE"), torque, body_ids + ) + elif self.device.type == "cpu" or ( + self._data is not None and self._data.is_newton_backend + ): for i, env_idx in enumerate(local_env_ids): if force is not None: self._entities[env_idx].add_force(force[i].cpu().numpy()) @@ -608,7 +725,19 @@ 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.is_newton_ready: + body_ids = self._data.newton_body_ids_for(local_env_ids) + if lin_vel is not None: + self._data._newton_view.apply_vec3( + newton_rigid_data_type("LINEAR_VELOCITY"), lin_vel, body_ids + ) + if ang_vel is not None: + self._data._newton_view.apply_vec3( + newton_rigid_data_type("ANGULAR_VELOCITY"), ang_vel, body_ids + ) + elif self.device.type == "cpu" or ( + self._data is not None and self._data.is_newton_backend + ): for i, env_idx in enumerate(local_env_ids): if lin_vel is not None: self._entities[env_idx].set_linear_velocity( @@ -941,7 +1070,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 +1092,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,7 +1215,26 @@ 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": + if self._data is not None and self._data.is_newton_ready: + zeros = torch.zeros( + (len(local_env_ids), 3), dtype=torch.float32, device=self.device + ) + body_ids = self._data.newton_body_ids_for(local_env_ids) + self._data._newton_view.apply_vec3( + newton_rigid_data_type("LINEAR_VELOCITY"), zeros, body_ids + ) + self._data._newton_view.apply_vec3( + newton_rigid_data_type("ANGULAR_VELOCITY"), zeros, body_ids + ) + self._data._newton_view.apply_force( + newton_rigid_data_type("FORCE"), zeros, body_ids + ) + self._data._newton_view.apply_force( + newton_rigid_data_type("TORQUE"), zeros, body_ids + ) + elif self._data is not None and self._data.is_newton_backend: + return + elif self.device.type == "cpu": for env_idx in local_env_ids: self._entities[env_idx].clear_dynamics() else: diff --git a/embodichain/lab/sim/objects/rigid_object_group.py b/embodichain/lab/sim/objects/rigid_object_group.py index e4cca592..4dc7f630 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/embodichain/lab/sim/objects/rigid_object_group.py @@ -28,6 +28,11 @@ RigidObjectGroupCfg, RigidBodyAttributesCfg, ) +from embodichain.lab.sim.objects.backends import ( + NewtonRigidBodyView, + is_newton_scene, + newton_rigid_data_type, +) from embodichain.lab.sim import ( BatchEntity, ) @@ -56,19 +61,34 @@ 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] + self._newton_view = ( + NewtonRigidBodyView(entities=self.flat_entities, scene=ps, device=device) + if is_newton_scene(ps) + else None + ) # 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, + self._newton_view.body_ids_tensor.reshape( + self.num_instances, self.num_objects + ) + if self.is_newton_backend + else ( + 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 ) - if self.device.type == "cuda" - else None + ) + self.newton_body_ids = ( + self._newton_view.body_ids if self.is_newton_backend else None ) # Initialize rigid body group data tensors. Shape of (num_instances, num_objects, data_dim) @@ -88,8 +108,35 @@ def __init__( device=self.device, ) + @property + def is_newton_backend(self) -> bool: + return self._newton_view is not None + + @property + def is_newton_ready(self) -> bool: + return self._newton_view is not None and self._newton_view.is_ready + + def newton_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 + body_ids = [] + for env_idx in env_ids: + for obj_idx in local_obj_ids: + flat_index = int(env_idx) * self.num_objects + int(obj_idx) + body_ids.append(self.newton_body_ids[flat_index]) + return body_ids + @property def pose(self) -> torch.Tensor: + if self.is_newton_ready: + self._pose = self._newton_view.fetch_pose().reshape( + self.num_instances, self.num_objects, 7 + ) + return self._pose + if self.device.type == "cpu": # Fetch pose from CPU entities xyzs = torch.as_tensor( @@ -97,6 +144,7 @@ def pose(self) -> torch.Tensor: [entity.get_location() for entity in instance] for instance in self.entities ], + dtype=torch.float32, device=self.device, ) quats = torch.as_tensor( @@ -104,12 +152,10 @@ def pose(self) -> torch.Tensor: [entity.get_rotation_quat() for entity in instance] for instance in self.entities ], + dtype=torch.float32, 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) + self._pose = torch.cat((xyzs, quats), dim=-1) else: pose = self._pose.reshape(-1, 7) self.ps.gpu_fetch_rigid_body_data( @@ -117,12 +163,20 @@ def pose(self) -> torch.Tensor: 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 + quat = pose[:, :4].clone() + xyz = pose[:, 4:7].clone() + pose[:, :3] = xyz + pose[:, 3:7] = quat + return self._pose @property def lin_vel(self) -> torch.Tensor: + if self.is_newton_ready: + self._lin_vel = self._newton_view.fetch_vec3( + newton_rigid_data_type("LINEAR_VELOCITY") + ).reshape(self.num_instances, self.num_objects, 3) + return self._lin_vel + if self.device.type == "cpu": # Fetch linear velocity from CPU entities self._lin_vel = torch.as_tensor( @@ -144,11 +198,17 @@ def lin_vel(self) -> torch.Tensor: @property def ang_vel(self) -> torch.Tensor: + if self.is_newton_ready: + self._ang_vel = self._newton_view.fetch_vec3( + newton_rigid_data_type("ANGULAR_VELOCITY") + ).reshape(self.num_instances, self.num_objects, 3) + return self._ang_vel + 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] + [entity.get_angular_velocity() for entity in instance] for instance in self.entities ], dtype=torch.float32, @@ -198,10 +258,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 +305,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 +359,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,7 +388,27 @@ 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": + if self._data.is_newton_ready: + if pose.dim() == 3 and pose.shape[2] == 7: + xyz = pose[..., :3].reshape(-1, 3) + quat = pose[..., 3:7].reshape(-1, 4) + 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") + else: + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, M, 7) or (N, M, 4, 4)." + ) + + newton_pose = torch.cat((xyz, quat), dim=-1).to( + device=self.device, dtype=torch.float32 + ) + body_ids = self._data.newton_body_ids_for(local_env_ids, local_obj_ids) + self._data._newton_view.apply_pose(newton_pose, body_ids) + return + + if self.device.type == "cpu" or self._data.is_newton_backend: pose = pose.cpu() if pose.dim() == 3 and pose.shape[2] == 7: reshape_pose = pose.reshape(-1, 7) @@ -329,7 +416,9 @@ def set_local_pose( 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_matrix[:, :3, :3] = matrix_from_quat( + convert_quat(reshape_pose[:, 3:7], to="wxyz") + ) pose = pose_matrix.reshape(-1, len(local_obj_ids), 4, 4) elif pose.dim() == 4 and pose.shape[2:] == (4, 4): pass @@ -346,7 +435,6 @@ def set_local_pose( 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) @@ -376,7 +464,7 @@ 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 +473,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,7 +510,28 @@ 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": + if self._data.is_newton_ready: + zeros = torch.zeros( + (len(local_env_ids) * self.num_objects, 3), + dtype=torch.float32, + device=self.device, + ) + body_ids = self._data.newton_body_ids_for(local_env_ids) + self._data._newton_view.apply_vec3( + newton_rigid_data_type("LINEAR_VELOCITY"), zeros, body_ids + ) + self._data._newton_view.apply_vec3( + newton_rigid_data_type("ANGULAR_VELOCITY"), zeros, body_ids + ) + self._data._newton_view.apply_force( + newton_rigid_data_type("FORCE"), zeros, body_ids + ) + self._data._newton_view.apply_force( + newton_rigid_data_type("TORQUE"), zeros, body_ids + ) + elif self._data.is_newton_backend: + return + elif self.device.type == "cpu": for env_idx in local_env_ids: for entity in self._entities[env_idx]: entity.clear_dynamics() 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..a56acc28 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 @@ -274,19 +273,21 @@ def load_mesh_objects_from_cfg( obj = env.load_actor( fpath, duplicate=True, attach_scene=True, option=option ) + 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(), ) else: obj = env.load_actor( fpath, duplicate=True, attach_scene=True, option=option ) - obj.add_rigidbody(body_type, RigidBodyShape.CONVEX) + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.CONVEX, cfg.attrs.attr()) obj.set_name(f"{cfg.uid}_{i}") obj_list.append(obj) @@ -305,7 +306,8 @@ 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) + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.BOX, cfg.attrs.attr()) elif isinstance(cfg.shape, SphereCfg): from embodichain.lab.sim.utility.sim_utils import create_sphere @@ -314,7 +316,8 @@ 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) + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.SPHERE, cfg.attrs.attr()) else: logger.log_error( f"Unsupported rigid object shape type: {type(cfg.shape)}. Supported types: MeshCfg, CubeCfg, SphereCfg." From 178c730dd5dcd561eca5ca256a3d12b79973dfe9 Mon Sep 17 00:00:00 2001 From: yuecideng Date: Fri, 22 May 2026 16:02:42 +0800 Subject: [PATCH 2/2] wip --- embodichain/lab/sim/cfg.py | 8 +- .../lab/sim/objects/backends/__init__.py | 11 +- embodichain/lab/sim/objects/backends/base.py | 128 ++++++ .../lab/sim/objects/backends/default.py | 283 +++++++++++++ .../lab/sim/objects/backends/newton.py | 173 +++++--- embodichain/lab/sim/objects/rigid_object.py | 380 +++++------------- .../lab/sim/objects/rigid_object_group.py | 308 +++++--------- embodichain/lab/sim/utility/sim_utils.py | 33 +- scripts/tutorials/sim/create_scene.py | 20 +- tests/sim/objects/test_rigid_object.py | 44 +- 10 files changed, 803 insertions(+), 585 deletions(-) create mode 100644 embodichain/lab/sim/objects/backends/base.py create mode 100644 embodichain/lab/sim/objects/backends/default.py diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index aa395589..84bc93fb 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -148,10 +148,16 @@ class NewtonPhysicsCfg: """Whether to enable Newton debug mode.""" solver_type: Literal["mjwarp", "xpbd", "semi_implicit", "featherstone", "vbd"] = ( - "mjwarp" + "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, diff --git a/embodichain/lab/sim/objects/backends/__init__.py b/embodichain/lab/sim/objects/backends/__init__.py index 076bad73..e65f00a5 100644 --- a/embodichain/lab/sim/objects/backends/__init__.py +++ b/embodichain/lab/sim/objects/backends/__init__.py @@ -14,14 +14,13 @@ # limitations under the License. # ---------------------------------------------------------------------------- -from .newton import ( - NewtonRigidBodyView, - is_newton_scene, - newton_rigid_data_type, -) +from .base import RigidBodyViewBase +from .default import DefaultRigidBodyView +from .newton import NewtonRigidBodyView, is_newton_scene __all__ = [ + "RigidBodyViewBase", + "DefaultRigidBodyView", "NewtonRigidBodyView", "is_newton_scene", - "newton_rigid_data_type", ] 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 index 89ac3ae0..122c44c7 100644 --- a/embodichain/lab/sim/objects/backends/newton.py +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- - from __future__ import annotations from typing import Sequence @@ -23,18 +22,15 @@ 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 newton_rigid_data_type(name: str): - from dexsim.engine.newton_physics.newton_physics_scene import NewtonRigidDataType - - return getattr(NewtonRigidDataType, name) - - def _normalize_native_handle(handle: int, owner: str) -> int: value = int(handle) if value < 0: @@ -54,8 +50,8 @@ def is_newton_scene(scene: object) -> bool: ) -class NewtonRigidBodyView: - """Thin adapter around DexSim Newton rigid body scene APIs. +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)``. @@ -63,6 +59,8 @@ class NewtonRigidBodyView: data API. """ + _DATA_TYPE = None # lazily resolved NewtonRigidDataType + def __init__( self, entities: Sequence[MeshObject], @@ -76,15 +74,28 @@ def __init__( _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(body_id < 0 or body_id > _INT32_MAX for body_id in self.body_ids): + 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 + 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) @@ -94,10 +105,75 @@ def is_ready(self) -> bool: == "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] + 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) @@ -115,60 +191,33 @@ def _resolve_body_id(self, entity: MeshObject) -> int: return body_id return -1 - 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._empty_warp((len(body_ids), 7)) - self.scene.gpu_fetch_rigid_body_data( - body_ids, - newton_rigid_data_type("POSE"), - out, - ) - return self._warp_to_torch(out) - - def apply_pose(self, pose: torch.Tensor, body_ids: Sequence[int]) -> None: - pose = pose.to(dtype=torch.float32) - self.scene.gpu_apply_rigid_body_data( - list(body_ids), - newton_rigid_data_type("POSE"), - self._to_numpy(pose), - ) - - 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._empty_warp((len(body_ids), 3)) - self.scene.gpu_fetch_rigid_body_data(body_ids, data_type, out) - return self._warp_to_torch(out) - - def apply_vec3( - self, data_type, data: torch.Tensor, body_ids: Sequence[int] - ) -> None: - self.scene.gpu_apply_rigid_body_data( - list(body_ids), - data_type, - self._to_numpy(data.to(dtype=torch.float32)), - ) - - def apply_force( - self, data_type, data: torch.Tensor, body_ids: Sequence[int] - ) -> None: - self.scene.gpu_apply_rigid_body_data( - list(body_ids), - data_type, - data.to(dtype=torch.float32, device=self.device), - ) - - def _empty_warp(self, shape: tuple[int, int]): + 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 _warp_to_torch(self, array) -> torch.Tensor: + 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 _to_numpy(self, tensor: torch.Tensor) -> np.ndarray: - return tensor.detach().cpu().numpy().astype(np.float32, copy=False) + 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 b4273296..9d2aa924 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -23,14 +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, - newton_rigid_data_type, ) +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase from embodichain.lab.sim import ( VisualMaterial, VisualMaterialInst, @@ -45,9 +45,8 @@ class RigidBodyData: """Data manager for rigid body with body type of dynamic or kinematic. - Note: - 1. The default DexSim GPU API stores pose as ``(qx, qy, qz, qw, x, y, z)``. - EmbodiChain and DexSim Newton use ``(x, y, z, qx, qy, qz, qw)``. + All pose/velocity/acceleration data uses EmbodiChain convention: + ``(x, y, z, qx, qy, qz, qw)``. """ def __init__( @@ -64,29 +63,19 @@ def __init__( self.ps = ps self.num_instances = len(entities) self.device = device - self._newton_view = ( - NewtonRigidBodyView(entities=entities, scene=ps, device=device) - if is_newton_scene(ps) - else None - ) - # get gpu indices for the entities. - self.gpu_indices = ( - self._newton_view.body_ids_tensor - if self.is_newton_backend - else ( - torch.as_tensor( - [entity.get_gpu_index() for entity in self.entities], - dtype=torch.int32, - device=self.device, - ) - if self.device.type == "cuda" - else None + # Create the appropriate backend view. + if is_newton_scene(ps): + self._body_view: RigidBodyViewBase = NewtonRigidBodyView( + entities=entities, scene=ps, device=device ) - ) - self.newton_body_ids = ( - self._newton_view.body_ids if self.is_newton_backend 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( @@ -114,93 +103,54 @@ def __init__( @property def is_newton_backend(self) -> bool: - return self._newton_view is not None + return isinstance(self._body_view, NewtonRigidBodyView) @property def is_newton_ready(self) -> bool: - return self._newton_view is not None and self._newton_view.is_ready + return self.is_newton_backend and self._body_view.is_ready - def newton_body_ids_for(self, env_ids: Sequence[int]) -> list[int]: - return self._newton_view.select_body_ids(env_ids) + 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.is_newton_ready: - self._pose = self._newton_view.fetch_pose() + if self._body_view.is_ready: + self._pose = self._body_view.fetch_pose() return self._pose - 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, + # 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 ) - 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 ) - quat = self._pose[:, :4].clone() - xyz = self._pose[:, 4:7].clone() - self._pose[:, :3] = xyz - self._pose[:, 3:7] = quat return self._pose @property def lin_vel(self) -> torch.Tensor: - if self.is_newton_ready: - self._lin_vel = self._newton_view.fetch_vec3( - newton_rigid_data_type("LINEAR_VELOCITY") - ) + if self._body_view.is_ready: + self._lin_vel = self._body_view.fetch_linear_velocity() return self._lin_vel - 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, + 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.is_newton_ready: - self._ang_vel = self._newton_view.fetch_vec3( - newton_rigid_data_type("ANGULAR_VELOCITY") - ) + if self._body_view.is_ready: + self._ang_vel = self._body_view.fetch_angular_velocity() return self._ang_vel - 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, + 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 @@ -215,50 +165,30 @@ def vel(self) -> torch.Tensor: @property def lin_acc(self) -> torch.Tensor: - if self.is_newton_ready: - self._lin_acc = self._newton_view.fetch_vec3( - newton_rigid_data_type("LINEAR_ACCELERATION") - ) + if self._body_view.is_ready: + self._lin_acc = self._body_view.fetch_linear_acceleration() return self._lin_acc - if self.device.type == "cpu": - self._lin_acc = torch.as_tensor( - np.array( - [entity.get_linear_acceleration() for entity in self.entities], - ), + 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.is_newton_ready: - self._ang_acc = self._newton_view.fetch_vec3( - newton_rigid_data_type("ANGULAR_ACCELERATION") - ) + if self._body_view.is_ready: + self._ang_acc = self._body_view.fetch_angular_acceleration() return self._ang_acc - if self.device.type == "cpu": - self._ang_acc = torch.as_tensor( - np.array( - [entity.get_angular_acceleration() for entity in self.entities], - ), + 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 @@ -278,8 +208,8 @@ def com_pose(self) -> torch.Tensor: torch.Tensor: The center of mass pose with shape (N, 7). """ if self.is_newton_backend: - manager = self._newton_view.scene.manager - for i, entity_handle in enumerate(self._newton_view.entity_handles): + 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) @@ -506,66 +436,41 @@ def set_local_pose( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." ) - if self._data is not None and self._data.is_newton_ready and not self.is_static: - if pose.dim() == 2 and pose.shape[1] == 7: - newton_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") - newton_pose = torch.cat((xyz, quat), dim=-1) - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - body_ids = self._data.newton_body_ids_for(local_env_ids) - self._data._newton_view.apply_pose(newton_pose, body_ids) + # 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: + 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.device.type == "cpu" - or self.is_static - or (self._data is not None and self._data.is_newton_backend) + self._data is not None + and self._data._body_view.is_ready + and not 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( - convert_quat(pose[:, 3:7], to="wxyz") - ) - 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)." - ) - - else: - if pose.dim() == 2 and pose.shape[1] == 7: - xyz = pose[:, :3] - quat = pose[:, 3:7] - 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)." - ) + body_ids = self._data.body_ids_for(local_env_ids) + self._data._body_view.apply_pose(target_pose, body_ids) + return - # 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, - ) + # 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. @@ -655,41 +560,21 @@ def add_force_torque( f"Length of env_ids {len(local_env_ids)} does not match torque length {len(torque)}." ) - if self._data is not None and self._data.is_newton_ready: - body_ids = self._data.newton_body_ids_for(local_env_ids) + 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._newton_view.apply_force( - newton_rigid_data_type("FORCE"), force, body_ids - ) + self._data._body_view.apply_force(force, body_ids) if torque is not None: - self._data._newton_view.apply_force( - newton_rigid_data_type("TORQUE"), torque, body_ids - ) - elif self.device.type == "cpu" or ( - self._data is not None and self._data.is_newton_backend - ): + 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, @@ -725,19 +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._data is not None and self._data.is_newton_ready: - body_ids = self._data.newton_body_ids_for(local_env_ids) + 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._newton_view.apply_vec3( - newton_rigid_data_type("LINEAR_VELOCITY"), lin_vel, body_ids - ) + self._data._body_view.apply_linear_velocity(lin_vel, body_ids) if ang_vel is not None: - self._data._newton_view.apply_vec3( - newton_rigid_data_type("ANGULAR_VELOCITY"), ang_vel, body_ids - ) - elif self.device.type == "cpu" or ( - self._data is not None and self._data.is_newton_backend - ): + 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( @@ -747,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, @@ -1215,55 +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._data is not None and self._data.is_newton_ready: + 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 ) - body_ids = self._data.newton_body_ids_for(local_env_ids) - self._data._newton_view.apply_vec3( - newton_rigid_data_type("LINEAR_VELOCITY"), zeros, body_ids - ) - self._data._newton_view.apply_vec3( - newton_rigid_data_type("ANGULAR_VELOCITY"), zeros, body_ids - ) - self._data._newton_view.apply_force( - newton_rigid_data_type("FORCE"), zeros, body_ids - ) - self._data._newton_view.apply_force( - newton_rigid_data_type("TORQUE"), zeros, body_ids - ) + 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 - elif self.device.type == "cpu": + else: for env_idx in local_env_ids: self._entities[env_idx].clear_dynamics() - else: - # Apply zero force and torque to the rigid bodies. - 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, - ) def set_physical_visible( self, @@ -1343,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 4dc7f630..dc2007a0 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/embodichain/lab/sim/objects/rigid_object_group.py @@ -22,17 +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, - newton_rigid_data_type, ) +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase from embodichain.lab.sim import ( BatchEntity, ) @@ -62,33 +62,20 @@ def __init__( self.num_objects = len(entities[0]) self.device = device self.flat_entities = [entity for instance in entities for entity in instance] - self._newton_view = ( - NewtonRigidBodyView(entities=self.flat_entities, scene=ps, device=device) - if is_newton_scene(ps) - else None - ) - # get gpu indices for the rigid bodies with shape of (num_instances, num_objects) - self.gpu_indices = ( - self._newton_view.body_ids_tensor.reshape( - self.num_instances, self.num_objects + # Create the appropriate backend view. + if is_newton_scene(ps): + self._body_view: RigidBodyViewBase = NewtonRigidBodyView( + entities=self.flat_entities, scene=ps, device=device ) - if self.is_newton_backend - else ( - 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 + else: + self._body_view = DefaultRigidBodyView( + entities=self.flat_entities, ps=ps, device=device ) - ) - self.newton_body_ids = ( - self._newton_view.body_ids if self.is_newton_backend 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) @@ -110,117 +97,75 @@ def __init__( @property def is_newton_backend(self) -> bool: - return self._newton_view is not None + return isinstance(self._body_view, NewtonRigidBodyView) @property def is_newton_ready(self) -> bool: - return self._newton_view is not None and self._newton_view.is_ready + return self.is_newton_backend and self._body_view.is_ready - def newton_body_ids_for( + 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 - body_ids = [] + flat_indices = [] for env_idx in env_ids: for obj_idx in local_obj_ids: - flat_index = int(env_idx) * self.num_objects + int(obj_idx) - body_ids.append(self.newton_body_ids[flat_index]) - return body_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.is_newton_ready: - self._pose = self._newton_view.fetch_pose().reshape( + if self._body_view.is_ready: + self._pose = self._body_view.fetch_pose().reshape( self.num_instances, self.num_objects, 7 ) return self._pose - 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 - ], - dtype=torch.float32, - device=self.device, - ) - quats = torch.as_tensor( - [ - [entity.get_rotation_quat() for entity in instance] - for instance in self.entities - ], - dtype=torch.float32, - device=self.device, - ) - self._pose = 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, - ) - quat = pose[:, :4].clone() - xyz = pose[:, 4:7].clone() - pose[:, :3] = xyz - pose[:, 3:7] = quat + # 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.is_newton_ready: - self._lin_vel = self._newton_view.fetch_vec3( - newton_rigid_data_type("LINEAR_VELOCITY") - ).reshape(self.num_instances, self.num_objects, 3) + 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 - 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, - ) + 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.is_newton_ready: - self._ang_vel = self._newton_view.fetch_vec3( - newton_rigid_data_type("ANGULAR_VELOCITY") - ).reshape(self.num_instances, self.num_objects, 3) + 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 - 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, - ) + 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 @@ -388,77 +333,41 @@ def set_local_pose( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." ) - if self._data.is_newton_ready: - if pose.dim() == 3 and pose.shape[2] == 7: - xyz = pose[..., :3].reshape(-1, 3) - quat = pose[..., 3:7].reshape(-1, 4) - 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") - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, M, 7) or (N, M, 4, 4)." - ) - - newton_pose = torch.cat((xyz, quat), dim=-1).to( + # 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 ) - body_ids = self._data.newton_body_ids_for(local_env_ids, local_obj_ids) - self._data._newton_view.apply_pose(newton_pose, body_ids) + 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 + ) + else: + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, M, 7) or (N, M, 4, 4)." + ) return - if self.device.type == "cpu" or self._data.is_newton_backend: - 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( - convert_quat(reshape_pose[:, 3:7], to="wxyz") - ) - 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) - 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)." - ) + # 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 - # 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, - ) - self._world.sync_poses_gpu_to_cpu( - rigid_pose=CudaArray(pose), rigid_gpu_indices=CudaArray(indices) - ) + # 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. @@ -510,60 +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._data.is_newton_ready: + if self._data._body_view.is_ready: zeros = torch.zeros( (len(local_env_ids) * self.num_objects, 3), dtype=torch.float32, device=self.device, ) - body_ids = self._data.newton_body_ids_for(local_env_ids) - self._data._newton_view.apply_vec3( - newton_rigid_data_type("LINEAR_VELOCITY"), zeros, body_ids - ) - self._data._newton_view.apply_vec3( - newton_rigid_data_type("ANGULAR_VELOCITY"), zeros, body_ids - ) - self._data._newton_view.apply_force( - newton_rigid_data_type("FORCE"), zeros, body_ids - ) - self._data._newton_view.apply_force( - newton_rigid_data_type("TORQUE"), zeros, body_ids - ) + 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 - elif self.device.type == "cpu": + else: 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, - ) def set_visual_material( self, mat: VisualMaterial, env_ids: Sequence[int] | None = None diff --git a/embodichain/lab/sim/utility/sim_utils.py b/embodichain/lab/sim/utility/sim_utils.py index a56acc28..398cfe1c 100644 --- a/embodichain/lab/sim/utility/sim_utils.py +++ b/embodichain/lab/sim/utility/sim_utils.py @@ -43,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. @@ -219,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() @@ -273,7 +286,8 @@ def load_mesh_objects_from_cfg( obj = env.load_actor( fpath, duplicate=True, attach_scene=True, option=option ) - obj.set_body_scale(*cfg.body_scale) + 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( @@ -282,12 +296,17 @@ def load_mesh_objects_from_cfg( config=sdf_cfg, 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.set_body_scale(*cfg.body_scale) + 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) @@ -306,8 +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.set_body_scale(*cfg.body_scale) + 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 @@ -316,8 +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.set_body_scale(*cfg.body_scale) + 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__":