diff --git a/design/newton-backend-design.md b/design/newton-backend-design.md new file mode 100644 index 00000000..9b27b0a2 --- /dev/null +++ b/design/newton-backend-design.md @@ -0,0 +1,288 @@ +# EmbodiChain Newton Backend Integration Design + +This document records the current EmbodiChain integration state for the DexSim +Newton physics backend and the remaining work needed to complete it. + +Use these EmbodiChain backend names consistently: + +- `default`: the existing DexSim default physics backend. +- `newton`: the DexSim Newton physics backend. + +Avoid exposing lower-level DexSim implementation names in EmbodiChain-facing +configuration, docs, and conditionals. + +## Current State + +### Configuration + +Backend selection is currently inferred from `SimulationManagerCfg.physics_cfg`: + +- `DefaultPhysicsCfg` selects the `default` backend. +- `NewtonPhysicsCfg` selects the `newton` backend. +- `physics_cfg_for_backend("default" | "newton")` returns the matching config. +- `physics_backend_from_cfg(...)` maps a config instance to its backend name. + +`DefaultPhysicsCfg` owns default-backend PhysX settings and GPU-memory settings. +`NewtonPhysicsCfg` owns Newton settings: `physics_dt`, `device`, `num_substeps`, +`requires_grad`, `use_cuda_graph`, `debug_mode`, `solver_type`, `broad_phase`, +and `visualizer_enabled`. + +`NewtonPhysicsCfg.to_dexsim_cfg(...)` creates a DexSim `NewtonCfg`, uses +`physics_dt` for `NewtonCfg.dt`, disables CUDA graph when gradient mode is +enabled, and requires `solver_type="semi_implicit"` for gradient mode. + +### SimulationManager + +`SimulationManager` now tracks the active backend with: + +- `physics_backend` +- `is_default_backend` +- `is_newton_backend` +- `newton_manager` + +For the `default` backend, manager initialization keeps the existing DexSim +behavior: + +- apply `DefaultPhysicsCfg.to_dexsim_args()` +- apply default-backend GPU-memory config +- enable default GPU simulation only when the selected device is CUDA + +For the `newton` backend, manager initialization: + +- imports DexSim Newton lazily during world-config conversion +- sets `world_config.newton_cfg` +- obtains the per-world Newton manager through `get_newton_manager(self._world)` +- avoids default-backend GPU flags and default GPU memory APIs + +Newton finalization is separate from default-backend GPU initialization: + +- `finalize_newton_physics()` prepares or rebuilds the Newton model until the + manager reaches `READY`. +- `update(...)` finalizes Newton before stepping. +- `init_gpu_physics()` delegates to `finalize_newton_physics()` when Newton is + active. +- `set_manual_update(False)` is ignored for Newton because the backend does not + support switching to automatic update. + +Scene mutation invalidates Newton finalization with `_invalidate_newton_physics()`. +After finalization, `_reset_newton_entities_after_finalize()` reapplies rigid +object reset state. Rigid object groups are not yet supported on Newton. + +### Object Backend Adapters + +Rigid-body data access is routed through: + +```text +embodichain/lab/sim/objects/backends/ + base.py + default.py + newton.py +``` + +`RigidBodyViewBase` defines the backend-neutral rigid-body API. The default +adapter handles existing CPU/default-GPU paths. The Newton adapter uses DexSim +Newton batch APIs for body data and collision filters. + +EmbodiChain public rigid-body tensor convention is: + +```text +(x, y, z, qx, qy, qz, qw) +``` + +Current Newton rigid-object support includes: + +- dynamic and kinematic single `RigidObject` creation +- static single `RigidObject` creation +- local pose get/set +- body state get +- linear/angular velocity get/set +- linear/angular acceleration get +- force and torque at center of mass +- clear dynamics +- reset +- center-of-mass local pose get/set for dynamic rigid objects +- mass get/set +- friction get/set +- inertia diagonal get/set +- collision filter set for dynamic, kinematic, static, and pre-finalize bodies +- visual material, visibility, geometry, scale, and user-id APIs through the + existing MeshObject paths + +Static Newton bodies do not have `RigidBodyData`; static collision-filter writes +therefore use DexSim's per-entity metadata hook when a Newton body ID is not +available yet. + +### Currently Unsupported Newton APIs + +`SimulationManager` explicitly rejects these asset types on Newton: + +- `add_soft_object(...)` +- `add_cloth_object(...)` +- `add_rigid_object_group(...)` +- `add_articulation(...)` +- `add_robot(...)` + +`RigidObject` still does not support these runtime updates on Newton: + +- `set_attrs(...)` +- `set_body_type(...)` +- `set_damping(...)` + +`RigidObject.add_force_torque(pos=...)` ignores `pos` and applies force/torque at +the center of mass. + +Newton kinematic pose locking is not complete. The rigid-object test suite keeps +a Newton-specific allowance for kinematic bodies changing after stepping. + +Newton SDF rigid mesh support is not validated in EmbodiChain. The SDF rigid +object test is skipped for Newton. + +### Verified Tests + +The current rigid-object test file passes after the latest Newton integration +fixes: + +```bash +pytest -q tests/sim/objects/test_rigid_object.py +``` + +Observed result: + +```text +62 passed, 1 skipped, 41 warnings +``` + +## Improvements To Make + +### API Clarity + +- Add explicit capability checks for backend-specific support instead of relying + on scattered `is_newton_scene(...)` checks. +- Make unsupported Newton APIs fail consistently with either `NotImplementedError` + or a documented warning/no-op policy. +- Separate `is_use_gpu_physics` into clearer concepts: + - selected tensor/device location + - default-backend GPU API availability + - Newton GPU execution + +### Newton Lifecycle + +- Keep `finalize_newton_physics()` as the single Newton preparation API. +- Do not add a separate non-stepping synchronization method until DexSim exposes + a real Newton synchronization API. +- Track dirty scene/model state more explicitly so mutations after finalization + can choose between live batch updates and model rebuilds. +- Avoid global Newton teardown while another world may still use monkey-patched + DexSim classes. + +### RigidObject + +- Implement Newton `set_attrs(...)` by decomposing supported fields into batch + property updates and rejecting unsupported fields explicitly. +- Implement Newton damping get/set through DexSim Newton if a runtime API exists; + otherwise keep it metadata-only before finalization and document that runtime + damping changes require rebuild. +- Implement `set_body_type(...)` for Newton or keep a hard unsupported error if + DexSim cannot safely switch dynamic/kinematic/static bodies at runtime. +- Implement force-at-position when DexSim Newton exposes the needed API. +- Validate SDF rigid mesh creation and collision behavior on Newton. +- Fix or document kinematic pose-lock semantics. + +### Object Groups, Articulations, Robots, Soft, Cloth + +- Add Newton rigid-object-group support after single-object support is stable. +- Keep articulations and robots fail-fast until DexSim Newton articulation APIs + are ready and tested. +- Keep soft and cloth fail-fast until there is an explicit Newton design and + test coverage for those object types. + +### Gym Env Integration + +Use backend-specific initialization in env setup: + +```python +if self.sim.is_default_backend and self.sim.is_use_gpu_physics: + self.sim.init_gpu_physics() +elif self.sim.is_newton_backend: + self.sim.finalize_newton_physics() +``` + +For stepping, keep the existing high-level flow: + +```python +self._preprocess_action(action) +self._step_action(action) +self.sim.update(self.sim_cfg.physics_dt, self.cfg.sim_steps_per_control) +``` + +For reset, call object/manager reset methods and finalize Newton before reading +observations when the backend is Newton. Do not rely on a separate sync API. + +## Completion Plan + +1. Stabilize the single-rigid-object Newton API and keep + `tests/sim/objects/test_rigid_object.py` green. +2. Add backend capability declarations and use them in public object APIs. +3. Finish Newton `RigidObject` parity for attributes, damping, body type, + force-at-position, SDF meshes, and kinematic pose semantics. +4. Add tests for Newton lifecycle rebuild after scene mutation and runtime + property mutation after finalization. +5. Implement and test Newton `RigidObjectGroup`. +6. Update gym env initialization/reset paths to use `finalize_newton_physics()` + directly. +7. Add rigid-only Newton gym smoke tests. +8. Add gradient rollout wrapper and a minimal differentiable Newton smoke test. +9. Add articulation and robot support only after DexSim Newton exposes stable + articulation APIs. +10. Add soft/cloth support only after a dedicated Newton object design and tests. + +## Tests To Maintain + +Configuration: + +- `SimulationManagerCfg(physics_cfg=DefaultPhysicsCfg())` preserves current + default-backend behavior. +- `SimulationManagerCfg(physics_cfg=NewtonPhysicsCfg())` creates a Newton world. +- `physics_cfg_for_backend(...)` and `physics_backend_from_cfg(...)` return the + expected backend mapping. + +Simulation: + +- Newton world can be created, finalized, stepped, destroyed, and recreated. +- Default-backend GPU initialization does not run for Newton. +- Newton finalization does not call default-backend GPU fetch/apply APIs. +- Destroying a Newton simulation does not break subsequent default-backend + simulation creation. + +Rigid object: + +- Dynamic rigid bodies fall under Newton. +- Static and kinematic rigid bodies can be created under Newton. +- Pose, velocity, acceleration, force/torque, reset, COM pose, mass, friction, + inertia, collision filters, and geometry APIs behave consistently with the + documented support matrix. +- Unsupported APIs produce the documented warning or exception. + +Gym: + +- Rigid-only Newton env initializes, steps, resets, and reads observations. +- Robot/articulation env under Newton raises the expected unsupported error. + +Gradient: + +- `requires_grad=True` plus `solver_type="semi_implicit"` can create a gradient + rollout. +- A simple loss can backpropagate through a rollout without CPU/NumPy observation + paths. + +## Known Risks + +- DexSim Newton monkey-patches global classes. Global teardown can affect other + worlds if used at the wrong time. +- Public body/articulation ID mapping APIs may still need DexSim improvements. +- Newton gravity and contact configuration may not yet match every default-backend + setting. +- Some object constructors still contain default-backend assumptions such as + warmup updates; keep Newton guarded from those paths. +- Runtime shape/property mutations may require model rebuilds rather than live + updates. diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst index dacdae33..394968d9 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst @@ -9,9 +9,11 @@ .. autosummary:: ArticulationCfg + DefaultPhysicsCfg GPUMemoryCfg JointDrivePropertiesCfg LightCfg + NewtonPhysicsCfg ObjectBaseCfg PhysicsCfg RigidBodyAttributesCfg diff --git a/docs/source/features/interaction/preview_asset.md b/docs/source/features/interaction/preview_asset.md index df3aa040..fa1541cc 100644 --- a/docs/source/features/interaction/preview_asset.md +++ b/docs/source/features/interaction/preview_asset.md @@ -73,7 +73,7 @@ asset.set_root_pose(pos=[0, 0, 1.0], rot=[0, 0, 0]) | `--body_type` | Body type for rigid objects: `dynamic`, `kinematic`, `static` | `kinematic` | | `--use_usd_properties` | Use physical properties from the USD file instead of defaults | `False` | | `--fix_base` | Fix the base of articulations | `True` | -| `--sim_device` | Simulation device | `cpu` | +| `--device` | Simulation device | `cpu` | | `--headless` | Run without rendering window | `False` | | `--renderer` | Renderer backend: `hybrid`, `fast-rt` or `rt` | `hybrid` | | `--preview` | Enter interactive embed mode after loading | `False` | diff --git a/docs/source/features/workspace_analyzer/workspace_analyzer.md b/docs/source/features/workspace_analyzer/workspace_analyzer.md index 133ee0eb..ee096fbc 100644 --- a/docs/source/features/workspace_analyzer/workspace_analyzer.md +++ b/docs/source/features/workspace_analyzer/workspace_analyzer.md @@ -24,7 +24,7 @@ from embodichain.lab.sim.utility.workspace_analyzer import ( ) # Setup simulation -sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) +sim = SimulationManager(SimulationManagerCfg(headless=False, device="cpu")) # Add robot robot = sim.add_robot(DexforceW1Cfg.from_dict({ @@ -167,7 +167,7 @@ from embodichain.lab.sim.utility.workspace_analyzer import ( from embodichain.lab.sim.utility.workspace_analyzer.configs import VisualizationConfig # Setup simulation -sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) +sim = SimulationManager(SimulationManagerCfg(headless=False, device="cpu")) # Add robot robot = sim.add_robot(DexforceW1Cfg.from_dict({ diff --git a/docs/source/overview/sim/planners/motion_generator.md b/docs/source/overview/sim/planners/motion_generator.md index bee23cb6..f166f12e 100644 --- a/docs/source/overview/sim/planners/motion_generator.md +++ b/docs/source/overview/sim/planners/motion_generator.md @@ -35,7 +35,7 @@ sim_cfg = SimulationManagerCfg( width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device="cpu", + device="cpu", ) sim = SimulationManager(sim_cfg) diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index e8605047..0950a6fb 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -77,7 +77,7 @@ from embodichain.lab.sim.objects import Articulation, ArticulationCfg # 1. Initialize Simulation device = "cuda" if torch.cuda.is_available() else "cpu" -sim_cfg = SimulationManagerCfg(sim_device=device) +sim_cfg = SimulationManagerCfg(device=device) sim = SimulationManager(sim_config=sim_cfg) # 2. Configure Articulation diff --git a/docs/source/overview/sim/sim_cloth.md b/docs/source/overview/sim/sim_cloth.md index 78cc4bf5..36e33cab 100644 --- a/docs/source/overview/sim/sim_cloth.md +++ b/docs/source/overview/sim/sim_cloth.md @@ -94,7 +94,7 @@ def create_2d_grid_mesh(width: float, height: float, nx: int = 1, ny: int = 1): # 1. Initialize Simulation device = "cuda" if torch.cuda.is_available() else "cpu" -sim_cfg = SimulationManagerCfg(sim_device=device) +sim_cfg = SimulationManagerCfg(device=device) sim = SimulationManager(sim_config=sim_cfg) cloth_verts, cloth_faces = create_2d_grid_mesh(width=0.3, height=0.3, nx=12, ny=12) diff --git a/docs/source/overview/sim/sim_manager.md b/docs/source/overview/sim/sim_manager.md index 675f8a06..8ab2709a 100644 --- a/docs/source/overview/sim/sim_manager.md +++ b/docs/source/overview/sim/sim_manager.md @@ -15,13 +15,16 @@ The simulation is configured using the {class}`SimulationManagerCfg` class. ```python from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import DefaultPhysicsCfg sim_config = SimulationManagerCfg( width=1920, # Window width height=1080, # Window height num_envs=10, # Number of parallel environments - physics_dt=0.01, # Physics time step - sim_device="cpu", # Simulation device ("cpu" or "cuda:0", etc.) + device="cpu", # Simulation device ("cpu" or "cuda:0", etc.) + physics_cfg=DefaultPhysicsCfg( + physics_dt=0.01, # Physics time step + ), arena_space=5.0 # Spacing between environments ) ``` @@ -39,14 +42,20 @@ sim_config = SimulationManagerCfg( | `cpu_num` | `int` | `1` | The number of CPU threads to use for the simulation engine. | | `num_envs` | `int` | `1` | The number of parallel environments (arenas) to simulate. | | `arena_space` | `float` | `5.0` | The distance between each arena when building multiple arenas. | -| `physics_dt` | `float` | `0.01` | The time step for the physics simulation. | -| `sim_device` | `str` \| `torch.device` | `"cpu"` | The device for the physics simulation. | -| `physics_config` | `PhysicsCfg` | `PhysicsCfg()` | The physics configuration parameters. | -| `gpu_memory_config` | `GPUMemoryCfg` | `GPUMemoryCfg()` | The GPU memory configuration parameters. | +| `physics_cfg` | `DefaultPhysicsCfg` \| `NewtonPhysicsCfg` | `DefaultPhysicsCfg()` | Physics backend configuration (class selects default vs Newton). | ### Physics Configuration -The {class}`~cfg.PhysicsCfg` class controls the global physics simulation parameters. +Use {class}`~cfg.DefaultPhysicsCfg` for the default PhysX backend or {class}`~cfg.NewtonPhysicsCfg` for Newton. GPU memory settings are on {class}`~cfg.DefaultPhysicsCfg` as ``gpu_memory``. + +All physics backends inherit these base parameters from {class}`~cfg.PhysicsCfg`: + +| Parameter | Type | Default | Description | +| :--- | :--- | :--- | :--- | +| `physics_dt` | `float` | `0.01` | The time step for the physics simulation. | +| `device` | `str` \| `torch.device` | `"cpu"` | The device for the physics simulation. | + +The {class}`~cfg.DefaultPhysicsCfg` class controls the global default-backend physics simulation parameters. | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | @@ -190,7 +199,7 @@ while True: In this mode, the physics simulation stepping is automatically handling by the physics thread running in dexsim engine, which makes it easier to use for visualization and interactive applications. -> When in automatic update mode, user are recommanded to use CPU `sim_device` for simulation. +> When in automatic update mode, user are recommanded to use CPU `device` for simulation. ## Mainly used methods @@ -210,9 +219,37 @@ In this mode, the physics simulation stepping is automatically handling by the p > Currently, multiple instances are not supported for ray tracing rendering backend. Good news is that we are working on adding this feature in future releases. +## Newton Physics Backend + +EmbodiChain supports the DexSim Newton physics backend as an alternative to the default PhysX backend. Select the Newton backend by passing a `NewtonPhysicsCfg` to `physics_cfg`: + +```python +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg + +sim_config = SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg(), +) +``` + +### Supported Runtime Operations + +The Newton backend supports runtime mutation of the following physical properties on rigid objects: + +| Property | `get_*` | `set_*` | Notes | +| :--- | :---: | :---: | :--- | +| Mass | ✅ | ✅ | Per-body mass via batch GPU API | +| Friction | ✅ | ✅ | Dynamic friction coefficient | +| Inertia | ✅ | ✅ | Diagonal inertia tensor (3-vector) | +| Restitution | ✅ | ✅ | Bounce coefficient | +| Damping | ✅ | ❌ | Read from initial metadata only | +| Body Type | ✅ | ❌ | Cannot change dynamic ↔ kinematic at runtime | +| Bulk `set_attrs` | — | ❌ | Use individual setters instead | + + For more methods and details, refer to the [SimulationManager](https://dexforce.github.io/EmbodiChain/api_reference/embodichain/embodichain.lab.sim.html#embodichain.lab.sim.SimulationManager) documentation. ### Related Tutorials - [Basic scene creation](https://dexforce.github.io/EmbodiChain/tutorial/create_scene.html) -- [Interactive simulation with Gizmo](https://dexforce.github.io/EmbodiChain/tutorial/gizmo.html) \ No newline at end of file +- [Interactive simulation with Gizmo](https://dexforce.github.io/EmbodiChain/tutorial/gizmo.html) diff --git a/docs/source/overview/sim/sim_rigid_object.md b/docs/source/overview/sim/sim_rigid_object.md index 185a533d..4d83a435 100644 --- a/docs/source/overview/sim/sim_rigid_object.md +++ b/docs/source/overview/sim/sim_rigid_object.md @@ -49,7 +49,7 @@ from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # 1. Initialize Simulation device = "cuda" if torch.cuda.is_available() else "cpu" -sim_cfg = SimulationManagerCfg(sim_device=device) +sim_cfg = SimulationManagerCfg(device=device) sim = SimulationManager(sim_cfg) # 2. Configure a rigid object (cube) @@ -192,7 +192,7 @@ N denotes the number of parallel environments when using vectorized simulation ( - Use `static` body type for fixed obstacles or environment pieces (they do not consume dynamic simulation resources). - Use `kinematic` for objects whose pose is driven by code (teleporting or animation) but still interact with dynamic objects. - For complex meshes, enabling convex decomposition (`RigidObjectCfg.max_convex_hull_num`) or providing a simplified collision mesh improves stability and performance. -- To use GPU physics, ensure `SimulationManagerCfg.sim_device` is set to `cuda` and call `sim.init_gpu_physics()` before large-batch simulations. +- To use GPU physics, ensure `SimulationManagerCfg.device` is set to `cuda` and call `sim.init_gpu_physics()` before large-batch simulations. ## Example: Applying Force and Torque diff --git a/docs/source/overview/sim/sim_rigid_object_group.md b/docs/source/overview/sim/sim_rigid_object_group.md index d6d22838..d5c7fb50 100644 --- a/docs/source/overview/sim/sim_rigid_object_group.md +++ b/docs/source/overview/sim/sim_rigid_object_group.md @@ -44,7 +44,7 @@ from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # 1. Initialize Simulation device = "cuda" if torch.cuda.is_available() else "cpu" -sim_cfg = SimulationManagerCfg(sim_device=device) +sim_cfg = SimulationManagerCfg(device=device) sim = SimulationManager(sim_cfg) # 2. Define shared physics attributes @@ -109,7 +109,7 @@ Use these shapes when collecting vectorized observations for multi-environment t - Prefer providing simplified collision meshes or enabling convex decomposition (`max_convex_hull_num` > 1) for complex visual meshes to improve physics stability. - `RigidObjectGroup` only supports `dynamic` and `kinematic` body types (not `static`). - When teleporting many members, batch pose updates and call `sim.update()` once to avoid synchronization overhead. -- For GPU physics, set `SimulationManagerCfg.sim_device` to `cuda` and call `sim.init_gpu_physics()` before running simulations. +- For GPU physics, set `SimulationManagerCfg.device` to `cuda` and call `sim.init_gpu_physics()` before running simulations. - Use `clear_dynamics()` to reset velocities without changing poses. ## Example: Working with Group Poses diff --git a/docs/source/overview/sim/sim_robot.md b/docs/source/overview/sim/sim_robot.md index e0ab5992..e184c28c 100644 --- a/docs/source/overview/sim/sim_robot.md +++ b/docs/source/overview/sim/sim_robot.md @@ -25,9 +25,9 @@ from embodichain.lab.sim.objects import Robot, RobotCfg from embodichain.lab.sim.solvers import SolverCfg # 1. Initialize Simulation Environment -# Note: Use 'sim_device' to specify device (e.g., "cuda:0" or "cpu") +# Note: Use 'device' to specify device (e.g., "cuda:0" or "cpu") device = "cuda" if torch.cuda.is_available() else "cpu" -sim_cfg = SimulationManagerCfg(sim_device=device, physics_dt=0.01) +sim_cfg = SimulationManagerCfg(device=device, physics_dt=0.01) sim = SimulationManager(sim_config=sim_cfg) # 2. Configure Robot diff --git a/docs/source/overview/sim/sim_soft_object.md b/docs/source/overview/sim/sim_soft_object.md index d8b9f510..7d2d7586 100644 --- a/docs/source/overview/sim/sim_soft_object.md +++ b/docs/source/overview/sim/sim_soft_object.md @@ -55,7 +55,7 @@ from embodichain.lab.sim.objects import SoftObject, SoftObjectCfg # 1. Initialize Simulation device = "cuda" if torch.cuda.is_available() else "cpu" -sim_cfg = SimulationManagerCfg(sim_device=device) +sim_cfg = SimulationManagerCfg(device=device) sim = SimulationManager(sim_config=sim_cfg) # 2. Configure Soft Object diff --git a/docs/source/resources/robot/cobotmagic.md b/docs/source/resources/robot/cobotmagic.md index de23dd2a..b60d7802 100644 --- a/docs/source/resources/robot/cobotmagic.md +++ b/docs/source/resources/robot/cobotmagic.md @@ -39,7 +39,7 @@ CobotMagic is a versatile dual-arm collaborative robot developed by AgileX Robot from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.robots import CobotMagicCfg -config = SimulationManagerCfg(headless=False, sim_device="cpu", num_envs=2) +config = SimulationManagerCfg(headless=False, device="cpu", num_envs=2) sim = SimulationManager(config) sim.set_manual_update(False) diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 1a0fa89e..2d4a0de6 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -129,15 +129,16 @@ def __init__( self._setup_scene(**kwargs) - # TODO: To be removed. - if self.device.type == "cuda": + if self.sim.is_default_backend and self.sim.is_use_gpu_physics: self.sim.init_gpu_physics() + elif self.sim.is_newton_backend: + self.sim.finalize_newton_physics() if not self.sim_cfg.headless: self.sim.open_window() self._elapsed_steps = torch.zeros( - self._num_envs, dtype=torch.int32, device=self.sim_cfg.sim_device + self._num_envs, dtype=torch.int32, device=self.sim_cfg.device ) # -1 means no limit on episode length, and the episode will only end when the task is successfully completed or failed. @@ -250,7 +251,7 @@ def _setup_scene(self, **kwargs): self.sim_cfg.headless = headless logger.log_info( - f"Initializing {self.num_envs} environments on {self.sim_cfg.sim_device}." + f"Initializing {self.num_envs} environments on {self.sim_cfg.device}." ) self.robot = self._setup_robot(**kwargs) diff --git a/embodichain/lab/gym/envs/managers/observations.py b/embodichain/lab/gym/envs/managers/observations.py index 3729ec8d..d3fcb984 100644 --- a/embodichain/lab/gym/envs/managers/observations.py +++ b/embodichain/lab/gym/envs/managers/observations.py @@ -1157,14 +1157,9 @@ def __call__( device=env.device, ) else: - ( - stiffness, - damping, - max_effort, - max_velocity, - friction, - armature, - ) = art.get_joint_drive() + stiffness, damping, max_effort, max_velocity, friction, armature = ( + art.get_joint_drive() + ) result = TensorDict( { "stiffness": stiffness, diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index bbef9ba1..ba133359 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -738,6 +738,7 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: --device: Device to run the environment on (default: 'cpu') --headless: Whether to perform the simulation in headless mode (default: False) --renderer: Renderer backend to use for the simulation. Options are 'hybrid', 'fast-rt', and 'rt'. (default: 'hybrid') + --physics: Physics backend configuration to use. Options are 'default' and 'newton'. (default: 'default') --gpu_id: The GPU ID to use for the simulation (default: 0) --gym_config: Path to gym config file (default: '') --action_config: Path to action config file (default: None) @@ -776,6 +777,13 @@ def add_env_launcher_args_to_parser(parser: argparse.ArgumentParser) -> None: default="auto", help="Renderer backend to use for the simulation.", ) + parser.add_argument( + "--physics", + type=str, + choices=["default", "newton"], + default="default", + help="Physics backend configuration to use for the simulation.", + ) parser.add_argument( "--arena_space", help="The size of the arena space.", @@ -838,6 +846,7 @@ def merge_args_with_gym_config(args: argparse.Namespace, gym_config: dict) -> di merged_config["device"] = args.device merged_config["headless"] = args.headless merged_config["renderer"] = args.renderer + merged_config["physics"] = args.physics merged_config["gpu_id"] = args.gpu_id merged_config["arena_space"] = args.arena_space return merged_config @@ -858,7 +867,7 @@ def build_env_cfg_from_args( from embodichain.utils.utility import load_config from embodichain.lab.gym.envs import EmbodiedEnvCfg from embodichain.lab.sim import SimulationManagerCfg - from embodichain.lab.sim.cfg import RenderCfg + from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend gym_config = load_config(args.gym_config) gym_config = merge_args_with_gym_config(args, gym_config) @@ -880,8 +889,9 @@ def build_env_cfg_from_args( cfg.sim_cfg = SimulationManagerCfg( headless=gym_config["headless"], - sim_device=gym_config["device"], + device=gym_config["device"], render_cfg=RenderCfg(renderer=gym_config["renderer"]), + physics_cfg=physics_cfg_for_backend(gym_config["physics"]), gpu_id=gym_config["gpu_id"], arena_space=gym_config["arena_space"], ) diff --git a/embodichain/lab/scripts/preview_asset.py b/embodichain/lab/scripts/preview_asset.py index 49c86de5..3dea22fe 100644 --- a/embodichain/lab/scripts/preview_asset.py +++ b/embodichain/lab/scripts/preview_asset.py @@ -53,6 +53,7 @@ from typing import TYPE_CHECKING +from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.utils.logger import log_info, log_warning, log_error if TYPE_CHECKING: @@ -68,13 +69,17 @@ def build_sim_cfg(args: argparse.Namespace): Returns: SimulationManagerCfg: Simulation configuration. """ - from embodichain.lab.sim.cfg import RenderCfg + from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend from embodichain.lab.sim.sim_manager import SimulationManagerCfg return SimulationManagerCfg( headless=args.headless, - sim_device=args.sim_device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), + gpu_id=args.gpu_id, + num_envs=args.num_envs, + arena_space=args.arena_space, ) @@ -248,6 +253,7 @@ def cli(): parser = argparse.ArgumentParser( description="Preview a USD or mesh asset in the EmbodiChain simulation." ) + add_env_launcher_args_to_parser(parser) parser.add_argument( "--asset_path", @@ -313,25 +319,6 @@ def cli(): default=True, help="Fix the base of articulations (default: True).", ) - parser.add_argument( - "--sim_device", - type=str, - default="cpu", - help="Simulation device (default: cpu).", - ) - parser.add_argument( - "--headless", - action="store_true", - default=False, - help="Run without rendering window.", - ) - parser.add_argument( - "--renderer", - type=str, - choices=["hybrid", "fast-rt", "rt"], - default="hybrid", - help="Renderer backend (default: hybrid).", - ) parser.add_argument( "--env_map", type=str, @@ -341,13 +328,6 @@ def cli(): "name (e.g. 'Studio') or an absolute file path (.hdr/.png/.exr)." ), ) - parser.add_argument( - "--preview", - action="store_true", - default=False, - help="Enter interactive embed mode after loading.", - ) - args = parser.parse_args() main(args) diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py index 0b2fe922..a5ad87cf 100644 --- a/embodichain/lab/sim/cfg.py +++ b/embodichain/lab/sim/cfg.py @@ -15,11 +15,12 @@ # ---------------------------------------------------------------------------- from __future__ import annotations +from collections.abc import Mapping import os import numpy as np import torch -from typing import Sequence, Union, Dict, Literal, List, Any, Optional +from typing import Sequence, Union, Dict, Literal, List, Any, Optional, TYPE_CHECKING from dataclasses import field, MISSING from dexsim.types import ( @@ -41,6 +42,9 @@ from .shapes import ShapeCfg, MeshCfg +if TYPE_CHECKING: + from dexsim.engine.newton_physics.solvers_cfg import NewtonSolverCfg + # Global default renderer settings for simulation. # # The sentinel value ``"auto"`` defers the choice to GPU-based auto-selection @@ -66,8 +70,11 @@ class RenderCfg: - 'rt' is an offline ray-traced renderer for maximum visual fidelity, suitable for high-quality rendering tasks. """ - spp: int = 1 - """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid', 'fast-rt' or 'rt'.""" + enable_denoiser: bool = True + """Whether to enable denoising. Only valid when renderer is 'hybrid' or 'fast-rt'.""" + + spp: int = 64 + """Samples per pixel for ray tracing rendering. This parameter is only valid when renderer is 'hybrid' or 'fast-rt' and enable_denoiser is False.""" def to_dexsim_flags(self): if self.renderer == "hybrid": @@ -90,8 +97,45 @@ def to_dexsim_flags(self): ) +@configclass +class GPUMemoryCfg: + """GPU memory configuration for default-backend GPU physics simulation.""" + + temp_buffer_capacity: int = 2**24 + """Increase this if you get 'PxgPinnedHostLinearMemoryAllocator: overflowing initial allocation size, increase capacity to at least %.' """ + + max_rigid_contact_count: int = 2**19 + """Increase this if you get 'Contact buffer overflow detected'""" + + max_rigid_patch_count: int = ( + 2**18 + ) # 81920 is DexSim default but most tasks work with 2**18 + """Increase this if you get 'Patch buffer overflow detected'""" + + heap_capacity: int = 2**26 + + found_lost_pairs_capacity: int = ( + 2**25 + ) # 262144 is DexSim default but most tasks work with 2**25 + found_lost_aggregate_pairs_capacity: int = 2**10 + total_aggregate_pairs_capacity: int = 2**10 + + @configclass class PhysicsCfg: + """Base configuration for DexSim physics backends.""" + + physics_dt: float = 1.0 / 100.0 + """The time step for the physics simulation.""" + + device: str | torch.device = "cpu" + """The device for the physics simulation. Can be 'cpu', 'cuda', or a torch.device object.""" + + +@configclass +class DefaultPhysicsCfg(PhysicsCfg): + """Configuration for the DexSim default (PhysX) physics backend.""" + gravity: np.ndarray = field(default_factory=lambda: np.array([0, 0, -9.81])) """Gravity vector for the simulation environment.""" @@ -115,15 +159,18 @@ class PhysicsCfg: length_tolerance: float = 0.05 """The length tolerance for the simulation. - - Note: the larger the tolerance, the faster the simulation will be. + + Note: the larger the tolerance, the faster the simulation will be. """ speed_tolerance: float = 0.25 """The speed tolerance for the simulation. - + Note: the larger the tolerance, the faster the simulation will be. """ + gpu_memory: GPUMemoryCfg = field(default_factory=GPUMemoryCfg) + """GPU memory configuration for GPU physics simulation.""" + def to_dexsim_args(self) -> Dict[str, Any]: """Convert to dexsim physics args dictionary.""" args = { @@ -138,6 +185,160 @@ def to_dexsim_args(self) -> Dict[str, Any]: return args +@configclass +class NewtonPhysicsCfg(PhysicsCfg): + """Configuration for DexSim Newton physics backend.""" + + device: str | torch.device = "cuda:0" + """The device for Newton physics simulation (e.g. ``cuda:0``).""" + + num_substeps: int = 10 + """Number of Newton solver substeps per EmbodiChain physics step.""" + + requires_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_cfg: Mapping[str, Any] | NewtonSolverCfg | None = None + """Optional Newton solver configuration. + + A mapping is converted to the matching DexSim Newton solver config. Include + ``solver_type`` or ``class_type`` to select the solver, then add any + parameters accepted by that DexSim solver config. If omitted, the Newton + backend uses DexSim's MuJoCo Warp solver config by default. + """ + + 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, + 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(self.device) if isinstance(self.device, str) else self.device + ) + device = ( + f"cuda:{gpu_id}" + if torch_device.type == "cuda" and torch_device.index is None + else str(torch_device) + ) + + solver_cfg_map = { + "mujoco_warp": MJWarpSolverCfg, + "xpbd": XPBDSolverCfg, + "semi_implicit": SemiImplicitSolverCfg, + "featherstone": FeatherstoneSolverCfg, + "vbd": VBDSolverCfg, + } + solver_cfg = _newton_solver_cfg_to_dexsim( + solver_cfg=self.solver_cfg, + solver_cfg_map=solver_cfg_map, + ) + + if self.requires_grad and solver_cfg.solver_type != "semi_implicit": + logger.log_error( + "Newton gradient mode requires solver_type='semi_implicit'." + ) + + cfg = NewtonCfg( + dt=self.physics_dt, + num_substeps=self.num_substeps, + device=device, + debug_mode=self.debug_mode, + requires_grad=self.requires_grad, + solver_cfg=solver_cfg, + collision_pipeline_cfg=NewtonCollisionPipelineCfg( + broad_phase=self.broad_phase, + requires_grad=self.requires_grad, + ), + ) + cfg.use_cuda_graph = self.use_cuda_graph and not self.requires_grad + cfg._visualizer_enabled = self.visualizer_enabled + return cfg + + +def _normalize_newton_solver_type(solver_type: str) -> str: + """Normalize public EmbodiChain and DexSim Newton solver aliases.""" + key = solver_type.replace("-", "_").lower() + aliases = { + "mjwarp": "mujoco_warp", + "mjwarpsolver": "mujoco_warp", + "mjwarpsolvercfg": "mujoco_warp", + "mjwarp_solver": "mujoco_warp", + "mjwarp_solver_cfg": "mujoco_warp", + "mujoco_warp": "mujoco_warp", + "mujocowarp": "mujoco_warp", + "mujocowarpsolver": "mujoco_warp", + "mujocowarpsolvercfg": "mujoco_warp", + "xpbdsolver": "xpbd", + "xpbdsolvercfg": "xpbd", + "xpbd": "xpbd", + "semiimplicit": "semi_implicit", + "semi_implicit": "semi_implicit", + "semiimplicitsolver": "semi_implicit", + "semiimplicitsolvercfg": "semi_implicit", + "featherstone": "featherstone", + "featherstonesolver": "featherstone", + "featherstonesolvercfg": "featherstone", + "vbd": "vbd", + "vbdsolver": "vbd", + "vbdsolvercfg": "vbd", + } + if key not in aliases: + logger.log_error( + f"Unsupported Newton solver type '{solver_type}'. " + "Expected one of 'mjwarp', 'xpbd', 'semi_implicit', " + "'featherstone', or 'vbd'." + ) + return aliases[key] + + +def _newton_solver_cfg_to_dexsim( + solver_cfg: Mapping[str, Any] | object | None, + solver_cfg_map: Mapping[str, type], +) -> object: + """Convert EmbodiChain Newton solver config input to a DexSim config.""" + if solver_cfg is None: + return solver_cfg_map["mujoco_warp"]() + + if not isinstance(solver_cfg, Mapping): + if not hasattr(solver_cfg, "solver_type"): + logger.log_error( + "Newton solver_cfg must be a mapping or a DexSim Newton solver " + "config object with a 'solver_type' attribute." + ) + return solver_cfg + + solver_cfg_data = dict(solver_cfg) + configured_solver_type = ( + solver_cfg_data.pop("solver_type", None) + or solver_cfg_data.pop("class_type", None) + or "mujoco_warp" + ) + normalized_solver_type = _normalize_newton_solver_type(str(configured_solver_type)) + return solver_cfg_map[normalized_solver_type](**solver_cfg_data) + + @configclass class MarkerCfg: """Configuration for visual markers in the simulation. @@ -195,28 +396,32 @@ class WindowRecordCfg: """Video file prefix used when no explicit save path is provided.""" -@configclass -class GPUMemoryCfg: - """A gpu memory configuration dataclass that neatly holds all parameters that configure physics GPU memory for simulation""" - - temp_buffer_capacity: int = 2**24 - """Increase this if you get 'PxgPinnedHostLinearMemoryAllocator: overflowing initial allocation size, increase capacity to at least %.' """ - - max_rigid_contact_count: int = 2**19 - """Increase this if you get 'Contact buffer overflow detected'""" - - max_rigid_patch_count: int = ( - 2**18 - ) # 81920 is DexSim default but most tasks work with 2**18 - """Increase this if you get 'Patch buffer overflow detected'""" +def physics_cfg_for_backend( + backend: Literal["default", "newton"], +) -> DefaultPhysicsCfg | NewtonPhysicsCfg: + """Return a default physics configuration instance for the given backend.""" + if backend == "newton": + return NewtonPhysicsCfg() + return DefaultPhysicsCfg() + + +def physics_backend_from_cfg( + physics_cfg: PhysicsCfg, +) -> Literal["default", "newton"]: + """Infer the physics backend name from a physics configuration instance.""" + if isinstance(physics_cfg, NewtonPhysicsCfg): + return "newton" + if isinstance(physics_cfg, DefaultPhysicsCfg): + return "default" + logger.log_error( + f"Unsupported physics_cfg type '{type(physics_cfg).__name__}'. " + "Expected DefaultPhysicsCfg or NewtonPhysicsCfg." + ) - heap_capacity: int = 2**26 - found_lost_pairs_capacity: int = ( - 2**25 - ) # 262144 is DexSim default but most tasks work with 2**25 - found_lost_aggregate_pairs_capacity: int = 2**10 - total_aggregate_pairs_capacity: int = 2**10 +def validate_physics_cfg(physics_cfg: PhysicsCfg) -> None: + """Validate that ``physics_cfg`` is a supported backend configuration.""" + physics_backend_from_cfg(physics_cfg) @configclass diff --git a/embodichain/lab/sim/common.py b/embodichain/lab/sim/common.py index f1380ed6..ff36ba5e 100644 --- a/embodichain/lab/sim/common.py +++ b/embodichain/lab/sim/common.py @@ -54,6 +54,7 @@ def __init__( cfg: ObjectBaseCfg, entities: List[T] = None, device: torch.device = torch.device("cpu"), + auto_reset: bool = True, ) -> None: if entities is None or len(entities) == 0: @@ -66,7 +67,8 @@ def __init__( self._entities = entities self.device = device - self.reset() + if auto_reset: + self.reset() def __str__(self) -> str: return f"{self.__class__}: managing {self.num_instances} {self._entities[0].__class__} objects | uid: {self.uid} | device: {self.device}" diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index ed5a9e32..186efde5 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import dexsim import numpy as np @@ -23,11 +25,7 @@ from typing import List, Sequence, Dict, Union, Tuple, Optional from dexsim.engine import Articulation as _Articulation -from dexsim.types import ( - ArticulationFlag, - ArticulationGPUAPIWriteType, - ArticulationGPUAPIReadType, -) +from dexsim.types import ArticulationFlag from dexsim.engine import CudaArray, PhysicsScene from embodichain.lab.sim import VisualMaterialInst, VisualMaterial @@ -40,10 +38,14 @@ from dexsim.types import PhysicalAttr from embodichain.utils.string import resolve_matching_names from embodichain.lab.sim.common import BatchEntity +from embodichain.lab.sim.objects.backends import ( + DefaultArticulationView, + NewtonArticulationView, + is_newton_scene, +) from embodichain.utils.math import ( matrix_from_quat, quat_from_matrix, - convert_quat, matrix_from_euler, ) from embodichain.lab.sim.utility.sim_utils import ( @@ -75,18 +77,17 @@ def __init__( self.ps = ps self.num_instances = len(entities) self.device = device - - # get gpu indices for the entities. - # only meaningful when using GPU physics. - self.gpu_indices = ( - torch.as_tensor( - [entity.get_gpu_index() for entity in self.entities], - dtype=torch.int32, - device=self.device, + if is_newton_scene(ps): + self.articulation_view = NewtonArticulationView( + entities=entities, scene=ps, device=device + ) + else: + self.articulation_view = DefaultArticulationView( + entities=entities, ps=ps, device=device ) - if self.device.type == "cuda" - else None - ) + + # Backward-compatible alias for callers that use GPU/articulation ids. + self.gpu_indices = self.articulation_view.articulation_ids_tensor self.dof = self.entities[0].get_dof() self.num_links = self.entities[0].get_links_num() @@ -104,7 +105,7 @@ def __init__( max_num_links = ( self.ps.gpu_get_articulation_max_link_count() - if self.device.type == "cuda" + if self.device.type == "cuda" and not self.is_newton_backend else self.num_links ) self._body_link_pose = torch.zeros( @@ -131,7 +132,7 @@ def __init__( max_dof = ( self.ps.gpu_get_articulation_max_dof() - if self.device.type == "cuda" + if self.device.type == "cuda" and not self.is_newton_backend else self.dof ) self._target_qpos = torch.zeros( @@ -153,6 +154,14 @@ def __init__( (self.num_instances, max_dof), dtype=torch.float32, device=self.device ) + @property + def is_newton_backend(self) -> bool: + return self.articulation_view.is_newton_backend + + @property + def is_ready(self) -> bool: + return self.articulation_view.is_ready + @property def root_pose(self) -> torch.Tensor: """Get the root pose of the articulation. @@ -160,24 +169,7 @@ def root_pose(self) -> torch.Tensor: Returns: torch.Tensor: The root pose of the articulation with shape of (num_instances, 7). """ - if self.device.type == "cpu": - # Fetch pose from CPU entities - root_pose = torch.as_tensor( - np.array([entity.get_local_pose() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - xyzs = root_pose[:, :3, 3] - quats = quat_from_matrix(root_pose[:, :3, :3]) - return torch.cat((xyzs, quats), dim=-1) - else: - self.ps.gpu_fetch_root_data( - data=self._root_pose, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.ROOT_GLOBAL_POSE, - ) - self._root_pose[:, :4] = convert_quat(self._root_pose[:, :4], to="wxyz") - return self._root_pose[:, [4, 5, 6, 0, 1, 2, 3]] + return self.articulation_view.fetch_root_pose(self._root_pose) @property def root_lin_vel(self) -> torch.Tensor: @@ -186,22 +178,7 @@ def root_lin_vel(self) -> torch.Tensor: Returns: torch.Tensor: The linear velocity of the root link with shape of (num_instances, 3). """ - if self.device.type == "cpu": - # Fetch linear velocity from CPU entities - return torch.as_tensor( - np.array( - [entity.get_root_link_velocity()[:3] for entity in self.entities] - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_root_data( - data=self._root_lin_vel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.ROOT_LINEAR_VELOCITY, - ) - return self._root_lin_vel.clone() + return self.articulation_view.fetch_root_linear_velocity(self._root_lin_vel) @property def root_ang_vel(self) -> torch.Tensor: @@ -210,22 +187,7 @@ def root_ang_vel(self) -> torch.Tensor: Returns: torch.Tensor: The angular velocity of the root link with shape of (num_instances, 3). """ - if self.device.type == "cpu": - # Fetch angular velocity from CPU entities - return torch.as_tensor( - np.array( - [entity.get_root_link_velocity()[3:] for entity in self.entities] - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_root_data( - data=self._root_ang_vel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.ROOT_ANGULAR_VELOCITY, - ) - return self._root_ang_vel.clone() + return self.articulation_view.fetch_root_angular_velocity(self._root_ang_vel) @property def root_vel(self) -> torch.Tensor: @@ -243,22 +205,7 @@ def qpos(self) -> torch.Tensor: Returns: torch.Tensor: The current positions of the articulation with shape of (num_instances, dof). """ - if self.device.type == "cpu": - # Fetch qpos from CPU entities - return torch.as_tensor( - np.array( - [entity.get_current_qpos() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_joint_data( - data=self._qpos, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.JOINT_POSITION, - ) - return self._qpos[:, : self.dof].clone() + return self.articulation_view.fetch_qpos(self._qpos) @property def target_qpos(self) -> torch.Tensor: @@ -267,22 +214,7 @@ def target_qpos(self) -> torch.Tensor: Returns: torch.Tensor: The target positions of the articulation with shape of (num_instances, dof). """ - if self.device.type == "cpu": - # Fetch target_qpos from CPU entities - return torch.as_tensor( - np.array( - [entity.get_target_qpos() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_joint_data( - data=self._target_qpos, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.JOINT_TARGET_POSITION, - ) - return self._target_qpos[:, : self.dof].clone() + return self.articulation_view.fetch_target_qpos(self._target_qpos) @property def qvel(self) -> torch.Tensor: @@ -291,20 +223,7 @@ def qvel(self) -> torch.Tensor: Returns: torch.Tensor: The current velocities of the articulation with shape of (num_instances, dof). """ - if self.device.type == "cpu": - # Fetch qvel from CPU entities - return torch.as_tensor( - np.array([entity.get_current_qvel() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_joint_data( - data=self._qvel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.JOINT_VELOCITY, - ) - return self._qvel[:, : self.dof].clone() + return self.articulation_view.fetch_qvel(self._qvel) @property def target_qvel(self) -> torch.Tensor: @@ -312,22 +231,7 @@ def target_qvel(self) -> torch.Tensor: Returns: torch.Tensor: The target velocities of the articulation with shape of (num_instances, dof). """ - if self.device.type == "cpu": - # Fetch target_qvel from CPU entities - return torch.as_tensor( - np.array( - [entity.get_target_qvel() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_joint_data( - data=self._target_qvel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.JOINT_TARGET_VELOCITY, - ) - return self._target_qvel[:, : self.dof].clone() + return self.articulation_view.fetch_target_qvel(self._target_qvel) @property def qacc(self) -> torch.Tensor: @@ -336,20 +240,7 @@ def qacc(self) -> torch.Tensor: Returns: torch.Tensor: The current accelerations of the articulation with shape of (num_instances, dof). """ - if self.device.type == "cpu": - # Fetch qacc from CPU entities - return torch.as_tensor( - np.array([entity.get_current_qacc() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_joint_data( - data=self._qacc, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.JOINT_ACCELERATION, - ) - return self._qacc[:, : self.dof].clone() + return self.articulation_view.fetch_qacc(self._qacc) @property def qf(self) -> torch.Tensor: @@ -358,20 +249,7 @@ def qf(self) -> torch.Tensor: Returns: torch.Tensor: The current forces of the articulation with shape of (num_instances, dof). """ - if self.device.type == "cpu": - # Fetch qf from CPU entities - return torch.as_tensor( - np.array([entity.get_current_qf() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_joint_data( - data=self._qf, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.JOINT_FORCE, - ) - return self._qf[:, : self.dof].clone() + return self.articulation_view.fetch_qf(self._qf) @property def body_link_pose(self) -> torch.Tensor: @@ -380,34 +258,7 @@ def body_link_pose(self) -> torch.Tensor: Returns: torch.Tensor: The poses of the links in the articulation with shape (N, num_links, 7). """ - if self.device.type == "cpu": - from embodichain.lab.sim.utility import get_dexsim_arenas - - arenas = get_dexsim_arenas() - for j, entity in enumerate(self.entities): - - link_pose = np.zeros((self.num_links, 4, 4), dtype=np.float32) - for i, link_name in enumerate(self.link_names): - pose = entity.get_link_pose(link_name) - arena_pose = arenas[j].get_root_node().get_local_pose() - pose[:2, 3] -= arena_pose[:2, 3] - link_pose[i] = pose - - link_pose = torch.from_numpy(link_pose) - xyz = link_pose[:, :3, 3] - quat = quat_from_matrix(link_pose[:, :3, :3]) - self._body_link_pose[j][: self.num_links, :] = torch.cat( - (xyz, quat), dim=-1 - ) - return self._body_link_pose[:, : self.num_links, :] - else: - self.ps.gpu_fetch_link_data( - data=self._body_link_pose, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.LINK_GLOBAL_POSE, - ) - quat = convert_quat(self._body_link_pose[..., :4], to="wxyz") - return torch.cat((self._body_link_pose[..., 4:], quat), dim=-1) + return self.articulation_view.fetch_link_pose(self._body_link_pose) @property def body_link_vel(self) -> torch.Tensor: @@ -416,26 +267,11 @@ def body_link_vel(self) -> torch.Tensor: Returns: torch.Tensor: The poses of the links in the articulation with shape (N, num_links, 6). """ - if self.device.type == "cpu": - for i, entity in enumerate(self.entities): - self._body_link_vel[i][: self.num_links] = torch.from_numpy( - entity.get_link_general_velocities() - ) - return self._body_link_vel[:, : self.num_links, :] - else: - self.ps.gpu_fetch_link_data( - data=self._body_link_lin_vel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.LINK_LINEAR_VELOCITY, - ) - self.ps.gpu_fetch_link_data( - data=self._body_link_ang_vel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.LINK_ANGULAR_VELOCITY, - ) - self._body_link_vel[..., :3] = self._body_link_lin_vel - self._body_link_vel[..., 3:] = self._body_link_ang_vel - return self._body_link_vel[:, : self.num_links, :] + return self.articulation_view.fetch_link_velocity( + self._body_link_vel, + self._body_link_lin_vel, + self._body_link_ang_vel, + ) @property def joint_stiffness(self) -> torch.Tensor: @@ -577,7 +413,9 @@ def __init__( ) -> None: # Initialize world and physics scene self._world = dexsim.default_world() - self._ps = self._world.get_physics_scene() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() self.cfg = cfg self._entities = entities @@ -807,6 +645,16 @@ def body_data(self) -> ArticulationData: """ return self._data + def _entity_link_name(self, env_idx: int, link_name: str) -> str: + """Resolve a canonical link name to the backend entity's local name.""" + if isinstance(env_idx, torch.Tensor): + env_idx = int(env_idx.detach().cpu().item()) + entity = self._entities[int(env_idx)] + view = self._data.articulation_view + if hasattr(view, "entity_link_name"): + return view.entity_link_name(entity, link_name) + return link_name + @property def root_state(self) -> torch.Tensor: """Get the root state of the articulation. @@ -915,47 +763,23 @@ def set_local_pose( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." ) - if self.device.type == "cpu": - pose = pose.cpu() - if pose.dim() == 2 and pose.shape[1] == 7: - pose_matrix = torch.eye(4).unsqueeze(0).repeat(pose.shape[0], 1, 1) - pose_matrix[:, :3, 3] = pose[:, :3] - pose_matrix[:, :3, :3] = matrix_from_quat(pose[:, 3:7]) - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_local_pose(pose_matrix[i]) - elif pose.dim() == 3 and pose.shape[1:] == (4, 4): - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_local_pose(pose[i]) - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - # TODO: in manual physics mode, the update should be explicitly called after - # setting the pose to synchronize the state to renderer. - + 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 = quat_from_matrix(pose[:, :3, :3]) + target_pose = torch.cat((xyz, quat), dim=-1).to( + device=self.device, dtype=torch.float32 + ) else: - if pose.dim() == 2 and pose.shape[1] == 7: - xyz = pose[:, :3] - quat = convert_quat(pose[:, 3:7], to="xyzw") - elif pose.dim() == 3 and pose.shape[1:] == (4, 4): - xyz = pose[:, :3, 3] - quat = quat_from_matrix(pose[:, :3, :3]) - quat = convert_quat(quat, to="xyzw") - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - # we should keep `pose_` life cycle to the end of the function. - pose_ = torch.cat((quat, xyz), dim=-1) - indices = self.body_data.gpu_indices[local_env_ids] - self._ps.gpu_apply_root_data( - data=pose_, - gpu_indices=indices, - data_type=ArticulationGPUAPIWriteType.ROOT_GLOBAL_POSE, + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." ) - self._ps.gpu_compute_articulation_kinematic(gpu_indices=indices) - self._world.update(0.001) + return + + self._data.articulation_view.apply_root_pose(target_pose, local_env_ids) + if self.device.type == "cpu" and not self._data.is_newton_backend: + self._world.update(0.001) def get_local_pose(self, to_matrix=False) -> torch.Tensor: """Get local pose (root link pose) of the articulation. @@ -1095,44 +919,16 @@ def set_qpos( f"env_ids: {local_env_ids}, qpos.shape: {qpos.shape}" ) - if self.device.type == "cpu": - for i, env_idx in enumerate(local_env_ids): - setter = ( - self._entities[env_idx].set_target_qpos - if target - else self._entities[env_idx].set_current_qpos - ) - setter(qpos[i].numpy(), local_joint_ids.numpy()) - else: - limits = self.body_data.qpos_limits[0].T - # clamp qpos to limits - lower_limits = limits[0][local_joint_ids] - upper_limits = limits[1][local_joint_ids] - qpos = qpos.clamp(lower_limits, upper_limits) - - data_type = ( - ArticulationGPUAPIWriteType.JOINT_TARGET_POSITION - if target - else ArticulationGPUAPIWriteType.JOINT_POSITION - ) - - # Always fetch the latest data to avoid stale values - if target: - qpos_set = self.body_data._target_qpos - else: - qpos_set = self.body_data._qpos - - if not isinstance(local_env_ids, torch.Tensor): - local_env_ids = torch.as_tensor( - local_env_ids, dtype=torch.long, device=self.device - ) - indices = self.body_data.gpu_indices[local_env_ids] - qpos_set[local_env_ids[:, None], local_joint_ids] = qpos - self._ps.gpu_apply_joint_data( - data=qpos_set, - gpu_indices=indices, - data_type=data_type, - ) + limits = self.body_data.qpos_limits[0].T + lower_limits = limits[0][local_joint_ids] + upper_limits = limits[1][local_joint_ids] + qpos = qpos.clamp(lower_limits, upper_limits) + self._data.articulation_view.apply_qpos( + qpos, + local_env_ids, + local_joint_ids, + target=target, + ) def get_qvel(self, target: bool = False) -> torch.Tensor: """Get the current velocities (qvel) or target velocities (target_qvel) of the articulation. @@ -1165,6 +961,14 @@ def set_qvel( """ local_env_ids = self._all_indices if env_ids is None else env_ids + if not isinstance(qvel, torch.Tensor): + qvel = torch.as_tensor(qvel, dtype=torch.float32, device=self.device) + else: + qvel = qvel.to(device=self.device, dtype=torch.float32) + + if qvel.dim() == 1: + qvel = qvel.unsqueeze(0) + if len(local_env_ids) != len(qvel): logger.log_error( f"Length of env_ids {len(local_env_ids)} does not match qvel length {len(qvel)}." @@ -1179,40 +983,14 @@ def set_qvel( joint_ids, dtype=torch.int32, device=self.device ) else: - local_joint_ids = joint_ids - - if self.device.type == "cpu": - for i, env_idx in enumerate(local_env_ids): - setter = ( - self._entities[env_idx].set_target_qvel - if target - else self._entities[env_idx].set_current_qvel - ) - setter(qvel[i].numpy(), local_joint_ids) - else: - data_type = ( - ArticulationGPUAPIWriteType.JOINT_TARGET_VELOCITY - if target - else ArticulationGPUAPIWriteType.JOINT_VELOCITY - ) - - # Always fetch the latest data to avoid stale values - if target: - qvel_set = self.body_data._target_qvel - else: - qvel_set = self.body_data._qvel + local_joint_ids = joint_ids.to(device=self.device, dtype=torch.int32) - if not isinstance(local_env_ids, torch.Tensor): - local_env_ids = torch.as_tensor( - local_env_ids, dtype=torch.long, device=self.device - ) - indices = self.body_data.gpu_indices[local_env_ids] - qvel_set[local_env_ids[:, None], local_joint_ids] = qvel - self._ps.gpu_apply_joint_data( - data=qvel_set, - gpu_indices=indices, - data_type=data_type, - ) + self._data.articulation_view.apply_qvel( + qvel, + local_env_ids, + local_joint_ids, + target=target, + ) def set_qf( self, @@ -1229,30 +1007,31 @@ def set_qf( """ local_env_ids = self._all_indices if env_ids is None else env_ids + if not isinstance(qf, torch.Tensor): + qf = torch.as_tensor(qf, dtype=torch.float32, device=self.device) + else: + qf = qf.to(device=self.device, dtype=torch.float32) + + if qf.dim() == 1: + qf = qf.unsqueeze(0) + if len(local_env_ids) != len(qf): logger.log_error( f"Length of env_ids {len(local_env_ids)} does not match qf length {len(qf)}." ) - if self.device.type == "cpu": - local_joint_ids = np.arange(self.dof) if joint_ids is None else joint_ids - for i, env_idx in enumerate(local_env_ids): - setter = self._entities[env_idx].set_current_qf - setter(qf[i].numpy(), local_joint_ids) - else: - indices = self.body_data.gpu_indices[local_env_ids] - if joint_ids is None: - qf_set = self.body_data._qf[local_env_ids] - qf_set[:, : self.dof] = qf - else: - self.body_data.qf - qf_set = self.body_data._qf[local_env_ids] - qf_set[:, joint_ids] = qf - self._ps.gpu_apply_joint_data( - data=qf_set, - gpu_indices=indices, - data_type=ArticulationGPUAPIWriteType.JOINT_FORCE, + if joint_ids is None: + local_joint_ids = torch.arange( + self.dof, device=self.device, dtype=torch.int32 ) + elif not isinstance(joint_ids, torch.Tensor): + local_joint_ids = torch.as_tensor( + joint_ids, dtype=torch.int32, device=self.device + ) + else: + local_joint_ids = joint_ids.to(device=self.device, dtype=torch.int32) + + self._data.articulation_view.apply_qf(qf, local_env_ids, local_joint_ids) def set_mass( self, @@ -1282,7 +1061,11 @@ def set_mass( for i, env_idx in enumerate(local_env_ids): for j, name in enumerate(link_names): - self._entities[env_idx].set_mass(name, mass[i, j].item()) + if self._data.is_newton_backend: + local_name = self._entity_link_name(env_idx, name) + self._entities[env_idx].set_link_mass(local_name, mass[i, j].item()) + else: + self._entities[env_idx].set_mass(name, mass[i, j].item()) def get_mass( self, @@ -1316,9 +1099,15 @@ def get_mass( ) for i, env_idx in enumerate(local_env_ids): for j, name in enumerate(link_names): - mass_tensor[i, j] = ( - self._entities[env_idx].get_physical_body(name).get_mass() - ) + if self._data.is_newton_backend: + local_name = self._entity_link_name(env_idx, name) + mass_tensor[i, j] = self._entities[env_idx].get_link_mass( + local_name + ) + else: + mass_tensor[i, j] = ( + self._entities[env_idx].get_physical_body(name).get_mass() + ) return mass_tensor def get_link_physical_attr( @@ -1351,7 +1140,11 @@ def get_link_physical_attr( attrs: list[PhysicalAttr] = [] for env_idx in local_env_ids: for name in matched_link_names: - attrs.append(self._entities[env_idx].get_physical_attr(name)) + attrs.append( + self._entities[env_idx].get_physical_attr( + self._entity_link_name(env_idx, name) + ) + ) return attrs def set_link_physical_attr( @@ -1398,7 +1191,9 @@ def set_link_physical_attr( for env_idx in local_env_ids: for name in matched_link_names: self._entities[env_idx].set_physical_attr( - physical_attr, name, is_replace_inertial=replace_inertial + physical_attr, + self._entity_link_name(env_idx, name), + is_replace_inertial=replace_inertial, ) def set_joint_drive( @@ -1578,10 +1373,7 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used. """ local_env_ids = self._all_indices if env_ids is None else env_ids - zeros = torch.zeros((len(local_env_ids), self.dof), device=self.device) - self.set_qvel(zeros, env_ids=local_env_ids) - self.set_qvel(zeros, env_ids=local_env_ids, target=True) - self.set_qf(zeros, env_ids=local_env_ids) + self._data.articulation_view.clear_dynamics(local_env_ids) def reallocate_body_data(self) -> None: """Reallocate body data tensors to match the current articulation state in the GPU physics scene.""" @@ -1666,11 +1458,9 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: self.clear_dynamics(env_ids=local_env_ids) - if self.device.type == "cuda": - self._ps.gpu_compute_articulation_kinematic( - gpu_indices=self.body_data.gpu_indices[local_env_ids] - ) - self._world.update(0.001) + self._data.articulation_view.compute_kinematics(local_env_ids) + if self.device.type == "cpu" and not self._data.is_newton_backend: + self._world.update(0.001) def _set_default_joint_drive(self) -> None: """Set default joint drive parameters based on the configuration.""" @@ -2050,4 +1840,7 @@ def destroy(self) -> None: if len(arenas) == 0: arenas = [env] for i, entity in enumerate(self._entities): - arenas[i].remove_articulation(entity) + if self._data.is_newton_backend: + arenas[i].remove_skeleton(entity) + else: + arenas[i].remove_articulation(entity) diff --git a/embodichain/lab/sim/objects/backends/__init__.py b/embodichain/lab/sim/objects/backends/__init__.py new file mode 100644 index 00000000..538afeb1 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/__init__.py @@ -0,0 +1,37 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from .base import ArticulationViewBase, RigidBodyViewBase +from .default import DefaultArticulationView, DefaultRigidBodyView +from .newton import ( + NewtonArticulationView, + NewtonRigidBodyView, + apply_collision_filter_for_entities, + apply_collision_filter_for_envs, + is_newton_scene, +) + +__all__ = [ + "ArticulationViewBase", + "RigidBodyViewBase", + "DefaultArticulationView", + "DefaultRigidBodyView", + "NewtonArticulationView", + "NewtonRigidBodyView", + "apply_collision_filter_for_entities", + "apply_collision_filter_for_envs", + "is_newton_scene", +] diff --git a/embodichain/lab/sim/objects/backends/base.py b/embodichain/lab/sim/objects/backends/base.py new file mode 100644 index 00000000..654eb701 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/base.py @@ -0,0 +1,348 @@ +# ---------------------------------------------------------------------------- +# 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 +from functools import cached_property + +import torch + +__all__ = ["RigidBodyViewBase", "ArticulationViewBase"] + + +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.""" + ... + + @property + def can_apply_pose(self) -> bool: + """Whether world poses can be written through the backend view.""" + return self.is_ready + + @property + def can_fetch_pose(self) -> bool: + """Whether world poses can be read through the backend view.""" + return self.is_ready + + # -- Body ID Management ------------------------------------------------- + + @cached_property + @abstractmethod + def body_ids(self) -> list[int]: + """Backend body IDs for all managed entities.""" + ... + + @cached_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) -> torch.Tensor: + """Return body IDs for the given entity indices.""" + ... + + # -- Pose --------------------------------------------------------------- + + @abstractmethod + def fetch_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch poses into ``data`` as ``(N, 7)`` in ``(x, y, z, qx, qy, qz, qw)``.""" + ... + + @abstractmethod + def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply poses from ``(N, 7)`` tensor in ``(x, y, z, qx, qy, qz, qw)``.""" + ... + + # -- Center of Mass (local) --------------------------------------------- + + @abstractmethod + def fetch_com_local_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch center-of-mass local poses into ``data`` as ``(N, 7)``.""" + ... + + @abstractmethod + def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply center-of-mass local poses from ``(N, 7)`` tensor.""" + ... + + # -- Velocity ----------------------------------------------------------- + + @abstractmethod + def fetch_linear_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch linear velocities into ``data`` as ``(N, 3)``.""" + ... + + @abstractmethod + def fetch_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch angular velocities into ``data`` as ``(N, 3)``.""" + ... + + @abstractmethod + def apply_linear_velocity(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Set linear velocities from ``(N, 3)`` tensor.""" + ... + + @abstractmethod + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + """Set angular velocities from ``(N, 3)`` tensor.""" + ... + + # -- Acceleration ------------------------------------------------------- + + @abstractmethod + def fetch_linear_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch linear accelerations into ``data`` as ``(N, 3)``.""" + ... + + @abstractmethod + def fetch_angular_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch angular accelerations into ``data`` as ``(N, 3)``.""" + ... + + # -- Force & Torque ----------------------------------------------------- + + @abstractmethod + def apply_force(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply external forces ``(N, 3)``. One-shot — consumed on next step.""" + ... + + @abstractmethod + def apply_torque(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply external torques ``(N, 3)``. One-shot — consumed on next step.""" + ... + + # -- Physical Properties ------------------------------------------------- + + @abstractmethod + def fetch_mass( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch masses into ``data`` as ``(N, 1)``.""" + ... + + @abstractmethod + def apply_mass(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply masses from ``(N, 1)`` tensor.""" + ... + + @abstractmethod + def fetch_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch inertia diagonals into ``data`` as ``(N, 3)``.""" + ... + + @abstractmethod + def apply_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + """Apply inertia diagonals from ``(N, 3)`` tensor.""" + ... + + @abstractmethod + def fetch_friction( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch friction coefficients into ``data`` as ``(N, 1)``.""" + ... + + @abstractmethod + def apply_friction(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply friction coefficients from ``(N, 1)`` tensor.""" + ... + + @abstractmethod + def fetch_restitution( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch restitution coefficients into ``data`` as ``(N, 1)``.""" + ... + + @abstractmethod + def apply_restitution(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply restitution coefficients from ``(N, 1)`` tensor.""" + ... + + +class ArticulationViewBase(ABC): + """Abstract interface for physics-backend articulation data access. + + Public root/link poses use EmbodiChain convention: + ``(x, y, z, qx, qy, qz, qw)``. + """ + + @property + @abstractmethod + def is_ready(self) -> bool: + """Whether backend runtime data can be accessed through batch APIs.""" + ... + + @property + def is_newton_backend(self) -> bool: + """Whether this view targets the DexSim Newton backend.""" + return False + + @property + @abstractmethod + def articulation_ids_tensor(self) -> torch.Tensor | None: + """Backend articulation ids as an int32 tensor, if the backend uses ids.""" + ... + + @abstractmethod + def select_articulation_ids( + self, env_ids: Sequence[int] | torch.Tensor + ) -> torch.Tensor: + """Return backend articulation ids for the given environment ids.""" + ... + + @abstractmethod + def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: + """Fetch root poses into ``data`` and return a view/result tensor.""" + ... + + @abstractmethod + def fetch_root_linear_velocity(self, data: torch.Tensor) -> torch.Tensor: + """Fetch root linear velocities into ``data`` and return a tensor.""" + ... + + @abstractmethod + def fetch_root_angular_velocity(self, data: torch.Tensor) -> torch.Tensor: + """Fetch root angular velocities into ``data`` and return a tensor.""" + ... + + @abstractmethod + def fetch_qpos(self, data: torch.Tensor) -> torch.Tensor: + """Fetch current joint positions into ``data``.""" + ... + + @abstractmethod + def fetch_target_qpos(self, data: torch.Tensor) -> torch.Tensor: + """Fetch target joint positions into ``data``.""" + ... + + @abstractmethod + def fetch_qvel(self, data: torch.Tensor) -> torch.Tensor: + """Fetch current joint velocities into ``data``.""" + ... + + @abstractmethod + def fetch_target_qvel(self, data: torch.Tensor) -> torch.Tensor: + """Fetch target joint velocities into ``data``.""" + ... + + @abstractmethod + def fetch_qacc(self, data: torch.Tensor) -> torch.Tensor: + """Fetch current joint accelerations into ``data``.""" + ... + + @abstractmethod + def fetch_qf(self, data: torch.Tensor) -> torch.Tensor: + """Fetch current joint forces into ``data``.""" + ... + + @abstractmethod + def fetch_link_pose(self, data: torch.Tensor) -> torch.Tensor: + """Fetch link poses into ``data``.""" + ... + + @abstractmethod + def fetch_link_velocity( + self, + data: torch.Tensor, + linear_data: torch.Tensor, + angular_data: torch.Tensor, + ) -> torch.Tensor: + """Fetch link velocities into ``data`` using provided scratch buffers.""" + ... + + @abstractmethod + def apply_root_pose( + self, pose: torch.Tensor, env_ids: Sequence[int] | torch.Tensor + ) -> None: + """Apply root poses from ``(N, 7)`` or equivalent backend convention.""" + ... + + @abstractmethod + def apply_qpos( + self, + qpos: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + """Apply joint positions for selected envs and joints.""" + ... + + @abstractmethod + def apply_qvel( + self, + qvel: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + """Apply joint velocities for selected envs and joints.""" + ... + + @abstractmethod + def apply_qf( + self, + qf: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + ) -> None: + """Apply joint forces for selected envs and joints.""" + ... + + @abstractmethod + def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + """Clear joint velocities, target velocities, and forces.""" + ... + + @abstractmethod + def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + """Refresh articulation kinematics if required by the backend.""" + ... diff --git a/embodichain/lab/sim/objects/backends/default.py b/embodichain/lab/sim/objects/backends/default.py new file mode 100644 index 00000000..0c9fc8be --- /dev/null +++ b/embodichain/lab/sim/objects/backends/default.py @@ -0,0 +1,725 @@ +# ---------------------------------------------------------------------------- +# 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 +from functools import cached_property + +import numpy as np +import torch + +from dexsim.models import MeshObject +from dexsim.engine import Articulation, PhysicsScene +from dexsim.types import ( + ArticulationGPUAPIReadType, + ArticulationGPUAPIWriteType, + RigidBodyGPUAPIReadType, + RigidBodyGPUAPIWriteType, +) +from embodichain.lab.sim.objects.backends.base import ( + ArticulationViewBase, + RigidBodyViewBase, +) +from embodichain.utils.math import ( + convert_quat, + matrix_from_quat, + quat_from_matrix, +) + +__all__ = ["DefaultRigidBodyView", "DefaultArticulationView"] + + +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: PhysicsScene, + 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 ----------------------------------------- + + @cached_property + def body_ids(self) -> list[int]: + if self._is_gpu: + return self._gpu_indices.cpu().tolist() + return list(range(len(self.entities))) + + @cached_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) -> torch.Tensor: + return self.body_ids_tensor[indices] + + # -- RigidBodyViewBase: pose --------------------------------------------- + + def fetch_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + if self._is_gpu: + indices = self.body_ids_tensor if body_ids is None else body_ids + self.ps.gpu_fetch_rigid_body_data( + data=data, + gpu_indices=indices.to(device=self.device, dtype=torch.int32), + data_type=RigidBodyGPUAPIReadType.POSE, + ) + # Convert (qx, qy, qz, qw, x, y, z) -> (x, y, z, qx, qy, qz, qw) + quat = data[:, :4].clone() + xyz = data[:, 4:7].clone() + data[:, :3] = xyz + data[:, 3:7] = quat + return + + entities = self._select_entities(body_ids) + data_np = data.cpu().numpy() + for i, entity in enumerate(entities): + data_np[i, :3] = entity.get_location() + data_np[i, 3:7] = entity.get_rotation_quat() + + def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> 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) + torch.cuda.synchronize(self.device) + self.ps.gpu_apply_rigid_body_data( + data=gpu_pose.clone(), + gpu_indices=body_ids.to(device=self.device, dtype=torch.int32), + data_type=RigidBodyGPUAPIWriteType.POSE, + ) + return + + # CPU: convert (x, y, z, qx, qy, qz, qw) -> 4x4 matrix per entity + indices = body_ids.detach().cpu().tolist() + 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: center of mass (local) --------------------------- + + def fetch_com_local_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + entities = self._select_entities(body_ids) + for i, entity in enumerate(entities): + pos, quat = entity.get_physical_body().get_cmass_local_pose() + data[i, :3] = torch.as_tensor(pos, dtype=torch.float32, device=self.device) + data[i, 3:7] = torch.as_tensor( + convert_quat(quat, to="xyzw"), + dtype=torch.float32, + device=self.device, + ) + + def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + data = data.to(dtype=torch.float32) + indices = body_ids.detach().cpu().tolist() + data_cpu = data.cpu().numpy() + for i, idx in enumerate(indices): + pos = data_cpu[i, :3] + quat = convert_quat(data_cpu[i, 3:7], to="wxyz") + self.entities[idx].get_physical_body().set_cmass_local_pose(pos, quat) + + # -- RigidBodyViewBase: velocity ----------------------------------------- + + def fetch_linear_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3( + RigidBodyGPUAPIReadType.LINEAR_VELOCITY, + "get_linear_velocity", + data, + body_ids, + ) + + def fetch_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3( + RigidBodyGPUAPIReadType.ANGULAR_VELOCITY, + "get_angular_velocity", + data, + body_ids, + ) + + def apply_linear_velocity(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, + "set_linear_velocity", + data, + body_ids, + ) + + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, + "set_angular_velocity", + data, + body_ids, + ) + + # -- RigidBodyViewBase: acceleration ------------------------------------- + + def fetch_linear_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3( + RigidBodyGPUAPIReadType.LINEAR_ACCELERATION, + "get_linear_acceleration", + data, + body_ids, + ) + + def fetch_angular_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3( + RigidBodyGPUAPIReadType.ANGULAR_ACCELERATION, + "get_angular_acceleration", + data, + body_ids, + ) + + # -- RigidBodyViewBase: force & torque ----------------------------------- + + def apply_force(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.FORCE, + "add_force", + data, + body_ids, + ) + + def apply_torque(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_vec3( + RigidBodyGPUAPIWriteType.TORQUE, + "add_torque", + data, + body_ids, + ) + + # -- RigidBodyViewBase: physical properties ------------------------------ + + def fetch_mass( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + entities = self._select_entities(body_ids) + for i, entity in enumerate(entities): + data[i, 0] = entity.get_physical_body().get_mass() + + def apply_mass(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + data_cpu = data.to(dtype=torch.float32).cpu().numpy() + indices = body_ids.detach().cpu().tolist() + for i, idx in enumerate(indices): + self.entities[int(idx)].get_physical_body().set_mass(data_cpu[i, 0]) + + def fetch_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + entities = self._select_entities(body_ids) + for i, entity in enumerate(entities): + inertia = entity.get_physical_body().get_mass_space_inertia_tensor() + data[i, :3] = torch.as_tensor( + inertia, dtype=torch.float32, device=self.device + ) + + def apply_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + data_cpu = data.to(dtype=torch.float32).cpu().numpy() + indices = body_ids.detach().cpu().tolist() + for i, idx in enumerate(indices): + self.entities[int(idx)].get_physical_body().set_mass_space_inertia_tensor( + data_cpu[i] + ) + + def fetch_friction( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + entities = self._select_entities(body_ids) + for i, entity in enumerate(entities): + data[i, 0] = entity.get_physical_body().get_dynamic_friction() + + def apply_friction(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + data_cpu = data.to(dtype=torch.float32).cpu().numpy() + indices = body_ids.detach().cpu().tolist() + for i, idx in enumerate(indices): + self.entities[int(idx)].get_physical_body().set_dynamic_friction( + data_cpu[i, 0] + ) + self.entities[int(idx)].get_physical_body().set_static_friction( + data_cpu[i, 0] + ) + + def fetch_restitution( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + entities = self._select_entities(body_ids) + for i, entity in enumerate(entities): + data[i, 0] = entity.get_physical_body().get_restitution() + + def apply_restitution(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + data_cpu = data.to(dtype=torch.float32).cpu().numpy() + indices = body_ids.detach().cpu().tolist() + for i, idx in enumerate(indices): + self.entities[int(idx)].get_physical_body().set_restitution(data_cpu[i, 0]) + + # -- Internal helpers ---------------------------------------------------- + + def _select_entities(self, body_ids: torch.Tensor | None) -> list[MeshObject]: + """Select entities by body IDs (entity list indices for CPU).""" + if body_ids is None: + return self.entities + body_ids = body_ids.detach().cpu().tolist() + return [self.entities[int(i)] for i in body_ids] + + def _fetch_vec3( + self, + gpu_read_type, + cpu_method: str, + data: torch.Tensor, + body_ids: torch.Tensor | None, + ) -> None: + """Fetch a vec3 field from GPU or CPU entities.""" + if self._is_gpu: + indices = self.body_ids_tensor if body_ids is None else body_ids + self.ps.gpu_fetch_rigid_body_data( + data=data, + gpu_indices=indices.to(device=self.device, dtype=torch.int32), + data_type=gpu_read_type, + ) + return + + entities = self._select_entities(body_ids) + data_np = data.cpu().numpy() + for i, entity in enumerate(entities): + data_np[i] = getattr(entity, cpu_method)() + + def _apply_vec3( + self, + gpu_write_type, + cpu_method: str, + data: torch.Tensor, + body_ids: torch.Tensor, + ) -> None: + """Apply a vec3 field to GPU or CPU entities.""" + data = data.to(dtype=torch.float32) + if self._is_gpu: + torch.cuda.synchronize(self.device) + self.ps.gpu_apply_rigid_body_data( + data=data, + gpu_indices=body_ids.to(device=self.device, dtype=torch.int32), + data_type=gpu_write_type, + ) + return + + indices = body_ids.detach().cpu().tolist() + data_cpu = data.cpu().numpy() + for i, idx in enumerate(indices): + getattr(self.entities[idx], cpu_method)(data_cpu[i]) + + +class DefaultArticulationView(ArticulationViewBase): + """Default DexSim backend articulation data adapter.""" + + def __init__( + self, + entities: Sequence[Articulation], + ps: PhysicsScene, + device: torch.device, + ) -> None: + self.entities = list(entities) + self.ps = ps + self.device = device + self._is_gpu = device.type == "cuda" + + self.dof = self.entities[0].get_dof() + self.num_links = self.entities[0].get_links_num() + self.link_names = self.entities[0].get_link_names() + + if self._is_gpu: + self._gpu_indices = torch.as_tensor( + [entity.get_gpu_index() for entity in self.entities], + dtype=torch.int32, + device=self.device, + ) + max_dof = self.ps.gpu_get_articulation_max_dof() + else: + self._gpu_indices = None + max_dof = self.dof + + self._qpos_apply = torch.zeros( + (len(self.entities), max_dof), dtype=torch.float32, device=self.device + ) + self._target_qpos_apply = torch.zeros_like(self._qpos_apply) + self._qvel_apply = torch.zeros_like(self._qpos_apply) + self._target_qvel_apply = torch.zeros_like(self._qpos_apply) + self._qf_apply = torch.zeros_like(self._qpos_apply) + + @property + def is_ready(self) -> bool: + return True + + @property + def articulation_ids_tensor(self) -> torch.Tensor | None: + return self._gpu_indices + + def select_articulation_ids( + self, env_ids: Sequence[int] | torch.Tensor + ) -> torch.Tensor: + if self._gpu_indices is None: + return torch.as_tensor(env_ids, dtype=torch.int32, device=self.device) + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + return self._gpu_indices[env_ids.to(device=self.device, dtype=torch.long)] + + def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: + if self._is_gpu: + self.ps.gpu_fetch_root_data( + data=data, + gpu_indices=self._gpu_indices, + data_type=ArticulationGPUAPIReadType.ROOT_GLOBAL_POSE, + ) + data[:, :4] = convert_quat(data[:, :4], to="wxyz") + return data[:, [4, 5, 6, 0, 1, 2, 3]] + + root_pose = torch.as_tensor( + np.array([entity.get_local_pose() for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + xyzs = root_pose[:, :3, 3] + quats = quat_from_matrix(root_pose[:, :3, :3]) + return torch.cat((xyzs, quats), dim=-1) + + def fetch_root_linear_velocity(self, data: torch.Tensor) -> torch.Tensor: + if self._is_gpu: + self.ps.gpu_fetch_root_data( + data=data, + gpu_indices=self._gpu_indices, + data_type=ArticulationGPUAPIReadType.ROOT_LINEAR_VELOCITY, + ) + return data.clone() + return torch.as_tensor( + np.array([entity.get_root_link_velocity()[:3] for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + + def fetch_root_angular_velocity(self, data: torch.Tensor) -> torch.Tensor: + if self._is_gpu: + self.ps.gpu_fetch_root_data( + data=data, + gpu_indices=self._gpu_indices, + data_type=ArticulationGPUAPIReadType.ROOT_ANGULAR_VELOCITY, + ) + return data.clone() + return torch.as_tensor( + np.array([entity.get_root_link_velocity()[3:] for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + + def fetch_qpos(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_data(data, ArticulationGPUAPIReadType.JOINT_POSITION) + + def fetch_target_qpos(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_data( + data, ArticulationGPUAPIReadType.JOINT_TARGET_POSITION + ) + + def fetch_qvel(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_data(data, ArticulationGPUAPIReadType.JOINT_VELOCITY) + + def fetch_target_qvel(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_data( + data, ArticulationGPUAPIReadType.JOINT_TARGET_VELOCITY + ) + + def fetch_qacc(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_data( + data, ArticulationGPUAPIReadType.JOINT_ACCELERATION + ) + + def fetch_qf(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_data(data, ArticulationGPUAPIReadType.JOINT_FORCE) + + def fetch_link_pose(self, data: torch.Tensor) -> torch.Tensor: + if self._is_gpu: + self.ps.gpu_fetch_link_data( + data=data, + gpu_indices=self._gpu_indices, + data_type=ArticulationGPUAPIReadType.LINK_GLOBAL_POSE, + ) + quat = convert_quat(data[..., :4], to="wxyz") + return torch.cat((data[..., 4:], quat), dim=-1) + + from embodichain.lab.sim.utility import get_dexsim_arenas + + arenas = get_dexsim_arenas() + for j, entity in enumerate(self.entities): + link_pose = np.zeros((self.num_links, 4, 4), dtype=np.float32) + for i, link_name in enumerate(self.link_names): + pose = entity.get_link_pose(link_name) + arena_pose = arenas[j].get_root_node().get_local_pose() + pose[:2, 3] -= arena_pose[:2, 3] + link_pose[i] = pose + + link_pose_tensor = torch.from_numpy(link_pose) + xyz = link_pose_tensor[:, :3, 3] + quat = quat_from_matrix(link_pose_tensor[:, :3, :3]) + data[j][: self.num_links, :] = torch.cat((xyz, quat), dim=-1) + return data[:, : self.num_links, :] + + def fetch_link_velocity( + self, + data: torch.Tensor, + linear_data: torch.Tensor, + angular_data: torch.Tensor, + ) -> torch.Tensor: + if self._is_gpu: + self.ps.gpu_fetch_link_data( + data=linear_data, + gpu_indices=self._gpu_indices, + data_type=ArticulationGPUAPIReadType.LINK_LINEAR_VELOCITY, + ) + self.ps.gpu_fetch_link_data( + data=angular_data, + gpu_indices=self._gpu_indices, + data_type=ArticulationGPUAPIReadType.LINK_ANGULAR_VELOCITY, + ) + data[..., :3] = linear_data + data[..., 3:] = angular_data + return data[:, : self.num_links, :] + + for i, entity in enumerate(self.entities): + data[i][: self.num_links] = torch.from_numpy( + entity.get_link_general_velocities() + ) + return data[:, : self.num_links, :] + + def apply_root_pose( + self, pose: torch.Tensor, env_ids: Sequence[int] | torch.Tensor + ) -> None: + pose = pose.to(dtype=torch.float32) + if self._is_gpu: + xyz = pose[:, :3] + quat = convert_quat(pose[:, 3:7], to="xyzw") + data = torch.cat((quat, xyz), dim=-1) + indices = self.select_articulation_ids(env_ids) + self.ps.gpu_apply_root_data( + data=data, + gpu_indices=indices, + data_type=ArticulationGPUAPIWriteType.ROOT_GLOBAL_POSE, + ) + self.ps.gpu_compute_articulation_kinematic(gpu_indices=indices) + return + + pose_cpu = pose.cpu() + env_indices = self._env_indices_list(env_ids) + pose_matrix = torch.eye(4).unsqueeze(0).repeat(len(env_indices), 1, 1) + pose_matrix[:, :3, 3] = pose_cpu[:, :3] + pose_matrix[:, :3, :3] = matrix_from_quat(pose_cpu[:, 3:7]) + for i, env_idx in enumerate(env_indices): + self.entities[env_idx].set_local_pose(pose_matrix[i]) + + def apply_qpos( + self, + qpos: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + if self._is_gpu: + buffer = self._target_qpos_apply if target else self._qpos_apply + data_type = ( + ArticulationGPUAPIWriteType.JOINT_TARGET_POSITION + if target + else ArticulationGPUAPIWriteType.JOINT_POSITION + ) + self._apply_gpu_joint_rows(buffer, qpos, env_ids, joint_ids, data_type) + return + + joint_ids_np = self._joint_ids_numpy(joint_ids) + qpos_np = qpos.detach().cpu().numpy() + for i, env_idx in enumerate(self._env_indices_list(env_ids)): + entity = self.entities[env_idx] + setter = entity.set_target_qpos if target else entity.set_current_qpos + setter(qpos_np[i], joint_ids_np) + + def apply_qvel( + self, + qvel: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + if self._is_gpu: + buffer = self._target_qvel_apply if target else self._qvel_apply + data_type = ( + ArticulationGPUAPIWriteType.JOINT_TARGET_VELOCITY + if target + else ArticulationGPUAPIWriteType.JOINT_VELOCITY + ) + self._apply_gpu_joint_rows(buffer, qvel, env_ids, joint_ids, data_type) + return + + joint_ids_np = self._joint_ids_numpy(joint_ids) + qvel_np = qvel.detach().cpu().numpy() + for i, env_idx in enumerate(self._env_indices_list(env_ids)): + entity = self.entities[env_idx] + setter = entity.set_target_qvel if target else entity.set_current_qvel + setter(qvel_np[i], joint_ids_np) + + def apply_qf( + self, + qf: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + ) -> None: + if self._is_gpu: + self._apply_gpu_joint_rows( + self._qf_apply, + qf, + env_ids, + joint_ids, + ArticulationGPUAPIWriteType.JOINT_FORCE, + ) + return + + joint_ids_np = self._joint_ids_numpy(joint_ids) + qf_np = qf.detach().cpu().numpy() + for i, env_idx in enumerate(self._env_indices_list(env_ids)): + self.entities[env_idx].set_current_qf(qf_np[i], joint_ids_np) + + def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + zeros = torch.zeros( + (len(env_ids), self.dof), dtype=torch.float32, device=self.device + ) + joint_ids = torch.arange(self.dof, dtype=torch.int32, device=self.device) + self.apply_qvel(zeros, env_ids, joint_ids, target=False) + self.apply_qvel(zeros, env_ids, joint_ids, target=True) + self.apply_qf(zeros, env_ids, joint_ids) + + def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + if self._is_gpu: + self.ps.gpu_compute_articulation_kinematic( + gpu_indices=self.select_articulation_ids(env_ids) + ) + + def _fetch_joint_data(self, data: torch.Tensor, data_type) -> torch.Tensor: + if self._is_gpu: + self.ps.gpu_fetch_joint_data( + data=data, + gpu_indices=self._gpu_indices, + data_type=data_type, + ) + return data[:, : self.dof].clone() + + method_map = { + ArticulationGPUAPIReadType.JOINT_POSITION: lambda entity: entity.get_current_qpos(), + ArticulationGPUAPIReadType.JOINT_TARGET_POSITION: lambda entity: entity.get_current_qpos( + is_target=True + ), + ArticulationGPUAPIReadType.JOINT_VELOCITY: lambda entity: entity.get_current_qvel(), + ArticulationGPUAPIReadType.JOINT_TARGET_VELOCITY: lambda entity: entity.get_current_qvel( + is_target=True + ), + ArticulationGPUAPIReadType.JOINT_ACCELERATION: lambda entity: entity.get_current_qacc(), + ArticulationGPUAPIReadType.JOINT_FORCE: lambda entity: entity.get_current_qf(), + } + return torch.as_tensor( + np.array([method_map[data_type](entity) for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + + def _apply_gpu_joint_rows( + self, + buffer: torch.Tensor, + values: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + data_type, + ) -> None: + env_ids_tensor = self._env_ids_tensor(env_ids) + joint_ids_tensor = self._joint_ids_tensor(joint_ids) + buffer[env_ids_tensor[:, None], joint_ids_tensor] = values + self.ps.gpu_apply_joint_data( + data=buffer, + gpu_indices=self.select_articulation_ids(env_ids), + data_type=data_type, + ) + + def _env_ids_tensor(self, env_ids: Sequence[int] | torch.Tensor) -> torch.Tensor: + if not isinstance(env_ids, torch.Tensor): + return torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + return env_ids.to(device=self.device, dtype=torch.long) + + def _joint_ids_tensor( + self, joint_ids: Sequence[int] | torch.Tensor + ) -> torch.Tensor: + if not isinstance(joint_ids, torch.Tensor): + return torch.as_tensor(joint_ids, dtype=torch.long, device=self.device) + return joint_ids.to(device=self.device, dtype=torch.long) + + def _env_indices_list(self, env_ids: Sequence[int] | torch.Tensor) -> list[int]: + if isinstance(env_ids, torch.Tensor): + return env_ids.detach().cpu().to(dtype=torch.long).tolist() + return [int(env_idx) for env_idx in env_ids] + + def _joint_ids_numpy(self, joint_ids: Sequence[int] | torch.Tensor) -> np.ndarray: + if isinstance(joint_ids, torch.Tensor): + return joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) + return np.asarray(joint_ids, dtype=np.int32) diff --git a/embodichain/lab/sim/objects/backends/newton.py b/embodichain/lab/sim/objects/backends/newton.py new file mode 100644 index 00000000..f3b28b9c --- /dev/null +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -0,0 +1,822 @@ +# ---------------------------------------------------------------------------- +# 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.engine.newton_physics import NewtonPhysicsScene +from embodichain.lab.sim.objects.backends.base import ( + ArticulationViewBase, + RigidBodyViewBase, +) +from embodichain.utils import logger +from embodichain.utils.math import matrix_from_quat, quat_from_matrix + +__all__ = [ + "NewtonRigidBodyView", + "NewtonArticulationView", + "apply_collision_filter_for_entities", + "apply_collision_filter_for_envs", + "is_newton_scene", +] + +_UINT64_MAX = (1 << 64) - 1 +_INT32_MAX = (1 << 31) - 1 + + +def _normalize_native_handle(handle: int, owner: str) -> int: + value = int(handle) + if value < 0: + value &= _UINT64_MAX + if value > _UINT64_MAX: + logger.log_error(f"{owner} native handle is outside uint64 range: {value}.") + return value + + +def _collision_filter_rows(filter_data: torch.Tensor) -> torch.Tensor: + """Return contiguous ``(N, 4)`` int32 rows for the Newton scene API.""" + rows = filter_data.to(dtype=torch.int32) + if rows.ndim != 2 or rows.shape[-1] != 4: + logger.log_error( + "Collision filter data must have shape (N, 4), " f"got {tuple(rows.shape)}." + ) + if not rows.is_contiguous(): + rows = rows.contiguous() + return rows + + +def _resolve_body_ids_and_filter_rows_for_entities( + manager: object, + entities: Sequence[MeshObject], + filter_data: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + body_ids: list[int] = [] + rows: list[torch.Tensor] = [] + for i, entity in enumerate(entities): + entity_handle = _normalize_native_handle( + entity.get_native_handle(), "MeshObject" + ) + body_id = manager.body_id_for_entity(entity_handle) + if body_id is None: + entity.set_collision_filter_data( + filter_data[i].detach().cpu().numpy().astype(np.int64) + ) + continue + body_ids.append(int(body_id)) + rows.append(filter_data[i]) + + if len(rows) == 0: + empty_rows = filter_data.new_empty((0, filter_data.shape[-1])) + return torch.as_tensor(body_ids, dtype=torch.int32), empty_rows + + return torch.as_tensor(body_ids, dtype=torch.int32), torch.stack(rows, dim=0) + + +def apply_collision_filter_for_entities( + scene: NewtonPhysicsScene, + entities: Sequence[MeshObject], + filter_data: torch.Tensor, +) -> None: + """Batch-apply collision filters for a list of MeshObjects. + + Uses DexSim ``NewtonPhysicsScene.apply_collision_filter`` (vectorized meta + and shape-group writes on the DexSim side). + """ + if len(entities) == 0: + return + if len(entities) != len(filter_data): + logger.log_error( + "Entity count does not match collision filter row count " + f"({len(entities)} vs {len(filter_data)})." + ) + + rows = _collision_filter_rows(filter_data) + body_ids, valid_rows = _resolve_body_ids_and_filter_rows_for_entities( + scene.manager, entities, rows + ) + if len(body_ids) == 0: + return + body_ids = body_ids.to(device=rows.device) + scene.apply_collision_filter(body_ids, valid_rows.to(device=rows.device)) + + +def apply_collision_filter_for_envs( + scene: NewtonPhysicsScene, + entities_by_env: Sequence[Sequence[MeshObject]], + filter_data: torch.Tensor, + env_indices: Sequence[int], +) -> None: + """Batch-apply collision filters with one filter row per environment. + + Expands each env row to every ``MeshObject`` in that env (e.g. rigid groups). + """ + entities: list[MeshObject] = [] + rows: list[torch.Tensor] = [] + for i, env_idx in enumerate(env_indices): + row = filter_data[i] + for entity in entities_by_env[env_idx]: + entities.append(entity) + rows.append(row) + if not entities: + return + stacked = torch.stack(rows, dim=0) + apply_collision_filter_for_entities(scene, entities, stacked) + + +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, "batch_fetch_rigid_body_data") + and hasattr(scene, "batch_apply_rigid_body_data") + and hasattr(scene, "apply_collision_filter") + and hasattr(scene, "fetch_collision_filter") + ) + + +class NewtonRigidBodyView(RigidBodyViewBase): + """Adapter around DexSim Newton rigid body scene APIs. + + EmbodiChain public rigid-body pose convention is + ``(x, y, z, qx, qy, qz, qw)``. + DexSim Newton exposes the same pose convention through its unified rigid + data API. + """ + + _DATA_TYPE = None # lazily resolved NewtonRigidDataType + + def __init__( + self, + entities: Sequence[MeshObject], + scene: NewtonPhysicsScene, + 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 + ] + # Body IDs are resolved lazily because Newton's model is not built + # until finalization. Pre-finalization, ``body_id_for_entity()`` + # returns tentative IDs that may differ from the final interleaved + # layout. We track whether IDs have been resolved in the READY + # state and re-resolve once when the manager transitions. + self._body_ids: list[int] | None = None + self._body_ids_tensor: torch.Tensor | None = None + self._body_ids_finalized: bool = False + + # -- Lazy enum access --------------------------------------------------- + + @classmethod + def _get_data_type(cls): + """Lazily resolve *NewtonRigidDataType* to avoid eager import.""" + if cls._DATA_TYPE is None: + from dexsim.engine.newton_physics import NewtonRigidDataType + + cls._DATA_TYPE = NewtonRigidDataType + return cls._DATA_TYPE + + # -- RigidBodyViewBase: lifecycle ---------------------------------------- + + @property + def is_ready(self) -> bool: + manager = getattr(self.scene, "manager", None) + return ( + manager is not None + and getattr(getattr(manager, "lifecycle_state", None), "name", "") + == "READY" + ) + + @property + def _lifecycle_state_name(self) -> str: + manager = getattr(self.scene, "manager", None) + return getattr(getattr(manager, "lifecycle_state", None), "name", "") + + @property + def can_apply_pose(self) -> bool: + return self._lifecycle_state_name in ("BUILDER", "READY") + + @property + def can_fetch_pose(self) -> bool: + return self._lifecycle_state_name in ("BUILDER", "READY") + + # -- RigidBodyViewBase: body IDs ----------------------------------------- + + def _ensure_body_ids(self) -> None: + """Resolve body IDs from the Newton manager. + + Body IDs resolved before finalization may be tentative. Once the + manager transitions to READY, re-resolve to get the correct + interleaved layout. + """ + if self._body_ids_finalized: + return + if self._body_ids is not None and not self.is_ready: + return + ids = [self._resolve_body_id(entity) for entity in self.entities] + if any(bid < 0 or bid > _INT32_MAX for bid in ids): + logger.log_error( + "Newton rigid body view found an entity without a Newton body id." + ) + self._body_ids = ids + self._body_ids_tensor = torch.as_tensor( + ids, dtype=torch.int32, device=self.device + ) + if self.is_ready: + self._body_ids_finalized = True + + @property + def body_ids(self) -> list[int]: + self._ensure_body_ids() + return self._body_ids # type: ignore[return-value] + + @property + def body_ids_tensor(self) -> torch.Tensor: + self._ensure_body_ids() + return self._body_ids_tensor # type: ignore[return-value] + + def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> torch.Tensor: + self._ensure_body_ids() + if not isinstance(indices, torch.Tensor): + indices = torch.as_tensor(indices, dtype=torch.long, device=self.device) + return self._body_ids_tensor[indices.to(device=self.device, dtype=torch.long)] + + # -- RigidBodyViewBase: pose --------------------------------------------- + + def fetch_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self.scene.batch_fetch_rigid_body_data( + self._fetch_buffer(data), + self._resolve_body_ids(body_ids), + self._get_data_type().POSE, + ) + + def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().POSE, pose) + + # -- RigidBodyViewBase: center of mass (local) --------------------------- + + def fetch_com_local_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + data_type = getattr(self._get_data_type(), "COM_LOCAL_POSE", None) + self.scene.batch_fetch_rigid_body_data( + self._fetch_buffer(data), self._resolve_body_ids(body_ids), data_type + ) + + def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + data_type = getattr(self._get_data_type(), "COM_LOCAL_POSE", None) + self._apply_data(body_ids, data_type, data) + + # -- RigidBodyViewBase: velocity ----------------------------------------- + + def fetch_linear_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3(self._get_data_type().LINEAR_VELOCITY, data, body_ids) + + def fetch_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3(self._get_data_type().ANGULAR_VELOCITY, data, body_ids) + + def apply_linear_velocity(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().LINEAR_VELOCITY, data) + + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_data(body_ids, self._get_data_type().ANGULAR_VELOCITY, data) + + # -- RigidBodyViewBase: acceleration ------------------------------------- + + def fetch_linear_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3(self._get_data_type().LINEAR_ACCELERATION, data, body_ids) + + def fetch_angular_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3(self._get_data_type().ANGULAR_ACCELERATION, data, body_ids) + + # -- RigidBodyViewBase: force & torque ----------------------------------- + + def apply_force(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().FORCE, data) + + def apply_torque(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().TORQUE, data) + + # -- RigidBodyViewBase: physical properties ------------------------------ + + def fetch_mass( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_scalar(self._get_data_type().MASS, data, body_ids) + + def apply_mass(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().MASS, data) + + def fetch_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_vec3(self._get_data_type().INERTIA_DIAGONAL, data, body_ids) + + def apply_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_data(body_ids, self._get_data_type().INERTIA_DIAGONAL, data) + + def fetch_friction( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_scalar(self._get_data_type().FRICTION, data, body_ids) + + def apply_friction(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().FRICTION, data) + + def fetch_restitution( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_scalar(self._get_data_type().RESTITUTION, data, body_ids) + + def apply_restitution(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().RESTITUTION, data) + + # -- Collision filter ---------------------------------------------------- + + def fetch_collision_filter( + self, + data: torch.Tensor, + env_indices: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Fetch collision filter rows into ``data`` with shape ``(N, 4)``.""" + if env_indices is None: + env_indices = torch.arange(len(self.entities), device=self.device) + body_ids = self._resolve_body_ids(self.select_body_ids(env_indices)) + out = self._fetch_buffer(data) + self.scene.fetch_collision_filter(body_ids, out) + + def apply_collision_filter( + self, + filter_data: torch.Tensor, + env_indices: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Apply DexSim collision filter rows for selected env instances.""" + if env_indices is None: + env_indices = torch.arange(len(self.entities), device=self.device) + body_ids = self._resolve_body_ids(self.select_body_ids(env_indices)) + rows = _collision_filter_rows(filter_data.to(device=self.device)) + self.scene.apply_collision_filter(body_ids, rows) + + # -- Internal helpers ---------------------------------------------------- + + def _resolve_body_id(self, entity: MeshObject) -> int: + manager = getattr(self.scene, "manager", None) + if manager is not None and hasattr(entity, "get_native_handle"): + entity_handle = _normalize_native_handle( + entity.get_native_handle(), "MeshObject" + ) + body_id = manager.body_id_for_entity(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 _resolve_body_ids(self, body_ids: torch.Tensor | None) -> torch.Tensor: + """Return body IDs as a device int32 tensor for the Newton scene API. + + DexSim's batch API normalizes GPU-resident tensors without a host + round-trip, so the cached ``body_ids_tensor`` is passed straight + through. This avoids a per-call ``cuda -> cpu`` synchronization on the + per-step fetch/apply hot path. + """ + if body_ids is None: + self._ensure_body_ids() + return self._body_ids_tensor # type: ignore[return-value] + if not isinstance(body_ids, torch.Tensor): + body_ids = torch.as_tensor(body_ids, dtype=torch.int32, device=self.device) + return body_ids + + def _fetch_buffer(self, data: torch.Tensor) -> torch.Tensor: + """Validate and forward a caller-owned fetch buffer to the scene API.""" + if not data.is_contiguous(): + logger.log_error("Newton rigid body fetch buffers must be contiguous.") + return data + + def _fetch_vec3( + self, + data_type, + data: torch.Tensor, + body_ids: torch.Tensor | None = None, + ) -> None: + self.scene.batch_fetch_rigid_body_data( + self._fetch_buffer(data), self._resolve_body_ids(body_ids), data_type + ) + + # Scalar ``(N, 1)`` fields share the same fetch path as vec3 fields. + _fetch_scalar = _fetch_vec3 + + def _apply_data( + self, body_ids: torch.Tensor, data_type, data: torch.Tensor + ) -> None: + """Apply data to bodies via the unified Newton GPU API.""" + self.scene.batch_apply_rigid_body_data( + data.to(dtype=torch.float32).contiguous(), + self._resolve_body_ids(body_ids), + data_type, + ) + + +class NewtonArticulationView(ArticulationViewBase): + """Adapter around DexSim Newton articulation scene APIs.""" + + _DATA_TYPE = None + + def __init__( + self, + entities: Sequence[object], + scene: NewtonPhysicsScene, + device: torch.device, + ) -> None: + self.entities = list(entities) + self.scene = scene + self.device = device + self.dof = self.entities[0].get_dof() + self.num_links = self.entities[0].get_links_num() + self.link_names = self.entities[0].get_link_names() + self._articulation_ids = torch.as_tensor( + [entity.get_gpu_index() for entity in self.entities], + dtype=torch.int32, + device=self.device, + ) + self._link_body_ids: torch.Tensor | None = None + self._link_body_ids_finalized = False + + @classmethod + def _get_data_type(cls): + if cls._DATA_TYPE is None: + from dexsim.engine.newton_physics import NewtonArticulationDataType + + cls._DATA_TYPE = NewtonArticulationDataType + return cls._DATA_TYPE + + @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" + ) + + @property + def is_newton_backend(self) -> bool: + return True + + @property + def articulation_ids_tensor(self) -> torch.Tensor: + return self._articulation_ids + + def select_articulation_ids( + self, env_ids: Sequence[int] | torch.Tensor + ) -> torch.Tensor: + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + return self._articulation_ids[env_ids.to(device=self.device, dtype=torch.long)] + + def link_body_ids_for( + self, env_ids: Sequence[int] | torch.Tensor | None = None + ) -> torch.Tensor: + if self._link_body_ids_finalized is False: + rows = [] + for entity in self.entities: + row = [] + for link_name in self.link_names: + local_link_name = self.entity_link_name(entity, link_name) + link_meta = entity.dexsim_meta_links["links"][local_link_name] + body_id = ( + -1 if link_meta.body_id is None else int(link_meta.body_id) + ) + if body_id < 0 or body_id > _INT32_MAX: + logger.log_error( + f"Newton articulation link '{link_name}' has no valid body id." + ) + row.append(body_id) + rows.append(row) + self._link_body_ids = torch.as_tensor( + rows, dtype=torch.int32, device=self.device + ) + if self.is_ready: + self._link_body_ids_finalized = True + + assert self._link_body_ids is not None + if env_ids is None: + return self._link_body_ids.reshape(-1) + if not isinstance(env_ids, torch.Tensor): + env_ids = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + return self._link_body_ids[ + env_ids.to(device=self.device, dtype=torch.long) + ].reshape(-1) + + def entity_link_name(self, entity: object, link_name: str) -> str: + if link_name in getattr(entity, "dexsim_meta_links", {}).get("links", {}): + return link_name + link_idx = self.link_names.index(link_name) + return entity.get_link_names()[link_idx] + + def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: + if self.is_ready: + self._fetch(data, self._get_data_type().ROOT_GLOBAL_POSE) + return data.clone() + + root_pose = torch.as_tensor( + np.array([entity.get_local_pose() for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + xyzs = root_pose[:, :3, 3] + quats = quat_from_matrix(root_pose[:, :3, :3]) + return torch.cat((xyzs, quats), dim=-1) + + def fetch_root_linear_velocity(self, data: torch.Tensor) -> torch.Tensor: + if self.is_ready: + self._fetch(data, self._get_data_type().ROOT_LINEAR_VELOCITY) + return data.clone() + return torch.as_tensor( + np.array( + [ + entity.get_link_general_velocities(entity.get_root_link_name())[ + 0, :3 + ] + for entity in self.entities + ] + ), + dtype=torch.float32, + device=self.device, + ) + + def fetch_root_angular_velocity(self, data: torch.Tensor) -> torch.Tensor: + if self.is_ready: + self._fetch(data, self._get_data_type().ROOT_ANGULAR_VELOCITY) + return data.clone() + return torch.as_tensor( + np.array( + [ + entity.get_link_general_velocities(entity.get_root_link_name())[ + 0, 3: + ] + for entity in self.entities + ] + ), + dtype=torch.float32, + device=self.device, + ) + + def fetch_qpos(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_or_entity( + data, self._get_data_type().JOINT_POSITION, "get_current_qpos" + ) + + def fetch_target_qpos(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_or_entity( + data, self._get_data_type().JOINT_TARGET_POSITION, "get_target_qpos" + ) + + def fetch_qvel(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_or_entity( + data, self._get_data_type().JOINT_VELOCITY, "get_current_qvel" + ) + + def fetch_target_qvel(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_or_entity( + data, self._get_data_type().JOINT_TARGET_VELOCITY, "get_target_qvel" + ) + + def fetch_qacc(self, data: torch.Tensor) -> torch.Tensor: + return torch.zeros( + (len(self.entities), self.dof), dtype=torch.float32, device=self.device + ) + + def fetch_qf(self, data: torch.Tensor) -> torch.Tensor: + return self._fetch_joint_or_entity( + data, self._get_data_type().JOINT_FORCE, "get_current_qf" + ) + + def fetch_link_pose(self, data: torch.Tensor) -> torch.Tensor: + if self.is_ready: + flat_pose = data[:, : self.num_links, :].reshape(-1, 7) + self.scene.batch_fetch_articulation_data( + flat_pose, + self.link_body_ids_for(), + self._get_data_type().LINK_GLOBAL_POSE, + ) + return data[:, : self.num_links, :].clone() + + from embodichain.lab.sim.utility import get_dexsim_arenas + + arenas = get_dexsim_arenas() + for j, entity in enumerate(self.entities): + link_pose = np.zeros((self.num_links, 4, 4), dtype=np.float32) + for i, link_name in enumerate(self.link_names): + pose = entity.get_link_pose(self.entity_link_name(entity, link_name)) + arena_pose = arenas[j].get_root_node().get_local_pose() + pose[:2, 3] -= arena_pose[:2, 3] + link_pose[i] = pose + + link_pose_tensor = torch.from_numpy(link_pose) + xyz = link_pose_tensor[:, :3, 3] + quat = quat_from_matrix(link_pose_tensor[:, :3, :3]) + data[j][: self.num_links, :] = torch.cat((xyz, quat), dim=-1) + return data[:, : self.num_links, :] + + def fetch_link_velocity( + self, + data: torch.Tensor, + linear_data: torch.Tensor, + angular_data: torch.Tensor, + ) -> torch.Tensor: + if self.is_ready: + flat_lin = linear_data[:, : self.num_links, :].reshape(-1, 3) + flat_ang = angular_data[:, : self.num_links, :].reshape(-1, 3) + link_ids = self.link_body_ids_for() + self.scene.batch_fetch_articulation_data( + flat_lin, link_ids, self._get_data_type().LINK_LINEAR_VELOCITY + ) + self.scene.batch_fetch_articulation_data( + flat_ang, link_ids, self._get_data_type().LINK_ANGULAR_VELOCITY + ) + data[..., :3] = linear_data + data[..., 3:] = angular_data + return data[:, : self.num_links, :].clone() + + for i, entity in enumerate(self.entities): + data[i][: self.num_links] = torch.from_numpy( + entity.get_link_general_velocities() + ) + return data[:, : self.num_links, :] + + def apply_root_pose( + self, pose: torch.Tensor, env_ids: Sequence[int] | torch.Tensor + ) -> None: + pose_cpu = pose.to(dtype=torch.float32).cpu() + env_indices = self._env_indices_list(env_ids) + pose_matrix = torch.eye(4).unsqueeze(0).repeat(len(env_indices), 1, 1) + pose_matrix[:, :3, 3] = pose_cpu[:, :3] + pose_matrix[:, :3, :3] = matrix_from_quat(pose_cpu[:, 3:7]) + for i, env_idx in enumerate(env_indices): + self.entities[env_idx].set_local_pose(pose_matrix[i]) + + def apply_qpos( + self, + qpos: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + if self.is_ready: + data_type = ( + self._get_data_type().JOINT_TARGET_POSITION + if target + else self._get_data_type().JOINT_POSITION + ) + self._apply(qpos, data_type, env_ids, joint_ids) + return + + joint_ids_np = self._joint_ids_numpy(joint_ids) + qpos_np = qpos.detach().cpu().numpy() + for i, env_idx in enumerate(self._env_indices_list(env_ids)): + setter = ( + self.entities[env_idx].set_target_qpos + if target + else self.entities[env_idx].set_current_qpos + ) + setter(qpos_np[i], joint_ids_np) + + def apply_qvel( + self, + qvel: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + if self.is_ready: + data_type = ( + self._get_data_type().JOINT_TARGET_VELOCITY + if target + else self._get_data_type().JOINT_VELOCITY + ) + self._apply(qvel, data_type, env_ids, joint_ids) + return + + joint_ids_np = self._joint_ids_numpy(joint_ids) + qvel_np = qvel.detach().cpu().numpy() + for i, env_idx in enumerate(self._env_indices_list(env_ids)): + setter = ( + self.entities[env_idx].set_target_qvel + if target + else self.entities[env_idx].set_current_qvel + ) + setter(qvel_np[i], joint_ids_np) + + def apply_qf( + self, + qf: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + ) -> None: + if self.is_ready: + self._apply(qf, self._get_data_type().JOINT_FORCE, env_ids, joint_ids) + return + + joint_ids_np = self._joint_ids_numpy(joint_ids) + qf_np = qf.detach().cpu().numpy() + for i, env_idx in enumerate(self._env_indices_list(env_ids)): + self.entities[env_idx].set_current_qf(qf_np[i], joint_ids_np) + + def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + zeros = torch.zeros( + (len(env_ids), self.dof), dtype=torch.float32, device=self.device + ) + joint_ids = torch.arange(self.dof, dtype=torch.int32, device=self.device) + self.apply_qvel(zeros, env_ids, joint_ids, target=False) + self.apply_qvel(zeros, env_ids, joint_ids, target=True) + self.apply_qf(zeros, env_ids, joint_ids) + + def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + return + + def _fetch(self, data: torch.Tensor, data_type, joint_ids=None) -> None: + self.scene.batch_fetch_articulation_data( + data.contiguous(), + self._articulation_ids, + data_type, + self._joint_ids_numpy(joint_ids) if joint_ids is not None else None, + ) + + def _apply( + self, + data: torch.Tensor, + data_type, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + self.scene.batch_apply_articulation_data( + data.to(dtype=torch.float32).contiguous(), + self.select_articulation_ids(env_ids), + data_type, + self._joint_ids_numpy(joint_ids) if joint_ids is not None else None, + ) + + def _fetch_joint_or_entity( + self, data: torch.Tensor, data_type, entity_method: str + ) -> torch.Tensor: + if self.is_ready: + self._fetch(data, data_type) + return data[:, : self.dof].clone() + return torch.as_tensor( + np.array([getattr(entity, entity_method)() for entity in self.entities]), + dtype=torch.float32, + device=self.device, + ) + + def _joint_ids_numpy( + self, joint_ids: Sequence[int] | torch.Tensor | None + ) -> np.ndarray | None: + if joint_ids is None: + return None + if isinstance(joint_ids, torch.Tensor): + return joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) + return np.asarray(joint_ids, dtype=np.int32) + + def _env_indices_list(self, env_ids: Sequence[int] | torch.Tensor) -> list[int]: + if isinstance(env_ids, torch.Tensor): + return env_ids.detach().cpu().to(dtype=torch.long).tolist() + return [int(env_idx) for env_idx in env_ids] diff --git a/embodichain/lab/sim/objects/cloth_object.py b/embodichain/lab/sim/objects/cloth_object.py index 28db03cb..0a06138b 100644 --- a/embodichain/lab/sim/objects/cloth_object.py +++ b/embodichain/lab/sim/objects/cloth_object.py @@ -118,7 +118,9 @@ def __init__( device: torch.device = torch.device("cpu"), ) -> None: self._world: dexsim.World = dexsim.default_world() - self._ps = self._world.get_physics_scene() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() self._data = ClothBodyData(entities=entities, ps=self._ps, device=device) @@ -126,6 +128,9 @@ def __init__( self._world.update(0.001) super().__init__(cfg=cfg, entities=entities, device=device) + + self.reset() + self._set_default_collision_filter() def _set_default_collision_filter(self) -> None: diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index a44bb437..3ade71db 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import torch import dexsim import numpy as np @@ -23,9 +25,15 @@ 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, + apply_collision_filter_for_entities, + is_newton_scene, +) +from embodichain.lab.sim.objects.backends.base import RigidBodyViewBase from embodichain.lab.sim import ( VisualMaterial, VisualMaterialInst, @@ -35,13 +43,15 @@ from embodichain.utils.math import matrix_from_quat, quat_from_matrix, matrix_from_euler from embodichain.utils import logger +_UINT64_MAX = (1 << 64) - 1 + @dataclass class RigidBodyData: """Data manager for rigid body with body type of dynamic or kinematic. - Note: - 1. The pose data managed by dexsim is in the format of (qx, qy, qz, qw, x, y, z), but in SimulationManager, we use (x, y, z, qw, qx, qy, qz) format. + All pose/velocity/acceleration data uses EmbodiChain convention: + ``(x, y, z, qx, qy, qz, qw)``. """ def __init__( @@ -59,16 +69,19 @@ def __init__( self.num_instances = len(entities) self.device = device - # get gpu indices for the entities. - self.gpu_indices = ( - torch.as_tensor( - [entity.get_gpu_index() for entity in self.entities], - dtype=torch.int32, - device=self.device, + # Create the appropriate backend view. + if is_newton_scene(ps): + self.body_view: RigidBodyViewBase = NewtonRigidBodyView( + entities=entities, scene=ps, device=device ) - if self.device.type == "cuda" - else None - ) + else: + self.body_view = DefaultRigidBodyView( + entities=entities, ps=ps, device=device + ) + + # Kept for backward compatibility with callers that index gpu_indices directly. + # NOTE: for Newton, body IDs are lazily resolved after finalization. + # Use the ``gpu_indices`` property instead of caching here. # Initialize rigid body data. self._pose = torch.zeros( @@ -86,77 +99,59 @@ 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 ) self._com_pose = torch.zeros( (self.num_instances, 7), dtype=torch.float32, device=self.device ) + # Physical property buffers + self._mass = torch.zeros( + (self.num_instances, 1), dtype=torch.float32, device=self.device + ) + self._inertia = torch.zeros( + (self.num_instances, 3), dtype=torch.float32, device=self.device + ) + self._friction = torch.zeros( + (self.num_instances, 1), dtype=torch.float32, device=self.device + ) + + @property + def is_newton_backend(self) -> bool: + return isinstance(self.body_view, NewtonRigidBodyView) + + @property + def gpu_indices(self) -> torch.Tensor: + """Body ID tensor (backward-compatible alias for ``body_view.body_ids_tensor``).""" + return self.body_view.body_ids_tensor + + def body_ids_for(self, env_ids: Sequence[int]) -> torch.Tensor: + return self.body_view.select_body_ids(env_ids) @property def pose(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch pose from CPU entities - xyzs = torch.as_tensor( - np.array([entity.get_location() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - quats = torch.as_tensor( - np.array( - [entity.get_rotation_quat() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) - quats = convert_quat(quats, to="wxyz") - self._pose = torch.cat((xyzs, quats), dim=-1) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._pose, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.POSE, - ) - self._pose[:, :4] = convert_quat(self._pose[:, :4], to="wxyz") - self._pose = self._pose[:, [4, 5, 6, 0, 1, 2, 3]] - return self._pose + if self.body_view.can_fetch_pose: + self.body_view.fetch_pose(self._pose) + return self._pose + + logger.log_error(f"RigidBodyData pose requested but body view is not ready.") @property def lin_vel(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch linear velocity from CPU entities - self._lin_vel = torch.as_tensor( - np.array([entity.get_linear_velocity() for entity in self.entities]), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._lin_vel, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.LINEAR_VELOCITY, - ) - return self._lin_vel + if self.body_view.is_ready: + self.body_view.fetch_linear_velocity(self._lin_vel) + return self._lin_vel + + logger.log_error("RigidBodyData lin_vel requested but body view is not ready.") @property def ang_vel(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch angular velocity from CPU entities - self._ang_vel = torch.as_tensor( - np.array( - [entity.get_angular_velocity() for entity in self.entities], - ), - dtype=torch.float32, - device=self.device, - ) - else: - self.ps.gpu_fetch_rigid_body_data( - data=self._ang_vel, - gpu_indices=self.gpu_indices, - data_type=RigidBodyGPUAPIReadType.ANGULAR_VELOCITY, - ) - return self._ang_vel + if self.body_view.is_ready: + self.body_view.fetch_angular_velocity(self._ang_vel) + return self._ang_vel + + logger.log_error("RigidBodyData ang_vel requested but body view is not ready.") @property def vel(self) -> torch.Tensor: @@ -169,39 +164,19 @@ def vel(self) -> torch.Tensor: @property def lin_acc(self) -> torch.Tensor: - if self.device.type == "cpu": - self._lin_acc = torch.as_tensor( - np.array( - [entity.get_linear_acceleration() for entity in self.entities], - ), - 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 + if self.body_view.is_ready: + self.body_view.fetch_linear_acceleration(self._lin_acc) + return self._lin_acc + + logger.log_error("RigidBodyData lin_acc requested but body view is not ready.") @property def ang_acc(self) -> torch.Tensor: - if self.device.type == "cpu": - self._ang_acc = torch.as_tensor( - np.array( - [entity.get_angular_acceleration() for entity in self.entities], - ), - 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 + if self.body_view.is_ready: + self.body_view.fetch_angular_acceleration(self._ang_acc) + return self._ang_acc + + logger.log_error("RigidBodyData ang_acc requested but body view is not ready.") @property def acc(self) -> torch.Tensor: @@ -219,14 +194,7 @@ def com_pose(self) -> torch.Tensor: Returns: torch.Tensor: The center of mass pose with shape (N, 7). """ - 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 - ) + self.body_view.fetch_com_local_pose(self._com_pose) return self._com_pose @@ -249,7 +217,9 @@ def __init__( self.body_type = cfg.body_type self._world = dexsim.default_world() - self._ps = self._world.get_physics_scene() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() @@ -266,6 +236,11 @@ def __init__( if not cfg.use_usd_properties: for entity in entities: entity.set_body_scale(*cfg.body_scale) + if is_newton_scene(self._ps): + # TODO: DexSim Newton consumes the initial physical + # attributes during add_rigidbody(); MeshObject + # set_physical_attr() is still default-backend only. + continue entity.set_physical_attr(cfg.attrs.attr()) else: # Read current properties from USD-loaded entities and write back to cfg @@ -277,18 +252,16 @@ def __init__( first_entity.get_physical_attr().as_dict() ) - super().__init__(cfg, entities, device) + super().__init__(cfg, entities, device, auto_reset=False) # set default collision filter self._set_default_collision_filter() - if device.type == "cuda": - self._world.update(0.001) - self.reset() + self._apply_initial_state() # 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. @@ -332,12 +305,49 @@ def body_data(self) -> RigidBodyData | None: return self._data + def _get_newton_attr(self, env_idx: int): + """Return DexSim Newton metadata physical attributes for an entity.""" + entity = self._entities[env_idx] + entity_handle = int(entity.get_native_handle()) + if entity_handle < 0: + entity_handle &= _UINT64_MAX + + manager = getattr(self._ps, "manager", None) + attr = None + if manager is not None: + attr = ( + getattr(manager, "dexsim_meta", {}).get(entity_handle, {}).get("attr") + ) + if attr is None: + logger.log_error( + f"Newton physical attributes for rigid object '{self.uid}' env {env_idx} are unavailable." + ) + return attr + + def _warn_newton_unsupported(self, api_name: str) -> None: + logger.log_warning( + f"Newton backend does not support RigidObject.{api_name} runtime updates. " + "Skipping this call." + ) + + def _newton_lifecycle_state(self) -> str: + manager = getattr(self._ps, "manager", None) + return getattr(getattr(manager, "lifecycle_state", None), "name", "") + + def _can_use_newton_entity_dynamics_fallback(self) -> bool: + """Return whether per-entity Newton patches are safe before GPU view is ready. + + DexSim Newton only supports MeshObject force/torque helpers in ``BUILDER`` + state. Calling them while the model is ``STALE`` can index stale body ids. + """ + return self._newton_lifecycle_state() == "BUILDER" + @property def body_state(self) -> torch.Tensor: """Get the body state of the rigid object. The body state of a rigid object is represented as a tensor with the following format: - [x, y, z, qw, qx, qy, qz, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z] + [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. @@ -401,6 +411,16 @@ def set_collision_filter( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." ) + if is_newton_scene(self._ps): + if self._data is not None and isinstance( + self._data.body_view, NewtonRigidBodyView + ): + self._data.body_view.apply_collision_filter(filter_data, local_env_ids) + else: + entities = [self._entities[env_idx] for env_idx in local_env_ids] + apply_collision_filter_for_entities(self._ps, entities, filter_data) + return + 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( @@ -423,50 +443,47 @@ def set_local_pose( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." ) - if self.device.type == "cpu" or self.is_static: - pose = pose.cpu() - if pose.dim() == 2 and pose.shape[1] == 7: - pose_matrix = torch.eye(4).unsqueeze(0).repeat(pose.shape[0], 1, 1) - pose_matrix[:, :3, 3] = pose[:, :3] - pose_matrix[:, :3, :3] = matrix_from_quat(pose[:, 3:7]) - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_local_pose(pose_matrix[i]) - elif pose.dim() == 3 and pose.shape[1:] == (4, 4): - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_local_pose(pose[i]) - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - + # Normalize pose to (N, 7) format in (x, y, z, qx, qy, qz, qw). + if pose.dim() == 2 and pose.shape[1] == 7: + target_pose = pose.to(device=self.device, dtype=torch.float32) + elif pose.dim() == 3 and pose.shape[1:] == (4, 4): + xyz = pose[:, :3, 3] + quat = convert_quat(quat_from_matrix(pose[:, :3, :3]), to="xyzw") + target_pose = torch.cat((xyz, quat), dim=-1).to( + device=self.device, dtype=torch.float32 + ) else: - if pose.dim() == 2 and pose.shape[1] == 7: - xyz = pose[:, :3] - quat = convert_quat(pose[:, 3:7], to="xyzw") - elif pose.dim() == 3 and pose.shape[1:] == (4, 4): - xyz = pose[:, :3, 3] - quat = quat_from_matrix(pose[:, :3, :3]) - quat = convert_quat(quat, to="xyzw") - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - # we should keep `pose_` life cycle to the end of the function. - pose = torch.cat((quat, xyz), dim=-1) - indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=pose.clone(), - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.POSE, + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." ) + return + + # Use backend view when pose writes are supported (Newton BUILDER/READY). + if ( + self._data is not None + and self._data.body_view.can_apply_pose + and not self.is_static + ): + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_pose(target_pose, body_ids) + return + + # Static bodies and non-ready backends (notably Newton before finalize) + # still accept direct entity pose updates. + target_pose = target_pose.cpu() + pose_matrix = torch.eye(4).unsqueeze(0).repeat(len(local_env_ids), 1, 1) + pose_matrix[:, :3, 3] = target_pose[:, :3] + pose_matrix[:, :3, :3] = matrix_from_quat( + convert_quat(target_pose[:, 3:7], to="wxyz") + ) + for i, env_idx in enumerate(local_env_ids): + self._entities[env_idx].set_local_pose(pose_matrix[i]) def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: """Get local pose of the rigid object. Args: - to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qw, qx, qy, qz). Defaults to False. + to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qx, qy, qz, qw). Defaults to False. Returns: torch.Tensor: The local pose of the rigid object with shape (N, 7) or (N, 4, 4) depending on `to_matrix`. @@ -485,7 +502,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 @@ -493,10 +509,10 @@ def get_local_pose_cpu( if self.is_static: return get_local_pose_cpu(self._entities, to_matrix).to(self.device) - pose = self.body_data.pose + pose = self.body_data.pose.clone() 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) @@ -551,28 +567,38 @@ 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": - 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()) + if pos is not None: + logger.log_warning( + "RigidObject.add_force_torque(pos=...) is not supported yet; " + "applying wrench at center of mass." + ) - else: - indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) + 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._ps.gpu_apply_rigid_body_data( - data=force, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.FORCE, - ) + self._data.body_view.apply_force(force, body_ids) if torque is not None: - self._ps.gpu_apply_rigid_body_data( - data=torque, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.TORQUE, - ) + self._data.body_view.apply_torque(torque, body_ids) + elif ( + self._data is not None + and self._data.is_newton_backend + and self._can_use_newton_entity_dynamics_fallback() + ): + force_np = force.detach().cpu().numpy() if force is not None else None + torque_np = torque.detach().cpu().numpy() if torque is not None else None + for i, env_idx in enumerate(local_env_ids): + entity = self._entities[env_idx] + if force_np is not None: + entity.add_force(force_np[i]) + if torque_np is not None: + entity.add_torque(torque_np[i]) + elif self._data is not None and self._data.is_newton_backend: + logger.log_warning( + "Cannot apply force or torque while Newton model is stale or " + "unfinalized; call SimulationManager.finalize_newton_physics() first." + ) + else: + logger.log_error("Cannot apply force or torque before body view is ready.") def set_velocity( self, @@ -609,31 +635,32 @@ 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": - for i, env_idx in enumerate(local_env_ids): - if lin_vel is not None: - self._entities[env_idx].set_linear_velocity( - lin_vel[i].cpu().numpy() - ) - if ang_vel is not None: - 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 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._ps.gpu_apply_rigid_body_data( - data=lin_vel, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, - ) + self._data.body_view.apply_linear_velocity(lin_vel, body_ids) if ang_vel is not None: - self._ps.gpu_apply_rigid_body_data( - data=ang_vel, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, - ) + self._data.body_view.apply_angular_velocity(ang_vel, body_ids) + elif ( + self._data is not None + and self._data.is_newton_backend + and self._can_use_newton_entity_dynamics_fallback() + ): + lin_vel_np = lin_vel.detach().cpu().numpy() if lin_vel is not None else None + ang_vel_np = ang_vel.detach().cpu().numpy() if ang_vel is not None else None + for i, env_idx in enumerate(local_env_ids): + entity = self._entities[env_idx] + if lin_vel_np is not None: + entity.set_linear_velocity(lin_vel_np[i]) + if ang_vel_np is not None: + entity.set_angular_velocity(ang_vel_np[i]) + elif self._data is not None and self._data.is_newton_backend: + logger.log_warning( + "Cannot set velocity while Newton model is stale or unfinalized; " + "call SimulationManager.finalize_newton_physics() first." + ) + else: + logger.log_error("Cannot set velocity before body view is ready.") def set_attrs( self, @@ -653,6 +680,10 @@ def set_attrs( f"Length of env_ids {len(local_env_ids)} does not match attrs length {len(attrs)}." ) + if is_newton_scene(self._ps): + self._warn_newton_unsupported("set_attrs") + return + # TODO: maybe need to improve the physical attributes setter efficiency. if isinstance(attrs, RigidBodyAttributesCfg): for i, env_idx in enumerate(local_env_ids): @@ -677,9 +708,17 @@ def set_mass( f"Length of env_ids {len(local_env_ids)} does not match mass length {len(mass)}." ) - mass = mass.cpu().numpy() + if self._data is not None and self._data.body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_mass( + mass.to(dtype=torch.float32, device=self.device).unsqueeze(-1), + body_ids, + ) + return + + mass_np = mass.cpu().numpy() for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].get_physical_body().set_mass(mass[i]) + self._entities[env_idx].get_physical_body().set_mass(mass_np[i]) def get_mass(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """Get mass for the rigid object. @@ -692,9 +731,18 @@ def get_mass(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self._data is not None and self._data.body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + buf = self._data._mass[: len(local_env_ids)] + self._data.body_view.fetch_mass(buf, body_ids) + return buf.squeeze(-1) + masses = [] for _, env_idx in enumerate(local_env_ids): - mass = self._entities[env_idx].get_physical_body().get_mass() + if is_newton_scene(self._ps): + mass = self._get_newton_attr(env_idx).mass + else: + mass = self._entities[env_idx].get_physical_body().get_mass() masses.append(mass) return torch.as_tensor(masses, dtype=torch.float32, device=self.device) @@ -715,12 +763,22 @@ def set_friction( f"Length of env_ids {len(local_env_ids)} does not match friction length {len(friction)}." ) - friction = friction.cpu().numpy() + if self._data is not None and self._data.body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_friction( + friction.to(dtype=torch.float32, device=self.device).unsqueeze(-1), + body_ids, + ) + return + + friction_np = friction.cpu().numpy() for i, env_idx in enumerate(local_env_ids): self._entities[env_idx].get_physical_body().set_dynamic_friction( - friction[i] + friction_np[i] + ) + self._entities[env_idx].get_physical_body().set_static_friction( + friction_np[i] ) - self._entities[env_idx].get_physical_body().set_static_friction(friction[i]) def get_friction(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """Get friction for the rigid object. @@ -733,11 +791,20 @@ def get_friction(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self._data is not None and self._data.body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + buf = self._data._friction[: len(local_env_ids)] + self._data.body_view.fetch_friction(buf, body_ids) + return buf.squeeze(-1) + frictions = [] for _, env_idx in enumerate(local_env_ids): - friction = ( - self._entities[env_idx].get_physical_body().get_dynamic_friction() - ) + if is_newton_scene(self._ps): + friction = self._get_newton_attr(env_idx).dynamic_friction + else: + friction = ( + self._entities[env_idx].get_physical_body().get_dynamic_friction() + ) frictions.append(friction) return torch.as_tensor(frictions, dtype=torch.float32, device=self.device) @@ -758,6 +825,10 @@ def set_damping( f"Length of env_ids {len(local_env_ids)} does not match damping length {len(damping)}." ) + if is_newton_scene(self._ps): + self._warn_newton_unsupported("set_damping") + return + damping = damping.cpu().numpy() for i, env_idx in enumerate(local_env_ids): self._entities[env_idx].get_physical_body().set_linear_damping( @@ -780,12 +851,17 @@ def get_damping(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: dampings = [] for _, env_idx in enumerate(local_env_ids): - linear_damping = ( - self._entities[env_idx].get_physical_body().get_linear_damping() - ) - angular_damping = ( - self._entities[env_idx].get_physical_body().get_angular_damping() - ) + if is_newton_scene(self._ps): + attr = self._get_newton_attr(env_idx) + linear_damping = attr.linear_damping + angular_damping = attr.angular_damping + else: + linear_damping = ( + self._entities[env_idx].get_physical_body().get_linear_damping() + ) + angular_damping = ( + self._entities[env_idx].get_physical_body().get_angular_damping() + ) dampings.append([linear_damping, angular_damping]) return torch.as_tensor(dampings, dtype=torch.float32, device=self.device) @@ -806,10 +882,18 @@ def set_inertia( f"Length of env_ids {len(local_env_ids)} does not match inertia length {len(inertia)}." ) - inertia = inertia.cpu().numpy() + if self._data is not None and self._data.body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_inertia_diagonal( + inertia.to(dtype=torch.float32, device=self.device), + body_ids, + ) + return + + inertia_np = inertia.cpu().numpy() for i, env_idx in enumerate(local_env_ids): self._entities[env_idx].get_physical_body().set_mass_space_inertia_tensor( - inertia[i] + inertia_np[i] ) def get_inertia(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: @@ -823,13 +907,22 @@ def get_inertia(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self._data is not None and self._data.body_view.is_ready: + body_ids = self._data.body_ids_for(local_env_ids) + buf = self._data._inertia[: len(local_env_ids)] + self._data.body_view.fetch_inertia_diagonal(buf, body_ids) + return buf + inertias = [] for _, env_idx in enumerate(local_env_ids): - inertia = ( - self._entities[env_idx] - .get_physical_body() - .get_mass_space_inertia_tensor() - ) + if is_newton_scene(self._ps): + inertia = self._get_newton_attr(env_idx).inertia + else: + inertia = ( + self._entities[env_idx] + .get_physical_body() + .get_mass_space_inertia_tensor() + ) inertias.append(inertia) return torch.as_tensor(inertias, dtype=torch.float32, device=self.device) @@ -942,7 +1035,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). @@ -961,11 +1054,13 @@ def set_com_pose( f"Length of env_ids {len(local_env_ids)} does not match com_pose length {len(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) + if self._data is not None: + target_com_pose = com_pose.to(device=self.device, dtype=torch.float32) + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_com_local_pose(target_com_pose, body_ids) + return + + logger.log_error("Cannot set center of mass pose before body view is ready.") def set_body_type(self, body_type: str) -> None: """Set the body type of the rigid object. @@ -978,6 +1073,10 @@ def set_body_type(self, body_type: str) -> None: """ from dexsim.types import ActorType + if is_newton_scene(self._ps): + self._warn_newton_unsupported("set_body_type") + return + if body_type not in ("dynamic", "kinematic"): logger.log_error( f"Invalid body type {body_type}. Must be one of 'dynamic', or 'kinematic'." @@ -1082,36 +1181,29 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: local_env_ids = self._all_indices if env_ids is None else env_ids - if self.device.type == "cpu": - for env_idx in local_env_ids: - self._entities[env_idx].clear_dynamics() - else: - # Apply zero force and torque to the rigid bodies. + if self._data is not None and self._data.body_view.is_ready: zeros = torch.zeros( (len(local_env_ids), 3), dtype=torch.float32, device=self.device ) - indices = self.body_data.gpu_indices[local_env_ids] - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.FORCE, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.TORQUE, + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_linear_velocity(zeros, body_ids) + self._data.body_view.apply_angular_velocity(zeros, body_ids) + self._data.body_view.apply_force(zeros, body_ids) + self._data.body_view.apply_torque(zeros, body_ids) + elif ( + self._data is not None + and self._data.is_newton_backend + and self._can_use_newton_entity_dynamics_fallback() + ): + for env_idx in local_env_ids: + self._entities[env_idx].clear_dynamics() + elif self._data is not None and self._data.is_newton_backend: + logger.log_warning( + "Cannot clear dynamics while Newton model is stale or unfinalized; " + "call SimulationManager.finalize_newton_physics() first." ) + else: + logger.log_error("Cannot clear dynamics before body view is ready.") def set_physical_visible( self, @@ -1157,12 +1249,9 @@ def set_visible(self, visible: bool = True) -> None: for i, env_idx in enumerate(self._all_indices): self._entities[env_idx].set_visible(visible) - def reset(self, env_ids: Sequence[int] | None = None) -> None: - local_env_ids = self._all_indices if env_ids is None else env_ids - num_instances = len(local_env_ids) - - self.set_attrs(self.cfg.attrs, env_ids=local_env_ids) - + def _build_cfg_init_pose(self, env_ids: Sequence[int]) -> torch.Tensor: + """Build initial root poses from cfg as ``(N, 4, 4)`` matrices.""" + num_instances = len(env_ids) pos = torch.as_tensor( self.cfg.init_pos, dtype=torch.float32, device=self.device ) @@ -1181,14 +1270,47 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: ) pose[:, :3, 3] = pos pose[:, :3, :3] = mat - self.set_local_pose(pose, env_ids=local_env_ids) + return pose + + def _apply_initial_state(self) -> None: + """Apply cfg initial pose after construction. + + PhysX/default backends run a full reset. Newton applies init pose in + ``BUILDER`` via the scene batch API; velocities are cleared after + finalization through :meth:`SimulationManager.finalize_newton_physics`. + """ + if is_newton_scene(self._ps): + if self._newton_lifecycle_state() == "BUILDER": + self.set_local_pose( + self._build_cfg_init_pose(self._all_indices), + env_ids=self._all_indices, + ) + return + + if self.device.type == "cuda": + self._world.update(0.001) + self.reset() + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + local_env_ids = self._all_indices if env_ids is None else env_ids + + # TODO: support attributes setter for newton. + if not is_newton_scene(self._ps): + self.set_attrs(self.cfg.attrs, env_ids=local_env_ids) self.clear_dynamics(env_ids=local_env_ids) + self.set_local_pose( + self._build_cfg_init_pose(local_env_ids), env_ids=local_env_ids + ) + def destroy(self) -> None: env = self._world.get_env() arenas = env.get_all_arenas() 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/soft_object.py b/embodichain/lab/sim/objects/soft_object.py index a06a30f9..1a488fb2 100644 --- a/embodichain/lab/sim/objects/soft_object.py +++ b/embodichain/lab/sim/objects/soft_object.py @@ -150,7 +150,9 @@ def __init__( device: torch.device = torch.device("cpu"), ) -> None: self._world: dexsim.World = dexsim.default_world() - self._ps = self._world.get_physics_scene() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() self._data = SoftBodyData(entities=entities, ps=self._ps, device=device) diff --git a/embodichain/lab/sim/robots/cobotmagic.py b/embodichain/lab/sim/robots/cobotmagic.py index 7bfb9a59..05251d30 100644 --- a/embodichain/lab/sim/robots/cobotmagic.py +++ b/embodichain/lab/sim/robots/cobotmagic.py @@ -194,21 +194,19 @@ def build_pk_serial_chain( torch.set_printoptions(precision=5, sci_mode=False) config = SimulationManagerCfg( - headless=True, - sim_device="cuda", + headless=False, + device="cpu", num_envs=2, render_cfg=RenderCfg(renderer="fast-rt"), ) sim = SimulationManager(config) - config = {"init_pos": [0.0, 0.0, 1.0], "init_qpos": [0.1] * 16} + config = { + "init_pos": [0.0, 0.0, 1.0], + } cfg = CobotMagicCfg.from_dict(config) robot = sim.add_robot(cfg=cfg) - sim.open_window() - - if sim.is_use_gpu_physics: - sim.init_gpu_physics() print("CobotMagic added to the simulation.") diff --git a/embodichain/lab/sim/robots/dexforce_w1/cfg.py b/embodichain/lab/sim/robots/dexforce_w1/cfg.py index 1dd41e93..786e1489 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/cfg.py +++ b/embodichain/lab/sim/robots/dexforce_w1/cfg.py @@ -380,7 +380,7 @@ def build_pk_serial_chain( DexforceW1ArmKind, ) - config = SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=4) + config = SimulationManagerCfg(headless=True, device="cpu", num_envs=4) sim = SimulationManager(config) cfg = DexforceW1Cfg.from_dict( diff --git a/embodichain/lab/sim/sensors/contact_sensor.py b/embodichain/lab/sim/sensors/contact_sensor.py index 49ebbe1c..1cdb82b4 100644 --- a/embodichain/lab/sim/sensors/contact_sensor.py +++ b/embodichain/lab/sim/sensors/contact_sensor.py @@ -211,7 +211,9 @@ def _precompute_filter_ids(self, config: ContactSensorCfg): def _build_sensor_from_config(self, config: ContactSensorCfg, device: torch.device): self._precompute_filter_ids(config) self._world: dexsim.World = dexsim.default_world() - self._ps = self._world.get_physics_scene() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() world_config = dexsim.get_world_config() self.is_use_gpu_physics = device.type == "cuda" and world_config.enable_gpu_sim if self.is_use_gpu_physics: diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index cb734239..1e36cd0b 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 @@ -32,7 +33,7 @@ from copy import deepcopy from datetime import datetime from functools import cached_property -from typing import List, Union, Dict, Union, Sequence +from typing import List, Union, Dict, Sequence from dataclasses import dataclass, asdict, field, MISSING # Global cache directories @@ -47,14 +48,13 @@ PhysicalAttr, ActorType, RigidBodyShape, - RigidBodyGPUAPIReadType, - ArticulationGPUAPIReadType, ) from dexsim.core import TASK_RETURN -from dexsim.engine import CudaArray, Material +from dexsim.engine import Material, PhysicsScene from dexsim.models import MeshObject from dexsim.render import Light as _Light, LightType, Windows from dexsim.engine import GizmoController, ObjectManipulator +from dexsim.engine.newton_physics import NewtonManager, NewtonPhysicsScene from embodichain.lab.sim.objects import ( RigidObject, @@ -75,9 +75,11 @@ ) from embodichain.lab.sim.cfg import ( RenderCfg, - PhysicsCfg, + DefaultPhysicsCfg, + NewtonPhysicsCfg, + physics_backend_from_cfg, + validate_physics_cfg, MarkerCfg, - GPUMemoryCfg, WindowRecordCfg, LightCfg, RigidObjectCfg, @@ -93,6 +95,7 @@ __all__ = [ "SimulationManager", "SimulationManagerCfg", + "get_physics_scene", "SIM_CACHE_DIR", "MATERIAL_CACHE_DIR", "CONVEX_DECOMP_DIR", @@ -104,6 +107,51 @@ class SimulationManagerCfg: """Global robot simulation configuration.""" + def __init__( + self, + width: int = 1920, + height: int = 1080, + headless: bool = False, + render_cfg: RenderCfg | None = None, + gpu_id: int = 0, + thread_mode: ThreadMode = ThreadMode.RENDER_SHARE_ENGINE, + cpu_num: int = 1, + num_envs: int = 1, + arena_space: float = 5.0, + physics_dt: float | None = None, + device: str | torch.device | None = None, + physics_cfg: DefaultPhysicsCfg | NewtonPhysicsCfg | None = None, + window_record: WindowRecordCfg | None = None, + ) -> None: + self.width = width + self.height = height + self.headless = headless + self.render_cfg = RenderCfg() if render_cfg is None else render_cfg + self.gpu_id = gpu_id + self.thread_mode = thread_mode + self.cpu_num = cpu_num + self.num_envs = num_envs + self.arena_space = arena_space + self.physics_cfg = DefaultPhysicsCfg() if physics_cfg is None else physics_cfg + self.window_record = ( + WindowRecordCfg() if window_record is None else window_record + ) + + if physics_dt is not None: + self.physics_cfg.physics_dt = physics_dt + if device is not None: + # Env tensors may use CPU while Newton/Warp sim stays on CUDA for GPU render sync. + if isinstance(self.physics_cfg, NewtonPhysicsCfg): + torch_device = ( + torch.device(device) if isinstance(device, str) else device + ) + if torch_device.type != "cpu": + self.physics_cfg.device = device + else: + self.physics_cfg.device = device + + self.__post_init__() + width: int = 1920 """The width of the simulation window.""" @@ -138,20 +186,35 @@ class SimulationManagerCfg: arena_space: float = 5.0 """The distance between each arena when building multiple arenas.""" - physics_dt: float = 1.0 / 100.0 - """The time step for the physics simulation.""" - - 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.""" - gpu_memory_config: GPUMemoryCfg = field(default_factory=GPUMemoryCfg) - """The GPU memory configuration parameters.""" + physics_cfg: DefaultPhysicsCfg | NewtonPhysicsCfg = field( + default_factory=DefaultPhysicsCfg + ) + """Physics backend configuration (type selects default vs Newton backend).""" window_record: WindowRecordCfg = field(default_factory=WindowRecordCfg) """Viewer window recording settings (hotkey, paths, FPS, memory budget).""" + def __post_init__(self): + validate_physics_cfg(self.physics_cfg) + + @property + def physics_dt(self) -> float: + """The time step for the physics simulation.""" + return self.physics_cfg.physics_dt + + @physics_dt.setter + def physics_dt(self, value: float) -> None: + self.physics_cfg.physics_dt = value + + @property + def device(self) -> str | torch.device: + """The device for the physics simulation.""" + return self.physics_cfg.device + + @device.setter + def device(self, value: str | torch.device) -> None: + self.physics_cfg.device = value + @dataclass class _WindowRecordState: @@ -231,6 +294,10 @@ def __init__( self.sim_config = sim_config self.device = torch.device("cpu") + # Initialize physics backend. + self._physics_backend = physics_backend_from_cfg(sim_config.physics_cfg) + self._newton_manager: NewtonManager = None + world_config = self._convert_sim_config(sim_config) # Initialize warp runtime context before creating the world. @@ -254,14 +321,23 @@ def __init__( self._window_record_input_control: ObjectManipulator | None = None self._window_record_save_threads: list[threading.Thread] = [] - self._world.set_delta_time(sim_config.physics_dt) + self._world.set_delta_time(sim_config.physics_cfg.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: + default_physics_cfg = sim_config.physics_cfg + assert isinstance(default_physics_cfg, DefaultPhysicsCfg) + dexsim.set_physics_config(**default_physics_cfg.to_dexsim_args()) + dexsim.set_physics_gpu_memory_config( + **default_physics_cfg.gpu_memory.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() + self._is_finalized_newton_physics = False # activate physics self.enable_physics(True) @@ -405,9 +481,36 @@ def num_envs(self) -> int: @property def is_use_gpu_physics(self) -> bool: - """Check if the physics simulation is using GPU.""" + """Whether the active physics backend is running on GPU.""" return self.device.type == "cuda" + @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 newton_manager(self) -> NewtonManager: + """Return the DexSim Newton manager for this world, if active.""" + if not self.is_newton_backend: + logger.log_warning("Newton backend is not active.") + 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 is_physics_manually_update(self) -> bool: return self._world.is_physics_manually_update() @@ -446,8 +549,9 @@ 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 + if isinstance(sim_config.physics_cfg, DefaultPhysicsCfg): + world_config.length_tolerance = sim_config.physics_cfg.length_tolerance + world_config.speed_tolerance = sim_config.physics_cfg.speed_tolerance if sim_config.render_cfg.renderer == "auto": from embodichain.lab.sim.utility.render_utils import ( @@ -461,19 +565,16 @@ def _convert_sim_config( sim_config.render_cfg.renderer = resolved_renderer world_config.renderer = sim_config.render_cfg.to_dexsim_flags() - world_config.raytrace_config.render_iterations_per_frame = ( - sim_config.render_cfg.spp - ) + if sim_config.render_cfg.enable_denoiser is False: + world_config.raytrace_config.spp = sim_config.render_cfg.spp + world_config.raytrace_config.open_denoise = False - if type(sim_config.sim_device) is str: - self.device = torch.device(sim_config.sim_device) + if type(sim_config.device) is str: + self.device = torch.device(sim_config.device) else: - self.device = sim_config.sim_device + self.device = sim_config.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." @@ -482,6 +583,18 @@ 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") + + newton_physics_cfg = sim_config.physics_cfg + world_config.newton_cfg = newton_physics_cfg.to_dexsim_cfg( + gpu_id=sim_config.gpu_id, + ) + world_config.gpu_id = sim_config.gpu_id return world_config @@ -492,6 +605,22 @@ def _init_sim_resources(self) -> None: self._default_resources = SimResources() + def _invalidate_newton_physics(self) -> None: + """Mark the Newton scene as needing finalization after scene mutation.""" + if self.is_newton_backend: + self._is_finalized_newton_physics = False + + def _reset_newton_entities_after_finalize(self) -> None: + """Apply deferred initial resets once Newton runtime data is ready.""" + if not self.is_newton_backend: + return + + for rigid_obj in self._rigid_objects.values(): + rigid_obj.reset() + for articulation in self._articulations.values(): + articulation.reset() + # Rigid object groups are not supported on the Newton backend yet. + def enable_physics(self, enable: bool) -> None: """Enable or disable physics simulation. @@ -509,11 +638,23 @@ def set_manual_update(self, enable: bool) -> None: Args: enable (bool): whether to enable manual update. """ + if self.is_newton_backend and enable is False: + logger.log_warning( + "Newton physics backend does not support switching between manual and automatic update. Ignoring set_manual_update call." + ) + return self._world.set_manual_update(enable) def init_gpu_physics(self) -> None: """Initialize the GPU physics simulation.""" - if self.device.type != "cuda": + if self.is_newton_backend: + logger.log_debug( + "GPU physics initialization is handled by the Newton backend. Forcing finalization of Newton physics." + ) + self.finalize_newton_physics() + return + + if not self.is_use_gpu_physics: logger.log_warning( "The simulation device is not cuda, cannot initialize GPU physics." ) @@ -535,6 +676,58 @@ def init_gpu_physics(self) -> None: self._is_initialized_gpu_physics = True + def _newton_lifecycle_state(self) -> str: + """Return the Newton manager lifecycle state name, or empty string.""" + mgr = self.newton_manager + return getattr(getattr(mgr, "lifecycle_state", None), "name", "") + + def finalize_newton_physics(self) -> None: + """Finalize the Newton scene if it has not been finalized yet.""" + if not self.is_newton_backend: + logger.log_warning( + "Newton backend is not active, cannot finalize Newton physics." + ) + return + + if ( + self._is_finalized_newton_physics + and self._newton_lifecycle_state() == "READY" + ): + return + + mgr: NewtonManager = self.newton_manager + state = self._newton_lifecycle_state() + + if state != "READY": + from dexsim.engine.newton_physics.rebuild import ( + ensure_simulation_prepared_lazy, + rebuild_newton_from_scene, + ) + + safe_to_continue, _ = ensure_simulation_prepared_lazy( + mgr, + self._world, + rebuild_from_scene=rebuild_newton_from_scene, + warn=True, + ) + if not safe_to_continue: + logger.log_error( + "Failed to finalize Newton physics: model is not ready to build " + f"(lifecycle state {state!r})." + ) + return + + state = self._newton_lifecycle_state() + if state != "READY": + logger.log_error( + "Failed to finalize Newton physics: lifecycle state is " + f"{state!r} after simulation preparation." + ) + + self._is_finalized_newton_physics = True + self._is_initialized_gpu_physics = self.device.type == "cuda" + self._reset_newton_entities_after_finalize() + def render_camera_group(self, group_ids: list[int]) -> None: """Render all camera group in the simulation. @@ -546,14 +739,16 @@ def render_camera_group(self, group_ids: list[int]) -> None: self._world.render_camera_group(group_ids) - def update(self, physics_dt: float | None = None, step: int = 10) -> None: + def update(self, physics_dt: float | None = None, step: int = 1) -> None: """Update the physics. Args: 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. + step (int, optional): the number of :meth:`World.update` calls per invocation. Defaults to 1. """ - if self.is_use_gpu_physics and not self._is_initialized_gpu_physics: + if self.is_newton_backend: + self.finalize_newton_physics() + elif self.is_use_gpu_physics and not self._is_initialized_gpu_physics: logger.log_warning( f"Using GPU physics, but not initialized yet. Forcing initialization." ) @@ -565,6 +760,8 @@ def update(self, physics_dt: float | None = None, step: int = 10) -> None: for i in range(step): self._world.update(physics_dt) + # TODO: Maybe add newton manager forward kinematics update. + else: logger.log_warning("Physics simulation is not manually updated.") @@ -592,6 +789,15 @@ def get_env(self, arena_index: int = -1) -> dexsim.environment.Arena: def get_world(self) -> dexsim.World: return self._world + def get_physics_scene(self) -> PhysicsScene | NewtonPhysicsScene: + """Get the physics scene of the simulation.""" + if self.is_newton_backend: + physics_scene = self.newton_manager.scene + else: + physics_scene = self._world.get_physics_scene() + + return physics_scene + def open_window(self) -> None: """Open the simulation window.""" self._world.open_window() @@ -675,16 +881,9 @@ def _create_default_plane(self): 0, default_length, repeat_uv_size, repeat_uv_size ) self._default_plane.set_name("default_plane") - plane_collision = self._env.create_cube( - default_length, default_length, default_length / 10 - ) - plane_collision.set_visible(False) - plane_collision_pose = np.eye(4, dtype=float) - plane_collision_pose[2, 3] = -default_length / 20 - 0.001 - plane_collision.set_local_pose(plane_collision_pose) - plane_collision.add_rigidbody(ActorType.KINEMATIC, RigidBodyShape.CONVEX) - - # TODO: add default physics attributes for the plane. + attr = PhysicalAttr(dynamic_friction=0.5, static_friction=0.5) + self._default_plane.add_rigidbody(ActorType.STATIC, RigidBodyShape.PLANE, attr) + self._invalidate_newton_physics() def set_default_background(self) -> None: """Set default background.""" @@ -872,13 +1071,18 @@ def add_rigid_object( cache_dir=self._convex_decomp_dir, ) - rigid_obj = RigidObject(cfg=cfg, entities=obj_list, device=self.device) + rigid_obj = RigidObject( + cfg=cfg, + entities=obj_list, + device=self.device, + ) if cfg.shape.visual_material: mat = self.create_visual_material(cfg.shape.visual_material) rigid_obj.set_visual_material(mat) self._rigid_objects[uid] = rigid_obj + self._invalidate_newton_physics() return rigid_obj @@ -891,6 +1095,13 @@ 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.") @@ -921,6 +1132,13 @@ 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.") @@ -1014,6 +1232,13 @@ def add_rigid_object_group(self, cfg: RigidObjectGroupCfg) -> RigidObjectGroup: Args: cfg (RigidObjectGroupCfg): Configuration for the rigid object group. """ + if self.is_newton_backend: + logger.log_error( + "Rigid object group support for the Newton backend is not enabled " + "in EmbodiChain yet.", + error_type=NotImplementedError, + ) + from embodichain.lab.sim.utility.sim_utils import ( load_mesh_objects_from_cfg, ) @@ -1043,10 +1268,13 @@ def add_rigid_object_group(self, cfg: RigidObjectGroupCfg) -> RigidObjectGroup: # Convert [a1, a2, ...], [b1, b2, ...] to [(a1, b1, ...), (a2, b2, ...), ...] obj_group_list = list(zip(*obj_group_list)) rigid_obj_group = RigidObjectGroup( - cfg=cfg, entities=obj_group_list, device=self.device + cfg=cfg, + entities=obj_group_list, + device=self.device, ) self._rigid_object_groups[uid] = rigid_obj_group + self._invalidate_newton_physics() return rigid_obj_group @@ -1117,7 +1345,6 @@ def add_articulation( Returns: Articulation: The added articulation instance handle. """ - uid = cfg.uid if uid is None: uid = os.path.splitext(os.path.basename(cfg.fpath))[0] @@ -1163,6 +1390,7 @@ def add_articulation( articulation = Articulation(cfg=cfg, entities=obj_list, device=self.device) self._articulations[uid] = articulation + self._invalidate_newton_physics() return articulation @@ -1197,6 +1425,12 @@ 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( + "Robot support for the Newton backend is not enabled " + "in EmbodiChain yet.", + error_type=NotImplementedError, + ) uid = cfg.uid if cfg.fpath is None: @@ -2114,7 +2348,6 @@ def _sever_wrapper_refs(obj_registry): _sever_wrapper_refs("_lights") # Explicitly clear Python references to trigger C++ object destructors - self._ps = None self._env = None self._world = None self._default_plane = None @@ -2156,3 +2389,12 @@ def flush_cleanup_queue(): # At this point, wait for the C++ Scene to return to zero, since the stack is at the top level, there will definitely be no deadlock SimulationManager.wait_scene_destruction() + + +def get_physics_scene(instance_id: int = 0): + """Return the active physics scene from a SimulationManager instance. + + This is the unified EmbodiChain access point for code that previously + reached through ``dexsim.default_world().get_physics_scene()``. + """ + return SimulationManager.get_instance(instance_id).get_physics_scene() diff --git a/embodichain/lab/sim/utility/sim_utils.py b/embodichain/lab/sim/utility/sim_utils.py index 993e4df9..bf1576b0 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 @@ -46,6 +45,19 @@ 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.sim_manager import get_physics_scene + from embodichain.lab.sim.objects.backends import is_newton_scene + + return is_newton_scene(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. @@ -160,22 +172,33 @@ def get_drive_type(drive_pros): else: logger.log_error(f"Unknow drive type {drive_type}") + from embodichain.lab.sim.sim_manager import SimulationManager + + sim = SimulationManager.get_instance() + for i, art in enumerate(arts): - art.set_body_scale(cfg.body_scale) - art.set_physical_attr(cfg.attrs.attr()) + is_newton_art = hasattr(art, "dexsim_meta_links") + lifecycle_state = getattr(getattr(art, "_mgr", None), "_lifecycle_state", None) + lifecycle_name = getattr(lifecycle_state, "name", "") + if not is_newton_art or lifecycle_name == "BUILDER": + art.set_body_scale(cfg.body_scale) + link_names = art.get_link_names() _apply_link_physics_overrides(art, cfg, link_names) art.set_articulation_flag(ArticulationFlag.FIX_BASE, cfg.fix_base) art.set_articulation_flag( ArticulationFlag.DISABLE_SELF_COLLISION, cfg.disable_self_collision ) - art.set_solver_iteration_counts( - min_position_iters=cfg.min_position_iters, - min_velocity_iters=cfg.min_velocity_iters, - ) + if hasattr(art, "set_solver_iteration_counts"): + art.set_solver_iteration_counts( + min_position_iters=cfg.min_position_iters, + min_velocity_iters=cfg.min_velocity_iters, + ) # TODO: We should change this part after improving spawning of articulation. for name in link_names: + if not hasattr(art, "get_physical_body"): + continue physical_body = art.get_physical_body(name) inertia = physical_body.get_mass_space_inertia_tensor() inertia = np.maximum(inertia, 1e-4) @@ -267,6 +290,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() @@ -318,22 +342,29 @@ def load_mesh_objects_from_cfg( max_convex_hull_num=max_convex_hull_num, ) elif cfg.sdf_resolution > 0: + if not is_newton_backend and cfg.body_scale not in [ + (1.0, 1.0, 1.0), + [1.0, 1.0, 1.0], + ]: + logger.log_error( + f"Non-unit body scale {cfg.body_scale} is not supported for SDF collision yet. Please set body_scale to (1.0, 1.0, 1.0) for SDF collision." + ) obj = env.load_actor( fpath, duplicate=True, attach_scene=True, option=option ) - sdf_cfg = SDFConfig() - sdf_cfg.resolution = cfg.sdf_resolution + sdf_cfg = SDFConfig(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.add_rigidbody(body_type, RigidBodyShape.CONVEX, cfg.attrs.attr()) + obj.set_name(f"{cfg.uid}_{i}") obj_list.append(obj) @@ -352,7 +383,11 @@ def load_mesh_objects_from_cfg( obj_list = create_cube(env_list, cfg.shape.size, uid=cfg.uid) for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.BOX) + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.BOX, cfg.attrs.attr()) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) elif isinstance(cfg.shape, SphereCfg): from embodichain.lab.sim.utility.sim_utils import create_sphere @@ -361,7 +396,11 @@ def load_mesh_objects_from_cfg( env_list, cfg.shape.radius, cfg.shape.resolution, uid=cfg.uid ) for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.SPHERE) + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, RigidBodyShape.SPHERE, cfg.attrs.attr()) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) else: logger.log_error( f"Unsupported rigid object shape type: {type(cfg.shape)}. Supported types: MeshCfg, CubeCfg, SphereCfg." diff --git a/examples/sim/demo/grasp_cup_to_caffe.py b/examples/sim/demo/grasp_cup_to_caffe.py index 25a040e4..f26beab8 100644 --- a/examples/sim/demo/grasp_cup_to_caffe.py +++ b/examples/sim/demo/grasp_cup_to_caffe.py @@ -29,6 +29,7 @@ from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, JointDrivePropertiesCfg, RigidObjectCfg, @@ -69,8 +70,9 @@ def initialize_simulation(args) -> SimulationManager: """ config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, num_envs=args.num_envs, arena_space=2.5, @@ -428,9 +430,6 @@ def main(): if not args.headless: sim.open_window() - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - run_simulation(sim, robot, cup, caffe) logger.log_info("\n Press Ctrl+C to exit simulation loop.") diff --git a/examples/sim/demo/pick_up_cloth.py b/examples/sim/demo/pick_up_cloth.py index d6f8e3fa..6be3e3c1 100644 --- a/examples/sim/demo/pick_up_cloth.py +++ b/examples/sim/demo/pick_up_cloth.py @@ -36,6 +36,7 @@ from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, RigidObjectCfg, @@ -252,10 +253,11 @@ def main(): num_envs=args.num_envs, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", + device="cuda", render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals + physics_cfg=physics_cfg_for_backend(args.physics), ) # Create the simulation instance diff --git a/examples/sim/demo/press_softbody.py b/examples/sim/demo/press_softbody.py index f5fada63..7b6d23ac 100644 --- a/examples/sim/demo/press_softbody.py +++ b/examples/sim/demo/press_softbody.py @@ -35,6 +35,7 @@ from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, LightCfg, SoftObjectCfg, @@ -72,8 +73,9 @@ def initialize_simulation(args): """ config = SimulationManagerCfg( headless=True, - sim_device="cuda", + device="cuda", render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, num_envs=args.num_envs, ) diff --git a/examples/sim/demo/scoop_ice.py b/examples/sim/demo/scoop_ice.py index b80e8707..2f03afe8 100644 --- a/examples/sim/demo/scoop_ice.py +++ b/examples/sim/demo/scoop_ice.py @@ -30,6 +30,7 @@ from embodichain.lab.sim.objects import Robot, RigidObject, RigidObjectGroup from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, URDFCfg, @@ -61,6 +62,7 @@ def initialize_simulation(args): config = SimulationManagerCfg( headless=True, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, ) sim = SimulationManager(config) diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index 296c3be4..c8720c29 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -28,7 +28,12 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.sensors import Camera, CameraCfg -from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + RigidObjectCfg, + RigidBodyAttributesCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -49,8 +54,9 @@ def main(): width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) # Create simulation context diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index b0931f24..844fd565 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -23,7 +23,11 @@ import time from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg @@ -46,10 +50,11 @@ def main(): height=1080, headless=args.headless, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, + device=args.device, render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals + physics_cfg=physics_cfg_for_backend(args.physics), ) # Create the simulation instance diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index 40f0d0c1..f13a52aa 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -25,6 +25,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, @@ -50,8 +51,9 @@ def main(): width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) sim = SimulationManager(sim_cfg) diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index a37e6eb8..24be4691 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -31,6 +31,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, @@ -60,8 +61,9 @@ def main(): height=1080, headless=args.headless, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) sim = SimulationManager(sim_cfg) diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py index 09779c84..42a06e3e 100644 --- a/examples/sim/gizmo/gizmo_w1.py +++ b/examples/sim/gizmo/gizmo_w1.py @@ -25,6 +25,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, @@ -51,8 +52,9 @@ def main(): height=1080, headless=args.headless, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) sim = SimulationManager(sim_cfg) diff --git a/examples/sim/planners/motion_generator.py b/examples/sim/planners/motion_generator.py index b3069078..3f118e24 100644 --- a/examples/sim/planners/motion_generator.py +++ b/examples/sim/planners/motion_generator.py @@ -76,7 +76,7 @@ def main(interactive=False): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) + sim = SimulationManager(SimulationManagerCfg(headless=False, device="cpu")) sim.set_manual_update(False) # Robot configuration diff --git a/examples/sim/scene/scene_demo.py b/examples/sim/scene/scene_demo.py index 1c08af6a..5ad4d130 100644 --- a/examples/sim/scene/scene_demo.py +++ b/examples/sim/scene/scene_demo.py @@ -26,6 +26,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RigidBodyAttributesCfg, LightCfg, RobotCfg, @@ -116,8 +117,9 @@ def main(): height=1080, headless=True, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), num_envs=args.num_envs, arena_space=10.0, ) diff --git a/examples/sim/sensors/batch_camera.py b/examples/sim/sensors/batch_camera.py index b6eb4824..00709710 100644 --- a/examples/sim/sensors/batch_camera.py +++ b/examples/sim/sensors/batch_camera.py @@ -19,7 +19,12 @@ import matplotlib.pyplot as plt from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RenderCfg, RigidObjectCfg, LightCfg +from embodichain.lab.sim.cfg import ( + RenderCfg, + physics_cfg_for_backend, + RigidObjectCfg, + LightCfg, +) from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import RigidObject, Light from embodichain.lab.sim.sensors import ( @@ -35,10 +40,11 @@ def main(args): config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, num_envs=args.num_envs, arena_space=2, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) sim = SimulationManager(config) diff --git a/examples/sim/sensors/create_contact_sensor.py b/examples/sim/sensors/create_contact_sensor.py index 17c26caf..4d17235f 100644 --- a/examples/sim/sensors/create_contact_sensor.py +++ b/examples/sim/sensors/create_contact_sensor.py @@ -26,6 +26,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RigidBodyAttributesCfg, ) from embodichain.lab.sim.sensors import ( @@ -189,10 +190,11 @@ def main(): num_envs=args.num_envs, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, + device=args.device, render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals + physics_cfg=physics_cfg_for_backend(args.physics), ) # Create the simulation instance diff --git a/examples/sim/solvers/differential_solver.py b/examples/sim/solvers/differential_solver.py index 11efa65d..6065cde9 100644 --- a/examples/sim/solvers/differential_solver.py +++ b/examples/sim/solvers/differential_solver.py @@ -31,10 +31,10 @@ def main(visualize: bool = True): torch.set_printoptions(precision=5, sci_mode=False) # Set up simulation with specified device (CPU or CUDA) - sim_device = "cpu" + device = "cpu" num_envs = 9 # Number of parallel arenas/environments config = SimulationManagerCfg( - headless=False, sim_device=sim_device, arena_space=1.5, num_envs=num_envs + headless=False, device=device, arena_space=1.5, num_envs=num_envs ) sim = SimulationManager(config) sim.set_manual_update(False) diff --git a/examples/sim/solvers/opw_solver.py b/examples/sim/solvers/opw_solver.py index e8ae222c..89fce41b 100644 --- a/examples/sim/solvers/opw_solver.py +++ b/examples/sim/solvers/opw_solver.py @@ -31,8 +31,8 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" - config = SimulationManagerCfg(headless=False, sim_device=sim_device) + device = "cpu" + config = SimulationManagerCfg(headless=False, device=device) sim = SimulationManager(config) sim.set_manual_update(False) diff --git a/examples/sim/solvers/pink_solver.py b/examples/sim/solvers/pink_solver.py index 6308b612..cd8bfecf 100644 --- a/examples/sim/solvers/pink_solver.py +++ b/examples/sim/solvers/pink_solver.py @@ -31,8 +31,8 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Set up simulation with specified device (CPU or CUDA) - sim_device = "cpu" - config = SimulationManagerCfg(headless=False, sim_device=sim_device) + device = "cpu" + config = SimulationManagerCfg(headless=False, device=device) sim = SimulationManager(config) sim.set_manual_update(False) diff --git a/examples/sim/solvers/pinocchio_solver.py b/examples/sim/solvers/pinocchio_solver.py index 6d70305e..c25ed6b2 100644 --- a/examples/sim/solvers/pinocchio_solver.py +++ b/examples/sim/solvers/pinocchio_solver.py @@ -32,8 +32,8 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" - config = SimulationManagerCfg(headless=False, sim_device=sim_device) + device = "cpu" + config = SimulationManagerCfg(headless=False, device=device) sim = SimulationManager(config) sim.set_manual_update(False) diff --git a/examples/sim/solvers/pytorch_solver.py b/examples/sim/solvers/pytorch_solver.py index 5d954ff6..4217c115 100644 --- a/examples/sim/solvers/pytorch_solver.py +++ b/examples/sim/solvers/pytorch_solver.py @@ -17,10 +17,10 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation environment (CPU or CUDA) - sim_device = "cpu" + device = "cpu" num_envs = 9 # Number of parallel environments config = SimulationManagerCfg( - headless=False, sim_device=sim_device, arena_space=2.0, num_envs=num_envs + headless=False, device=device, arena_space=2.0, num_envs=num_envs ) sim = SimulationManager(config) sim.set_manual_update(False) diff --git a/examples/sim/solvers/srs_solver.py b/examples/sim/solvers/srs_solver.py index 502726de..8aa34bb2 100644 --- a/examples/sim/solvers/srs_solver.py +++ b/examples/sim/solvers/srs_solver.py @@ -31,11 +31,9 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" + device = "cpu" sim = SimulationManager( - SimulationManagerCfg( - headless=False, sim_device=sim_device, width=2200, height=1200 - ) + SimulationManagerCfg(headless=False, device=device, width=2200, height=1200) ) sim.set_manual_update(False) diff --git a/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py index 8d2b5b9c..6ee790dd 100644 --- a/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py +++ b/examples/sim/utility/workspace_analyzer/analyze_cartesian_workspace.py @@ -37,7 +37,7 @@ config = SimulationManagerCfg( headless=False, - sim_device="cuda", + device="cuda", width=1080, height=1080, ) diff --git a/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py index 5c658fa9..fd8a9839 100644 --- a/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py +++ b/examples/sim/utility/workspace_analyzer/analyze_joint_workspace.py @@ -29,7 +29,7 @@ np.set_printoptions(precision=5, suppress=True) torch.set_printoptions(precision=5, sci_mode=False) - config = SimulationManagerCfg(headless=False, sim_device="cpu") + config = SimulationManagerCfg(headless=False, device="cpu") sim_manager = SimulationManager(config) sim_manager.set_manual_update(False) diff --git a/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py b/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py index 8bd1b4ce..47181d15 100644 --- a/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py +++ b/examples/sim/utility/workspace_analyzer/analyze_plane_workspace.py @@ -37,7 +37,7 @@ config = SimulationManagerCfg( headless=False, - sim_device="cpu", + device="cpu", width=1080, height=1080, ) diff --git a/scripts/tutorials/grasp/grasp_generator.py b/scripts/tutorials/grasp/grasp_generator.py index 42fe250f..e2bf43b0 100644 --- a/scripts/tutorials/grasp/grasp_generator.py +++ b/scripts/tutorials/grasp/grasp_generator.py @@ -34,6 +34,7 @@ from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, LightCfg, @@ -77,8 +78,9 @@ def initialize_simulation(args) -> SimulationManager: """ config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, arena_space=2.5, ) diff --git a/scripts/tutorials/gym/modular_env.py b/scripts/tutorials/gym/modular_env.py index 4bfbb5b3..1b7b5146 100644 --- a/scripts/tutorials/gym/modular_env.py +++ b/scripts/tutorials/gym/modular_env.py @@ -34,6 +34,7 @@ from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, ArticulationCfg, RobotCfg, @@ -220,8 +221,9 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs): sim_cfg=SimulationManagerCfg( render_cfg=RenderCfg(renderer=args.renderer), headless=args.headless, - sim_device=args.device, + device=args.device, num_envs=args.num_envs, + physics_cfg=physics_cfg_for_backend(args.physics), ) ) diff --git a/scripts/tutorials/gym/random_reach.py b/scripts/tutorials/gym/random_reach.py index b55a7a8e..61b45f04 100644 --- a/scripts/tutorials/gym/random_reach.py +++ b/scripts/tutorials/gym/random_reach.py @@ -25,6 +25,7 @@ from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, RigidObjectCfg, RigidBodyAttributesCfg, @@ -45,14 +46,16 @@ def __init__( headless=False, device="cpu", renderer="hybrid", + physics_cfg="default", **kwargs, ): env_cfg = EnvCfg( sim_cfg=SimulationManagerCfg( headless=headless, arena_space=2.0, - sim_device=device, + device=device, render_cfg=RenderCfg(renderer=renderer), + physics_cfg=physics_cfg_for_backend(physics_cfg), ), num_envs=num_envs, ) @@ -131,6 +134,7 @@ def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: headless=args.headless, device=args.device, renderer=args.renderer, + physics_cfg=args.physics, ) for episode in range(10): diff --git a/scripts/tutorials/sim/atomic_actions.py b/scripts/tutorials/sim/atomic_actions.py index 2f346ed1..003f818b 100644 --- a/scripts/tutorials/sim/atomic_actions.py +++ b/scripts/tutorials/sim/atomic_actions.py @@ -44,6 +44,7 @@ from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, RenderCfg, + physics_cfg_for_backend, RobotCfg, RigidObjectCfg, RigidBodyAttributesCfg, @@ -99,13 +100,18 @@ def initialize_simulation(args): width=1920, height=1080, headless=True, - sim_device="cuda", + device="cuda", physics_dt=1.0 / 100.0, num_envs=args.num_envs, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) sim = SimulationManager(sim_cfg) + light = sim.add_light( + cfg=LightCfg(uid="main_light", intensity=50.0, init_pos=(0, 0, 2.0)) + ) + return sim diff --git a/scripts/tutorials/sim/create_cloth.py b/scripts/tutorials/sim/create_cloth.py index 1f0d883c..3fe99659 100644 --- a/scripts/tutorials/sim/create_cloth.py +++ b/scripts/tutorials/sim/create_cloth.py @@ -30,6 +30,7 @@ from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RigidObjectCfg, RigidBodyAttributesCfg, ClothObjectCfg, @@ -90,8 +91,9 @@ def main(): headless=True, num_envs=args.num_envs, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", # soft simulation only supports cuda device + device="cuda", # soft simulation only supports cuda device render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) # Create the simulation instance diff --git a/scripts/tutorials/sim/create_rigid_object_group.py b/scripts/tutorials/sim/create_rigid_object_group.py index d681dc91..5d26f6e9 100644 --- a/scripts/tutorials/sim/create_rigid_object_group.py +++ b/scripts/tutorials/sim/create_rigid_object_group.py @@ -23,7 +23,11 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.objects import ( RigidObjectGroup, @@ -48,10 +52,11 @@ def main(): height=1080, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, + device=args.device, render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals + physics_cfg=physics_cfg_for_backend(args.physics), num_envs=args.num_envs, arena_space=3.0, ) diff --git a/scripts/tutorials/sim/create_robot.py b/scripts/tutorials/sim/create_robot.py index 3fe3f9fd..d0e4aa19 100644 --- a/scripts/tutorials/sim/create_robot.py +++ b/scripts/tutorials/sim/create_robot.py @@ -32,6 +32,7 @@ from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, URDFCfg, @@ -54,9 +55,10 @@ def main(): print("Creating simulation...") config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, arena_space=3.0, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, num_envs=args.num_envs, ) diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index b8f6c727..c0331866 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -23,7 +23,11 @@ import time from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -38,6 +42,12 @@ def main(): description="Create a simulation scene with SimulationManager" ) add_env_launcher_args_to_parser(parser) + 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 @@ -46,7 +56,8 @@ def main(): height=1080, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, + device=args.device, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg( renderer=args.renderer, ), @@ -63,8 +74,9 @@ def main(): uid="cube", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="dynamic", + body_scale=[0.5, 0.5, 0.5], attrs=RigidBodyAttributesCfg( - mass=1.0, + mass=0.1, dynamic_friction=0.5, static_friction=0.5, restitution=0.1, @@ -81,11 +93,12 @@ def main(): shape=MeshCfg(fpath=path), body_type="dynamic", attrs=RigidBodyAttributesCfg( - mass=3.0, + mass=10.0, ), body_scale=[0.5, 0.5, 0.5], - init_pos=[0.0, 0.0, 0.2], - init_rot=[90.0, 0.0, 0.0], + init_pos=[0.0, 0.0, 0.5], + init_rot=[0.0, 0.0, 0.0], + max_convex_hull_num=32, ) ) @@ -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/scripts/tutorials/sim/create_sensor.py b/scripts/tutorials/sim/create_sensor.py index 39534d32..343ac854 100644 --- a/scripts/tutorials/sim/create_sensor.py +++ b/scripts/tutorials/sim/create_sensor.py @@ -34,6 +34,7 @@ from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, URDFCfg, @@ -87,9 +88,10 @@ def main(): print("Creating simulation...") config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, arena_space=3.0, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, num_envs=args.num_envs, ) diff --git a/scripts/tutorials/sim/create_softbody.py b/scripts/tutorials/sim/create_softbody.py index 3b8973ef..feaf6369 100644 --- a/scripts/tutorials/sim/create_softbody.py +++ b/scripts/tutorials/sim/create_softbody.py @@ -26,6 +26,7 @@ from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, ) @@ -53,10 +54,11 @@ def main(): headless=True, num_envs=args.num_envs, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", # soft simulation only supports cuda device + device="cuda", # soft simulation only supports cuda device render_cfg=RenderCfg( renderer=args.renderer ), # Enable ray tracing for better visuals + physics_cfg=physics_cfg_for_backend(args.physics), ) # Create the simulation instance diff --git a/scripts/tutorials/sim/export_usd.py b/scripts/tutorials/sim/export_usd.py index c6cb91c7..7e188d5e 100644 --- a/scripts/tutorials/sim/export_usd.py +++ b/scripts/tutorials/sim/export_usd.py @@ -25,6 +25,7 @@ from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, JointDrivePropertiesCfg, RigidObjectCfg, @@ -64,8 +65,9 @@ def initialize_simulation(args) -> SimulationManager: """ config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), physics_dt=1.0 / 100.0, num_envs=1, arena_space=2.5, diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 6d6613f9..aba36d4e 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -26,6 +26,7 @@ from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, @@ -51,8 +52,9 @@ def main(): width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), ) sim = SimulationManager(sim_cfg) diff --git a/scripts/tutorials/sim/import_usd.py b/scripts/tutorials/sim/import_usd.py index ada74edf..d350e2e0 100644 --- a/scripts/tutorials/sim/import_usd.py +++ b/scripts/tutorials/sim/import_usd.py @@ -25,7 +25,11 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + RigidBodyAttributesCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.objects import ( RigidObject, @@ -52,10 +56,11 @@ def main(): height=1080, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, + device=args.device, render_cfg=RenderCfg( renderer=args.renderer, ), # Enable ray tracing for better visuals + physics_cfg=physics_cfg_for_backend(args.physics), num_envs=1, arena_space=3.0, ) diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index e2698d9d..74f53107 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -77,7 +77,7 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim = SimulationManager(SimulationManagerCfg(headless=False, sim_device="cpu")) + sim = SimulationManager(SimulationManagerCfg(headless=False, device="cpu")) sim.set_manual_update(False) # Robot configuration diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 2fb8edda..e0c48861 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -31,11 +31,9 @@ def main(): torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" + device = "cpu" sim = SimulationManager( - SimulationManagerCfg( - headless=False, sim_device=sim_device, width=2200, height=1200 - ) + SimulationManagerCfg(headless=False, device=device, width=2200, height=1200) ) sim.set_manual_update(False) diff --git a/tests/agents/test_shared_rollout.py b/tests/agents/test_shared_rollout.py index 4701540f..c65d1506 100644 --- a/tests/agents/test_shared_rollout.py +++ b/tests/agents/test_shared_rollout.py @@ -186,7 +186,7 @@ def test_embodied_env_writes_next_fields_into_external_rollout(): env_cfg.num_envs = 2 env_cfg.sim_cfg = SimulationManagerCfg( headless=True, - sim_device=torch.device("cpu"), + device=torch.device("cpu"), render_cfg=RenderCfg(renderer="hybrid"), gpu_id=0, ) diff --git a/tests/gym/envs/test_base_env.py b/tests/gym/envs/test_base_env.py index 27767bef..bd353cd6 100644 --- a/tests/gym/envs/test_base_env.py +++ b/tests/gym/envs/test_base_env.py @@ -54,7 +54,7 @@ def __init__( env_cfg = EnvCfg( sim_cfg=SimulationManagerCfg( - headless=headless, arena_space=2.0, sim_device=device + headless=headless, arena_space=2.0, device=device ), num_envs=NUM_ENVS, ) @@ -117,14 +117,14 @@ class BaseEnvTest: """Shared test logic for CPU and CUDA.""" @classmethod - def setup_simulation_hook(cls, sim_device): + def setup_simulation_hook(cls, device): if hasattr(cls, "env"): return cls.env = gym.make( "RandomReach-v1", num_envs=NUM_ENVS, headless=True, - device=sim_device, + device=device, ) cls.device = cls.env.get_wrapper_attr("device") cls.num_envs = cls.env.get_wrapper_attr("num_envs") @@ -217,12 +217,12 @@ def setup_class(cls): import sys -def new_setup_simulation(cls, sim_device): +def new_setup_simulation(cls, device): print(">>> ENTERING setup_simulation", file=sys.stderr) if hasattr(cls, "env"): return cls.env = gym.make( - "RandomReach-v1", num_envs=NUM_ENVS, headless=True, device=sim_device + "RandomReach-v1", num_envs=NUM_ENVS, headless=True, device=device ) cls.device = cls.env.get_wrapper_attr("device") cls.num_envs = cls.env.get_wrapper_attr("num_envs") diff --git a/tests/gym/envs/test_embodied_env.py b/tests/gym/envs/test_embodied_env.py index 9539381e..7bad9b54 100644 --- a/tests/gym/envs/test_embodied_env.py +++ b/tests/gym/envs/test_embodied_env.py @@ -120,14 +120,14 @@ class EmbodiedEnvTest: """Shared test logic for CPU and CUDA.""" - def setup_simulation(self, sim_device): + def setup_simulation(self, device): cfg: EmbodiedEnvCfg = config_to_cfg( METADATA, manager_modules=DEFAULT_MANAGER_MODULES ) cfg.num_envs = NUM_ENVS cfg.sim_cfg = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=device, ) self.env = gym.make(id=METADATA["id"], cfg=cfg) diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index da6b9802..803de5c5 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -19,17 +19,43 @@ import pytest import torch +import argparse from tensordict import TensorDict from embodichain.lab.gym.utils.gym_utils import ( + add_env_launcher_args_to_parser, + init_rollout_buffer_from_config, + merge_args_with_gym_config, config_to_cfg, DEFAULT_MANAGER_MODULES, - init_rollout_buffer_from_config, ) from embodichain.utils.utility import load_config, save_config +def test_env_launcher_args_include_physics(): + """Test that launcher args expose the physics backend config selector.""" + parser = argparse.ArgumentParser() + add_env_launcher_args_to_parser(parser) + + default_args = parser.parse_args([]) + assert default_args.physics == "default" + + newton_args = parser.parse_args(["--physics", "newton"]) + assert newton_args.physics == "newton" + + +def test_merge_args_with_gym_config_includes_physics(): + """Test that CLI physics config overrides the gym config.""" + parser = argparse.ArgumentParser() + add_env_launcher_args_to_parser(parser) + args = parser.parse_args(["--physics", "newton"]) + + merged_config = merge_args_with_gym_config(args, {}) + + assert merged_config["physics"] == "newton" + + class TestInitRolloutBufferFromConfig: """Tests for init_rollout_buffer_from_config function.""" diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 89c37e72..b7ccf17d 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -28,6 +28,7 @@ ArticulationCfg, JointDrivePropertiesCfg, LinkPhysicsOverrideCfg, + physics_cfg_for_backend, RigidBodyAttributesCfg, RigidBodyAttributesOverrideCfg, ) @@ -39,6 +40,12 @@ NUM_ARENAS = 10 +def _teardown_newton_physics() -> None: + from dexsim.engine.newton_physics import teardown_newton_physics + + teardown_newton_physics() + + def _link_static_friction(art: Articulation, link_name: str, env_idx: int = 0) -> float: return art._entities[env_idx].get_physical_attr(link_name).static_friction @@ -77,11 +84,15 @@ def test_resolve_link_physics_overlap_raises(self): class BaseArticulationTest: """Shared test logic for CPU and CUDA.""" - def setup_simulation(self, sim_device): + def setup_simulation(self, device, physics: str = "default"): config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + headless=True, + device=device, + num_envs=NUM_ARENAS, + physics_cfg=physics_cfg_for_backend(physics), ) self.sim = SimulationManager(config) + self.physics = physics art_path = get_data_path(ART_PATH) assert os.path.isfile(art_path) @@ -91,8 +102,10 @@ def setup_simulation(self, sim_device): cfg=ArticulationCfg.from_dict(cfg_dict) ) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): + if device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): self.sim.init_gpu_physics() + if physics == "newton": + self.sim.finalize_newton_physics() def test_local_pose_behavior(self): """Test set_local_pose and get_local_pose: @@ -313,7 +326,7 @@ class BaseArticulationLinkPhysicsTest: """Tests for per-link physics configuration (isolated sim per test).""" def setup_simulation(self, sim_device: str) -> None: - config = SimulationManagerCfg(headless=True, sim_device=sim_device, num_envs=2) + config = SimulationManagerCfg(headless=True, device=sim_device, num_envs=2) self.sim = SimulationManager(config) self.art_path = get_data_path(ART_PATH) assert os.path.isfile(self.art_path) @@ -426,6 +439,69 @@ def setup_method(self): self.setup_simulation("cuda") +class TestArticulationNewton(BaseArticulationTest): + """Articulation coverage on the DexSim Newton physics backend.""" + + def setup_method(self): + self.setup_simulation("cuda", physics="newton") + + def teardown_method(self): + self.sim.destroy() + import embodichain.lab.sim as om + + om.SimulationManager.flush_cleanup_queue() + _teardown_newton_physics() + import gc + + gc.collect() + + def test_control_api(self): + """Newton articulation direct state and control buffers round-trip.""" + qpos_zero = torch.zeros( + (NUM_ARENAS, self.art.dof), dtype=torch.float32, device=self.sim.device + ) + qpos = qpos_zero.clone() + qpos[:, -1] = 0.1 + + self.art.set_qpos(qpos, env_ids=None, target=False) + assert torch.allclose(self.art.body_data.qpos, qpos, atol=1e-5) + + self.art.set_qpos(qpos_zero, env_ids=None, target=False) + self.art.set_qpos(qpos, env_ids=None, target=True) + assert torch.allclose(self.art.body_data.target_qpos, qpos, atol=1e-5) + + qvel = torch.full( + (NUM_ARENAS, self.art.dof), + 0.2, + dtype=torch.float32, + device=self.sim.device, + ) + self.art.set_qvel(qvel, env_ids=None, target=False) + assert torch.allclose(self.art.body_data.qvel, qvel, atol=1e-5) + + qf = torch.ones( + (NUM_ARENAS, self.art.dof), dtype=torch.float32, device=self.sim.device + ) + self.art.set_qf(qf, env_ids=None) + assert torch.allclose(self.art.body_data.qf, qf, atol=1e-5) + + self.art.clear_dynamics() + assert torch.allclose(self.art.body_data.qvel, qpos_zero, atol=1e-5) + assert torch.allclose(self.art.body_data.qf, qpos_zero, atol=1e-5) + + @pytest.mark.skip( + reason="DexSim Newton articulation visual-material helpers are render-Skeleton only." + ) + def test_set_visual_material(self): + super().test_set_visual_material() + + @pytest.mark.skip( + reason="DexSim Newton articulation physical-visible helpers are render-Skeleton only." + ) + def test_set_physical_visible(self): + super().test_set_physical_visible() + + if __name__ == "__main__": test = TestArticulationCPU() test.setup_method() diff --git a/tests/sim/objects/test_cloth_object.py b/tests/sim/objects/test_cloth_object.py index afa182e5..7b3aa313 100644 --- a/tests/sim/objects/test_cloth_object.py +++ b/tests/sim/objects/test_cloth_object.py @@ -67,7 +67,7 @@ def setup_simulation(self): height=1080, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", + device="cuda", num_envs=4, arena_space=3.0, ) diff --git a/tests/sim/objects/test_light.py b/tests/sim/objects/test_light.py index 7e9d58c4..2840567b 100644 --- a/tests/sim/objects/test_light.py +++ b/tests/sim/objects/test_light.py @@ -23,7 +23,7 @@ class TestLight: def setup_method(self): # Setup SimulationManager - config = SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=10) + config = SimulationManagerCfg(headless=True, device="cpu", num_envs=10) self.sim = SimulationManager(config) # Create batch of lights diff --git a/tests/sim/objects/test_rigid_object.py b/tests/sim/objects/test_rigid_object.py index 5beebe26..a756904e 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -13,23 +13,23 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations import os -import torch + import pytest +import torch from embodichain.lab.sim import ( SimulationManager, SimulationManagerCfg, VisualMaterialCfg, ) +from embodichain.data import get_data_path +from embodichain.lab.sim.cfg import RigidObjectCfg, physics_cfg_for_backend +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg 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 DUCK_PATH = "ToyDuck/toy_duck.glb" TABLE_PATH = "ShopTableSimple/shop_table_simple.ply" @@ -38,14 +38,36 @@ Z_TRANSLATION = 2.0 +def _make_test_com_pose(device: torch.device) -> torch.Tensor: + """Create per-env COM poses using EmbodiChain xyzw quaternion convention.""" + return torch.tensor( + [ + [0.04, -0.02, 0.03, 0.0, 0.0, 0.0, 1.0], + [-0.01, 0.05, 0.02, 0.0, 0.0, 0.70710677, 0.70710677], + ], + device=device, + dtype=torch.float32, + ) + + +def _teardown_newton_physics() -> None: + from dexsim.engine.newton_physics import teardown_newton_physics + + teardown_newton_physics() + + class BaseRigidObjectTest: """Shared test logic for CPU and CUDA.""" - def setup_simulation(self, sim_device): + def setup_simulation(self, sim_device: str, physics: str = "default"): config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + headless=True, + device=sim_device, + num_envs=NUM_ARENAS, + physics_cfg=physics_cfg_for_backend(physics), ) self.sim = SimulationManager(config) + self.physics = physics self.sim.enable_physics(False) duck_path = get_data_path(DUCK_PATH) assert os.path.isfile(duck_path) @@ -80,10 +102,16 @@ def setup_simulation(self, sim_device): ), ) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): + if ( + physics == "default" + and sim_device == "cuda" + and getattr(self.sim, "is_use_gpu_physics", False) + ): self.sim.init_gpu_physics() self.sim.enable_physics(True) + if physics == "newton": + self.sim.finalize_newton_physics() def test_is_static(self): """Test the is_static() method of duck, table, and chair objects.""" @@ -158,9 +186,11 @@ 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 != "newton": + assert torch.allclose( + chair_xyz_after, expected_chair_pos, atol=1e-5 + ), f"FAIL: Chair pose changed unexpectedly: {chair_xyz_after.tolist()}" + # Newton: kinematic bodies are not pose-locked yet (DexSim TODO). def test_add_force_torque(self): """Test that add_force applies force correctly to the duck object.""" @@ -404,6 +434,67 @@ 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 == "newton": + expected_mass = torch.ones(NUM_ARENAS, device=self.sim.device) + expected_friction = torch.full( + (NUM_ARENAS,), + self.duck.cfg.attrs.dynamic_friction, + device=self.sim.device, + ) + expected_damping = torch.tensor( + [ + self.duck.cfg.attrs.linear_damping, + self.duck.cfg.attrs.angular_damping, + ], + device=self.sim.device, + ).repeat(NUM_ARENAS, 1) + expected_inertia = self.duck.get_inertia() + assert expected_inertia.shape == (NUM_ARENAS, 3) + assert ( + expected_inertia >= 0 + ).all(), "Initial inertia should be non-negative" + + assert torch.allclose(self.duck.get_mass(), expected_mass) + assert torch.allclose(self.duck.get_friction(), expected_friction) + assert torch.allclose(self.duck.get_damping(), expected_damping) + + # set_attrs and set_body_type remain unsupported on Newton + self.duck.set_attrs(RigidBodyAttributesCfg(mass=2.5)) + self.duck.set_body_type("kinematic") + assert self.duck.body_type == "dynamic" + + # Mass: set and verify round-trip + new_mass = torch.full((NUM_ARENAS,), 2.5, device=self.sim.device) + self.duck.set_mass(new_mass) + assert torch.allclose( + self.duck.get_mass(), new_mass, atol=1e-5 + ), f"Newton set_mass round-trip failed: {self.duck.get_mass()}" + + # Friction: set and verify round-trip + new_friction = torch.full((NUM_ARENAS,), 0.7, device=self.sim.device) + self.duck.set_friction(new_friction) + assert torch.allclose( + self.duck.get_friction(), new_friction, atol=1e-5 + ), f"Newton set_friction round-trip failed: {self.duck.get_friction()}" + + # Inertia: set and verify round-trip + new_inertia = torch.full((NUM_ARENAS, 3), 0.3, device=self.sim.device) + self.duck.set_inertia(new_inertia) + assert torch.allclose( + self.duck.get_inertia(), new_inertia, atol=1e-5 + ), f"Newton set_inertia round-trip failed: {self.duck.get_inertia()}" + + # Damping: still unsupported on Newton + self.duck.set_damping( + torch.full((NUM_ARENAS, 2), 0.2, device=self.sim.device) + ) + + self.table.get_mass() + self.table.get_friction() + self.table.get_damping() + self.table.get_inertia() + return + # 3. body_type assert self.duck.body_type == "dynamic" self.duck.set_body_type("kinematic") @@ -481,29 +572,57 @@ def test_physical_attributes(self): self.duck.get_body_scale(), new_scale ), f"Body scale not set correctly" - # 6. COM pose - com_pose = torch.zeros((NUM_ARENAS, 7), device=self.sim.device) - com_pose[:, 3] = 1.0 # Unit quaternion - com_pose[0, :3] = torch.tensor([0.1, 0.1, 0.1], device=self.sim.device) - - self.duck.set_com_pose(com_pose) - - # Static object should not be able to set COM pose - self.table.set_com_pose(com_pose) # Should log warning but not crash - + def test_set_com_pose(self): + """Test setting full and partial center-of-mass poses.""" assert self.duck.body_data is not None assert self.duck.body_data.default_com_pose is not None assert self.duck.body_data.default_com_pose.shape == ( NUM_ARENAS, 7, - ), f"Default COM pose should have shape (NUM_ARENAS, 7)" + ), "Default COM pose should have shape (NUM_ARENAS, 7)" - com_pose = self.duck.body_data.com_pose - assert isinstance(com_pose, torch.Tensor), "com_pose should be a torch.Tensor" - assert com_pose.shape == ( + com_pose = _make_test_com_pose(self.sim.device) + + self.duck.set_com_pose(com_pose) + + actual_com_pose = self.duck.body_data.com_pose + assert isinstance( + actual_com_pose, torch.Tensor + ), "com_pose should be a torch.Tensor" + assert actual_com_pose.shape == ( NUM_ARENAS, 7, - ), f"COM pose should have shape (NUM_ARENAS, 7), got {com_pose.shape}" + ), f"COM pose should have shape (NUM_ARENAS, 7), got {actual_com_pose.shape}" + assert torch.allclose(actual_com_pose, com_pose, atol=1e-5), ( + "COM pose did not match after full set: " + f"expected {com_pose.tolist()}, got {actual_com_pose.tolist()}" + ) + + partial_com_pose = torch.tensor( + [[0.07, -0.03, 0.04, 0.0, 0.38268343, 0.0, 0.9238795]], + device=self.sim.device, + dtype=torch.float32, + ) + expected_com_pose = com_pose.clone() + expected_com_pose[1] = partial_com_pose[0] + + self.duck.set_com_pose(partial_com_pose, env_ids=[1]) + + actual_com_pose = self.duck.body_data.com_pose + assert torch.allclose(actual_com_pose, expected_com_pose, atol=1e-5), ( + "COM pose did not preserve untouched envs after partial set: " + f"expected {expected_com_pose.tolist()}, got {actual_com_pose.tolist()}" + ) + + assert self.chair.body_data is not None + chair_com_pose_before = self.chair.body_data.com_pose.clone() + self.chair.set_com_pose(com_pose) + assert torch.allclose( + self.chair.body_data.com_pose, chair_com_pose_before, atol=1e-5 + ), "Kinematic rigid object COM pose should not change" + + # Static object should not be able to set COM pose. + self.table.set_com_pose(com_pose) def test_misc_properties(self): """Test miscellaneous properties like collision filter, vertices, and visual materials.""" @@ -578,6 +697,229 @@ def test_misc_properties(self): 1.0, ], f"Material {i} base color incorrect" + def test_geometry_data(self): + """Test mesh-level read APIs: get_triangles and scaled get_vertices. + + Covers: + - ``get_triangles`` — shape ``(N, num_tris, 3)``, int32, partial env_ids. + - ``get_vertices(scale=True)`` — scaled vertices differ from unscaled. + """ + # --- get_triangles (full) --- + triangles = self.duck.get_triangles() + assert isinstance( + triangles, torch.Tensor + ), "get_triangles should return a torch.Tensor" + assert triangles.ndim == 3, "Triangles tensor should be 3-D (N, num_tris, 3)" + assert ( + triangles.shape[0] == NUM_ARENAS + ), f"First dim should be {NUM_ARENAS}, got {triangles.shape[0]}" + assert triangles.shape[2] == 3, "Last dim should be 3 (vertex indices)" + assert ( + triangles.dtype == torch.int32 + ), f"Triangles dtype should be int32, got {triangles.dtype}" + + # --- get_triangles (partial) --- + partial_tris = self.duck.get_triangles(env_ids=[0]) + assert ( + partial_tris.shape[0] == 1 + ), "Partial get_triangles should return 1 instance" + + # --- get_vertices(scale=True) --- + new_scale = torch.full( + (NUM_ARENAS, 3), 2.0, device=self.sim.device, dtype=torch.float32 + ) + self.duck.set_body_scale(new_scale) + + verts_raw = self.duck.get_vertices() + verts_scaled = self.duck.get_vertices(scale=True) + assert torch.allclose( + verts_scaled, verts_raw * 2.0, atol=1e-5 + ), "Scaled vertices should be 2x the raw vertices" + + def test_enable_collision(self): + """Test enable_collision toggle for individual arenas. + + Covers: + - ``enable_collision`` with ``enable=False`` (per-instance mask). + - ``enable_collision`` with ``enable=True`` (restore). + - partial ``env_ids`` subset. + """ + # Disable collision for all arenas and re-enable — no exception should be raised. + disable = torch.zeros(NUM_ARENAS, dtype=torch.bool, device=self.sim.device) + self.duck.enable_collision(disable) + + enable = torch.ones(NUM_ARENAS, dtype=torch.bool, device=self.sim.device) + self.duck.enable_collision(enable) + + # Partial: disable only env 0. + partial_disable = torch.zeros(1, dtype=torch.bool, device=self.sim.device) + self.duck.enable_collision(partial_disable, env_ids=[0]) + + # Restore env 0. + partial_enable = torch.ones(1, dtype=torch.bool, device=self.sim.device) + self.duck.enable_collision(partial_enable, env_ids=[0]) + + def test_reset(self): + """Test reset() restores initial pose and clears dynamics. + + Covers: + - ``reset()`` — all envs returned to ``cfg.init_pos`` (default origin). + - Velocities cleared to zero after reset. + - Partial ``env_ids`` reset: only the specified instance is restored. + """ + # Move duck far from origin and give it velocity. + pose_far = ( + torch.eye(4, device=self.sim.device).unsqueeze(0).repeat(NUM_ARENAS, 1, 1) + ) + pose_far[:, 2, 3] = 5.0 + self.duck.set_local_pose(pose_far) + + lin_vel = ( + torch.tensor([3.0, 0.0, 0.0], device=self.sim.device) + .unsqueeze(0) + .repeat(NUM_ARENAS, 1) + ) + self.duck.set_velocity(lin_vel=lin_vel) + + # Full reset. + self.duck.reset() + + pos_after = self.duck.get_local_pose()[:, :3] + origin = torch.zeros(NUM_ARENAS, 3, device=self.sim.device) + assert torch.allclose( + pos_after, origin, atol=1e-4 + ), f"Duck should be at origin after reset, got {pos_after.tolist()}" + + # Velocities should be zero after reset. + assert self.duck.body_data is not None + lin_vel_after = self.duck.body_data.lin_vel + assert torch.allclose( + lin_vel_after, torch.zeros_like(lin_vel_after), atol=1e-5 + ), f"Linear velocity should be zero after reset, got {lin_vel_after.tolist()}" + + # --- Partial reset: move duck again, reset only env 0 --- + self.duck.set_local_pose(pose_far) + self.duck.reset(env_ids=[0]) + + pos_partial = self.duck.get_local_pose()[:, :3] + assert torch.allclose( + pos_partial[0], origin[0], atol=1e-4 + ), f"Env 0 should be at origin after partial reset, got {pos_partial[0].tolist()}" + # Env 1 was not reset — it should still be displaced. + assert ( + pos_partial[1, 2].item() > 1.0 + ), f"Env 1 should remain displaced after partial reset, got z={pos_partial[1, 2].item()}" + + def test_local_pose_matrix(self): + """Test ``get_local_pose(to_matrix=True)`` returns correct shape and values. + + Covers: + - Shape ``(N, 4, 4)`` output. + - Rotation and translation columns are consistent with the 7-vec form. + - Partial ``env_ids``. + """ + pose_7 = torch.eye(4, device=self.sim.device) + pose_7[0, 3] = 1.0 + pose_7[1, 3] = 2.0 + pose_7[2, 3] = 3.0 + pose_mat_input = pose_7.unsqueeze(0).repeat(NUM_ARENAS, 1, 1) + self.duck.set_local_pose(pose_mat_input) + + # 7-vec form + pose_vec = self.duck.get_local_pose(to_matrix=False) + assert pose_vec.shape == ( + NUM_ARENAS, + 7, + ), f"7-vec pose shape should be ({NUM_ARENAS}, 7), got {pose_vec.shape}" + + # Matrix form + pose_mat = self.duck.get_local_pose(to_matrix=True) + assert pose_mat.shape == ( + NUM_ARENAS, + 4, + 4, + ), f"Matrix pose shape should be ({NUM_ARENAS}, 4, 4), got {pose_mat.shape}" + + # Translation columns must match. + assert torch.allclose( + pose_mat[:, :3, 3], pose_vec[:, :3], atol=1e-5 + ), "Matrix translation column should match 7-vec xyz" + + # Last row must be [0, 0, 0, 1]. + last_row = ( + torch.tensor([0.0, 0.0, 0.0, 1.0], device=self.sim.device) + .unsqueeze(0) + .repeat(NUM_ARENAS, 1) + ) + assert torch.allclose( + pose_mat[:, 3, :], last_row, atol=1e-5 + ), "Last row of pose matrix should be [0, 0, 0, 1]" + + # Rotation matrix must be orthogonal (R @ R.T ≈ I). + R = pose_mat[:, :3, :3] + eye = torch.eye(3, device=self.sim.device).unsqueeze(0).repeat(NUM_ARENAS, 1, 1) + assert torch.allclose( + torch.bmm(R, R.transpose(1, 2)), eye, atol=1e-5 + ), "Rotation sub-matrix should be orthogonal" + + # Partial env_ids. + pose_mat_partial = self.duck.get_local_pose(to_matrix=True) + assert pose_mat_partial.shape[0] == NUM_ARENAS + + def test_body_data_vel_clear(self): + """Test ``body_data.vel``, partial ``clear_dynamics``, and verify dynamics reset. + + Covers: + - ``body_data.vel`` — shape ``(N, 6)`` concatenated lin+ang vel. + - ``clear_dynamics()`` — verifies all velocities become zero (not just called). + - ``clear_dynamics(env_ids=[0])`` — partial clear; only env 0 is zeroed. + """ + assert self.duck.body_data is not None + + lin_vel = ( + torch.tensor([2.0, 0.0, 0.0], device=self.sim.device) + .unsqueeze(0) + .repeat(NUM_ARENAS, 1) + ) + ang_vel = ( + torch.tensor([0.0, 3.0, 0.0], device=self.sim.device) + .unsqueeze(0) + .repeat(NUM_ARENAS, 1) + ) + self.duck.set_velocity(lin_vel=lin_vel, ang_vel=ang_vel) + + # --- body_data.vel --- + vel = self.duck.body_data.vel + assert vel.shape == ( + NUM_ARENAS, + 6, + ), f"vel shape should be ({NUM_ARENAS}, 6), got {vel.shape}" + assert torch.allclose( + vel[:, :3], lin_vel, atol=1e-5 + ), f"First 3 columns of vel should match lin_vel" + assert torch.allclose( + vel[:, 3:], ang_vel, atol=1e-5 + ), f"Last 3 columns of vel should match ang_vel" + + # --- clear_dynamics() full — verify velocities go to zero --- + self.duck.clear_dynamics() + vel_after_clear = self.duck.body_data.vel + assert torch.allclose( + vel_after_clear, torch.zeros_like(vel_after_clear), atol=1e-5 + ), f"Velocities should be zero after clear_dynamics, got {vel_after_clear.tolist()}" + + # --- clear_dynamics(env_ids=[0]) partial --- + # Give env 1 non-zero velocity again. + self.duck.set_velocity(lin_vel=lin_vel, ang_vel=ang_vel) + self.duck.clear_dynamics(env_ids=[0]) + vel_partial = self.duck.body_data.vel + assert torch.allclose( + vel_partial[0], torch.zeros(6, device=self.sim.device), atol=1e-5 + ), f"Env 0 should be zeroed after partial clear_dynamics, got {vel_partial[0].tolist()}" + assert not torch.allclose( + vel_partial[1], torch.zeros(6, device=self.sim.device), atol=1e-5 + ), "Env 1 should still have non-zero velocity after partial clear_dynamics" + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() @@ -600,6 +942,27 @@ def setup_method(self): self.setup_simulation("cuda") +class TestRigidObjectNewton(BaseRigidObjectTest): + """Full rigid-object coverage on the DexSim Newton physics backend.""" + + def setup_method(self): + self.setup_simulation("cuda", physics="newton") + + def teardown_method(self): + super().teardown_method() + _teardown_newton_physics() + + def test_physical_attributes(self): + """Newton getters and setters for mass, friction, inertia work via batch API.""" + super().test_physical_attributes() + + @pytest.mark.skip( + reason="TODO: DexSim Newton SDF rigidbody path is not validated in EmbodiChain yet." + ) + def test_add_sdf_mesh(self): + super().test_add_sdf_mesh() + + if __name__ == "__main__": # pytest.main(["-s", __file__]) test = TestRigidObjectCPU() diff --git a/tests/sim/objects/test_rigid_object_group.py b/tests/sim/objects/test_rigid_object_group.py index 896f5ad3..fd2d6f7f 100644 --- a/tests/sim/objects/test_rigid_object_group.py +++ b/tests/sim/objects/test_rigid_object_group.py @@ -34,10 +34,8 @@ class BaseRigidObjectGroupTest: """Shared test logic for CPU and CUDA.""" - def setup_simulation(self, sim_device): - config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS - ) + def setup_simulation(self, device): + config = SimulationManagerCfg(headless=True, device=device, num_envs=NUM_ARENAS) self.sim = SimulationManager(config) duck_path = get_data_path(DUCK_PATH) @@ -66,7 +64,7 @@ def setup_simulation(self, sim_device): cfg=RigidObjectGroupCfg.from_dict(cfg_dict) ) - if sim_device == "cuda" and self.sim.is_use_gpu_physics: + if device == "cuda" and self.sim.is_use_gpu_physics: self.sim.init_gpu_physics() self.sim.enable_physics(True) diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index 83b1414d..39533490 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -50,11 +50,11 @@ # Base test class for CPU and CUDA class BaseRobotTest: @classmethod - def setup_simulation(cls, sim_device): + def setup_simulation(cls, device): if hasattr(cls, "sim"): return # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device=sim_device, num_envs=10) + config = SimulationManagerCfg(headless=True, device=device, num_envs=10) cls.sim = SimulationManager(config) cfg = DexforceW1Cfg.from_dict( @@ -68,7 +68,7 @@ def setup_simulation(cls, sim_device): cls.robot: Robot = cls.sim.add_robot(cfg=cfg) # Initialize GPU physics if needed - if sim_device == "cuda" and getattr(cls.sim, "is_use_gpu_physics", False): + if device == "cuda" and getattr(cls.sim, "is_use_gpu_physics", False): cls.sim.init_gpu_physics() def test_get_joint_ids(self): diff --git a/tests/sim/objects/test_soft_object.py b/tests/sim/objects/test_soft_object.py index 06b3c1dc..081ab526 100644 --- a/tests/sim/objects/test_soft_object.py +++ b/tests/sim/objects/test_soft_object.py @@ -39,7 +39,7 @@ def setup_simulation(self): height=1080, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device="cuda", + device="cuda", num_envs=4, arena_space=3.0, ) diff --git a/tests/sim/objects/test_usd.py b/tests/sim/objects/test_usd.py index a5558a39..6d307f5c 100644 --- a/tests/sim/objects/test_usd.py +++ b/tests/sim/objects/test_usd.py @@ -38,15 +38,15 @@ class BaseUsdTest: """Shared test logic for CPU and CUDA.""" - def setup_simulation(self, sim_device): + def setup_simulation(self, device): config = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=device, num_envs=NUM_ARENAS, ) self.sim = SimulationManager(config) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): + if device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): self.sim.init_gpu_physics() def test_import_rigid(self): diff --git a/tests/sim/planners/test_motion_generator.py b/tests/sim/planners/test_motion_generator.py index 300d191b..08758e79 100644 --- a/tests/sim/planners/test_motion_generator.py +++ b/tests/sim/planners/test_motion_generator.py @@ -50,7 +50,7 @@ def setup_simulation(self): cls = type(self) if hasattr(cls, "robot_sim"): return - cls.config = SimulationManagerCfg(headless=True, sim_device="cpu") + cls.config = SimulationManagerCfg(headless=True, device="cpu") cls.robot_sim = SimulationManager(cls.config) cls.robot_sim.set_manual_update(False) diff --git a/tests/sim/planners/test_toppra_planner.py b/tests/sim/planners/test_toppra_planner.py index 604581df..31662fc8 100644 --- a/tests/sim/planners/test_toppra_planner.py +++ b/tests/sim/planners/test_toppra_planner.py @@ -25,7 +25,7 @@ def setup_simulation(self): cls = type(self) if hasattr(cls, "sim"): return - cls.sim_config = SimulationManagerCfg(headless=True, sim_device="cpu") + cls.sim_config = SimulationManagerCfg(headless=True, device="cpu") cls.sim = SimulationManager(cls.sim_config) cfg_dict = { diff --git a/tests/sim/sensors/test_camera.py b/tests/sim/sensors/test_camera.py index d95f0c4f..3e9af436 100644 --- a/tests/sim/sensors/test_camera.py +++ b/tests/sim/sensors/test_camera.py @@ -31,11 +31,11 @@ class CameraTest: - def setup_simulation(self, sim_device, renderer="hybrid"): + def setup_simulation(self, device, renderer="hybrid"): # Setup SimulationManager config = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=device, render_cfg=RenderCfg(renderer=renderer), num_envs=NUM_ENVS, ) diff --git a/tests/sim/sensors/test_contact.py b/tests/sim/sensors/test_contact.py index aa38fc22..f53189a7 100644 --- a/tests/sim/sensors/test_contact.py +++ b/tests/sim/sensors/test_contact.py @@ -39,14 +39,14 @@ class ContactTest: - def setup_simulation(self, sim_device, renderer="hybrid"): + def setup_simulation(self, device, renderer="hybrid"): sim_cfg = SimulationManagerCfg( width=1920, height=1080, num_envs=2, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=sim_device, + device=device, render_cfg=RenderCfg(renderer=renderer), ) diff --git a/tests/sim/sensors/test_stereo.py b/tests/sim/sensors/test_stereo.py index 58c5caed..c59b8cb1 100644 --- a/tests/sim/sensors/test_stereo.py +++ b/tests/sim/sensors/test_stereo.py @@ -25,11 +25,11 @@ class StereoCameraTest: - def setup_simulation(self, sim_device, renderer="hybrid"): + def setup_simulation(self, device, renderer="hybrid"): # Setup SimulationManager config = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=device, num_envs=NUM_ENVS, render_cfg=RenderCfg(renderer=renderer), ) diff --git a/tests/sim/solvers/test_differential_solver.py b/tests/sim/solvers/test_differential_solver.py index 0e22a567..1c49c8e9 100644 --- a/tests/sim/solvers/test_differential_solver.py +++ b/tests/sim/solvers/test_differential_solver.py @@ -31,7 +31,7 @@ class BaseSolverTest: def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) # Load robot URDF file diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/solvers/test_opw_solver.py index 7dae255d..8636f354 100644 --- a/tests/sim/solvers/test_opw_solver.py +++ b/tests/sim/solvers/test_opw_solver.py @@ -68,8 +68,8 @@ def grid_sample_qpos_from_limits( class BaseSolverTest: sim = None # Define as a class attribute - def setup_simulation(self, sim_device): - config = SimulationManagerCfg(headless=True, sim_device=sim_device) + def setup_simulation(self, device): + config = SimulationManagerCfg(headless=True, device=device) self.sim = SimulationManager(config) self.sim.set_manual_update(False) diff --git a/tests/sim/solvers/test_pink_solver.py b/tests/sim/solvers/test_pink_solver.py index d5589fde..50c510b9 100644 --- a/tests/sim/solvers/test_pink_solver.py +++ b/tests/sim/solvers/test_pink_solver.py @@ -31,7 +31,7 @@ class BaseSolverTest: def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) self.sim.set_manual_update(False) diff --git a/tests/sim/solvers/test_pinocchio_solver.py b/tests/sim/solvers/test_pinocchio_solver.py index 698cb1f9..bd0236d8 100644 --- a/tests/sim/solvers/test_pinocchio_solver.py +++ b/tests/sim/solvers/test_pinocchio_solver.py @@ -31,7 +31,7 @@ class BaseSolverTest: def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) self.sim.set_manual_update(False) diff --git a/tests/sim/solvers/test_pytorch_solver.py b/tests/sim/solvers/test_pytorch_solver.py index 64bafee8..e3657648 100644 --- a/tests/sim/solvers/test_pytorch_solver.py +++ b/tests/sim/solvers/test_pytorch_solver.py @@ -71,7 +71,7 @@ class BaseSolverTest: def setup_simulation(self, solver_type: str): # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) # Load robot URDF file diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index ada04e84..6b9fdbc2 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -132,7 +132,7 @@ class BaseRobotSolverTest: def setup_simulation(self, solver_type: str, device: str = "cpu"): # Set up simulation with specified device (CPU or CUDA) - config = SimulationManagerCfg(headless=True, sim_device=device) + config = SimulationManagerCfg(headless=True, device=device) self.sim = SimulationManager(config) # Load robot URDF file diff --git a/tests/sim/test_batch_entity.py b/tests/sim/test_batch_entity.py new file mode 100644 index 00000000..78bd7cd0 --- /dev/null +++ b/tests/sim/test_batch_entity.py @@ -0,0 +1,55 @@ +# ---------------------------------------------------------------------------- +# 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 types import SimpleNamespace + +import torch + +from embodichain.lab.sim.common import BatchEntity + + +class _BatchEntityForTest(BatchEntity): + def __init__(self) -> None: + self.reset_calls = 0 + cfg = SimpleNamespace(uid="test_entity") + super().__init__( + cfg=cfg, + entities=[object()], + device=torch.device("cpu"), + ) + + def set_local_pose(self, pose, env_ids=None) -> None: + pass + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + return torch.empty(0) + + def reset(self, env_ids=None) -> None: + self.reset_calls += 1 + + +def test_batch_entity_does_not_reset_in_constructor() -> None: + entity = _BatchEntityForTest() + + assert entity.reset_calls == 0 + + +def test_batch_entity_reset_is_explicit() -> None: + entity = _BatchEntityForTest() + entity.reset() + + assert entity.reset_calls == 1 diff --git a/tests/sim/test_newton_finalize_lifecycle.py b/tests/sim/test_newton_finalize_lifecycle.py new file mode 100644 index 00000000..b43d9c8e --- /dev/null +++ b/tests/sim/test_newton_finalize_lifecycle.py @@ -0,0 +1,92 @@ +# ---------------------------------------------------------------------------- +# 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 types import SimpleNamespace + +from embodichain.lab.sim.sim_manager import SimulationManager + + +class _Resettable: + def __init__(self) -> None: + self.reset_calls = 0 + + def reset(self) -> None: + self.reset_calls += 1 + + +class _NewtonManager: + def __init__(self) -> None: + self.lifecycle_state = SimpleNamespace(name="BUILDER") + self.start_calls = 0 + + def start_simulation(self) -> None: + self.start_calls += 1 + self.lifecycle_state.name = "READY" + + +def _make_newton_sim() -> ( + tuple[SimulationManager, _NewtonManager, _Resettable, _Resettable] +): + sim = object.__new__(SimulationManager) + rigid_obj = _Resettable() + rigid_obj_group = _Resettable() + manager = _NewtonManager() + + sim._physics_backend = "newton" + sim._newton_manager = manager + sim._is_finalized_newton_physics = False + sim._is_initialized_gpu_physics = False + sim._has_reset_newton_entities_after_finalize = False + sim._rigid_objects = {"rigid": rigid_obj} + sim._rigid_object_groups = {"rigid_group": rigid_obj_group} + + return sim, manager, rigid_obj, rigid_obj_group + + +def test_finalize_newton_physics_resets_entities_after_ready() -> None: + sim, manager, rigid_obj, rigid_obj_group = _make_newton_sim() + + sim.finalize_newton_physics() + + assert manager.start_calls == 1 + assert rigid_obj.reset_calls == 1 + assert rigid_obj_group.reset_calls == 0 + assert sim._is_finalized_newton_physics + assert sim._is_initialized_gpu_physics + + +def test_finalize_newton_physics_does_not_repeat_deferred_reset() -> None: + sim, manager, rigid_obj, rigid_obj_group = _make_newton_sim() + + sim.finalize_newton_physics() + sim.finalize_newton_physics() + + assert manager.start_calls == 1 + assert rigid_obj.reset_calls == 1 + assert rigid_obj_group.reset_calls == 0 + + +def test_newton_invalidation_allows_next_finalize_to_reset_again() -> None: + sim, manager, rigid_obj, rigid_obj_group = _make_newton_sim() + + sim.finalize_newton_physics() + sim._invalidate_newton_physics() + sim.finalize_newton_physics() + + assert manager.start_calls == 1 + assert rigid_obj.reset_calls == 2 + assert rigid_obj_group.reset_calls == 0 diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py new file mode 100644 index 00000000..9fb6c616 --- /dev/null +++ b/tests/sim/test_sim_manager_cfg.py @@ -0,0 +1,106 @@ +# ---------------------------------------------------------------------------- +# 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 + +import torch + +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg + + +def test_physics_runtime_fields_are_stored_on_physics_cfg() -> None: + cfg = SimulationManagerCfg( + headless=True, + physics_dt=0.02, + device=torch.device("cpu"), + ) + + assert cfg.physics_dt == 0.02 + assert cfg.device == torch.device("cpu") + assert cfg.physics_cfg.physics_dt == 0.02 + assert cfg.physics_cfg.device == torch.device("cpu") + + serialized = cfg.to_dict() + assert "physics_dt" not in serialized + assert "device" not in serialized + assert serialized["physics_cfg"]["physics_dt"] == 0.02 + assert serialized["physics_cfg"]["device"] == torch.device("cpu") + + +def test_simulation_manager_cfg_keeps_legacy_physics_accessors() -> None: + cfg = SimulationManagerCfg(physics_cfg=NewtonPhysicsCfg()) + + cfg.physics_dt = 0.005 + cfg.device = "cuda:0" + + assert cfg.physics_cfg.physics_dt == 0.005 + assert cfg.physics_cfg.device == "cuda:0" + + +def test_newton_physics_cfg_uses_device() -> None: + cfg = NewtonPhysicsCfg(device="cuda:1") + + serialized = cfg.to_dict() + assert serialized["device"] == "cuda:1" + assert serialized["physics_dt"] == 1.0 / 100.0 + assert "solver_type" not in serialized + + +def test_newton_physics_cfg_uses_mujoco_warp_solver_by_default() -> None: + from dexsim.engine.newton_physics import MJWarpSolverCfg + + cfg = NewtonPhysicsCfg() + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + + assert isinstance(dexsim_cfg.solver_cfg, MJWarpSolverCfg) + assert dexsim_cfg.solver_cfg.solver_type == "mujoco_warp" + + +def test_newton_physics_cfg_converts_mapping_solver_cfg_to_dexsim_cfg() -> None: + from dexsim.engine.newton_physics import MJWarpSolverCfg + + cfg = NewtonPhysicsCfg( + device="cuda", + solver_cfg={ + "class_type": "MJWarpSolverCfg", + "iterations": 12, + "ls_iterations": 4, + "use_mujoco_contacts": False, + }, + ) + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=2) + + assert dexsim_cfg.device == "cuda:2" + assert isinstance(dexsim_cfg.solver_cfg, MJWarpSolverCfg) + assert dexsim_cfg.solver_cfg.iterations == 12 + assert dexsim_cfg.solver_cfg.ls_iterations == 4 + assert dexsim_cfg.solver_cfg.use_mujoco_contacts is False + + +def test_newton_physics_cfg_directly_accepts_dexsim_solver_cfg_object() -> None: + from dexsim.engine.newton_physics import XPBDSolverCfg + + solver_cfg = XPBDSolverCfg(iterations=8, enable_restitution=True) + cfg = NewtonPhysicsCfg(solver_cfg=solver_cfg) + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + + assert isinstance(dexsim_cfg.solver_cfg, XPBDSolverCfg) + assert dexsim_cfg.solver_cfg.iterations == 8 + assert dexsim_cfg.solver_cfg.enable_restitution is True diff --git a/tests/sim/utility/test_workspace_analyze.py b/tests/sim/utility/test_workspace_analyze.py index f6bc95bf..1c952899 100644 --- a/tests/sim/utility/test_workspace_analyze.py +++ b/tests/sim/utility/test_workspace_analyze.py @@ -33,7 +33,7 @@ class BaseWorkspaceAnalyzeTest: sim = None # Define as a class attribute def setup_simulation(self): - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) self.sim.set_manual_update(False)