diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index 885442e7..47a3f607 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -322,3 +322,33 @@ topics: - manager-functor - env-framework status: active + - id: differentiable-env + title: Differentiable Environment (APG) + aliases: + - differentiable env + - apg + - analytic policy gradient + - differentiable rl + - Warp tape autograd + - NewtonStepFunc + - 可微环境 + keywords: + - differentiable + - gradient + - apg + - autograd + - warp tape + - requires_grad + - semi_implicit + - DifferentiableEmbodiedEnv + - NewtonStepFunc + paths: + - topics/differentiable-env/differentiable-env.md + source_of_truth: + - embodichain/lab/gym/envs/differentiable_env.py + - embodichain/lab/sim/diff/ + - embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py + related_topics: + - env-framework + - rl-training + status: active diff --git a/agent_context/topics/differentiable-env/differentiable-env.md b/agent_context/topics/differentiable-env/differentiable-env.md new file mode 100644 index 00000000..89b31cee --- /dev/null +++ b/agent_context/topics/differentiable-env/differentiable-env.md @@ -0,0 +1,110 @@ +# differentiable-env + +> Topic: Differentiable environment for analytic policy gradient (APG) — +> `DifferentiableEmbodiedEnv` + the `embodichain.lab.sim.diff` Warp-tape +> ↔ PyTorch-autograd bridge. + +## Overview + +EmbodiChain supports analytic policy gradient (APG) via +`embodichain.lab.gym.envs.differentiable_env.DifferentiableEmbodiedEnv`. +The bridge wraps a Warp tape around one EmbodiChain physics step and +exposes a `torch.autograd.Function` +(`embodichain.lab.sim.diff.NewtonStepFunc`) so PyTorch-side `action` +tensors get a gradient from `tape.backward()`. + +## Required configuration + +- `NewtonPhysicsCfg(requires_grad=True, solver_cfg={"solver_type": "semi_implicit"})` +- `use_cuda_graph=False` (forced by dexsim when grad mode is on) + +The default backend and any other Newton solver are rejected at +construction time by `DifferentiableEmbodiedEnv._validate_diff_cfg`. + +## Subclass contract + +Task authors implement two methods on `DifferentiableEmbodiedEnv`: + +- `_apply_action_kernel(action_wp, tape)` — launch a Warp kernel that + writes joint/body targets into `nm._control` while the tape is open. + The `action_wp` argument is a `wp.array(dtype=wp.float32, + requires_grad=True)` of shape `[num_envs * action_dim]`. +- `_read_outputs(final_state)` — build the `obs` / `reward` / + `terminated` / `truncated` outputs as torch tensors via `wp.to_torch` + so the tape can record the dependency. Must return a dict with + `_order` (tuple of output keys) and `_grad_track` (mapping from output + key to the Warp array that backs its gradient, or `None` for outputs + that don't need grad). + +Optionally override `_make_step_fn()` to swap the per-substep advance +function. The default uses `dexsim.engine.newton_physics.DifferentiableStepper.step`; +the Franka APG example overrides it to call `newton.eval_fk` directly +(see "FK bypass" below). + +See `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` for +the canonical example. + +## Why reward must be computed inside the tape + +`NewtonStepFunc.forward` keeps the `wp.Tape` open while +`obs_reward_fn(final_state)` runs. Reward must be computed by a Warp +kernel that writes into a `wp.zeros(..., requires_grad=True)` array +inside the tape; `wp.to_torch(reward_wp)` then returns a torch tensor +that carries the tape's gradient. Computing reward in pure torch *after* +the tape closes would detach it from the autograd graph and +`action.grad` would come back as `None`. + +The same rule applies to any observation that needs to be +grad-tracked: build it from `wp.to_torch` of a tape-tracked Warp array. + +## FK bypass for the Franka task + +The `semi_implicit` Newton solver does not propagate gradient through +`joint_target_pos` to `body_q` (verified empirically; the reference +implementation at `/root/sources/analytic_policy_gradients/envs/franka_reach_env.py` +hits the same limitation and uses the same workaround). The Franka APG +example overrides `_make_step_fn()` to call `newton.eval_fk(model, +new_joint_q, joint_qd, fk_state)` directly, bypassing the dynamics +solver. The grad path is then: + + action → new_joint_q (action kernel) → eval_fk → body_q → reward kernel → reward_wp → tape.backward → action.grad + +The default `_make_step_fn` still uses the differentiable stepper, so +envs whose reward depends on dynamics (not just FK) can use it — but +they should verify the solver actually propagates grad for their +control inputs before relying on it. + +## Functor autograd compatibility + +Reward/observation functors that compose torch operations on tensors +obtained via `wp.to_torch` are automatically autograd-compatible. +Functors that detour through CPU / NumPy break the graph; those need +torch-only reimplementations for the differentiable path. For now, the +differentiable env computes reward via a dedicated Warp kernel rather +than reusing the standard reward-manager functors — a future task can +audit and port functors as needed. + +## Memory + +Each step records `sim_steps_per_control` substeps into the tape. For +long horizons or large `num_envs`, pass `truncate_backward_at=K` on the +env config to split the tape and detach at chunk boundaries. + +## Source of truth + +- `embodichain/lab/gym/envs/differentiable_env.py` — + `DifferentiableEmbodiedEnv` base class. +- `embodichain/lab/sim/diff/bridge.py` — `NewtonStepFunc`, + `tape_context`, `differentiable_step`. +- `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` — + example task. +- `embodichain/lab/sim/sim_manager.py` — + `SimulationManager.create_differentiable_stepper` / + `create_gradient_rollout` delegators. +- `/root/sources/dexsim/python/dexsim/engine/newton_physics/differentiable_stepper.py` + — the underlying dexsim primitive. + +## Related topics + +- env-framework +- rl-training diff --git a/design/newton-backend-design.md b/design/newton-backend-design.md new file mode 100644 index 00000000..55f51b22 --- /dev/null +++ b/design/newton-backend-design.md @@ -0,0 +1,457 @@ +# 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 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_cfg` (mapping or +`NewtonSolverCfg` selecting `mujoco_warp` / `xpbd` / `semi_implicit` / +`featherstone` / `vbd`), `broad_phase`, and `visualizer_enabled`. +`NewtonPhysicsCfg.to_dexsim_cfg(...)` builds a DexSim `NewtonCfg`, disables +CUDA graph when gradient mode is enabled, and requires +`solver_type="semi_implicit"` for gradient mode. + +### PhysicsBackend abstraction + +`SimulationManager` delegates backend-specific behavior to a +`PhysicsBackend` instance held as `self.physics` (selected by `physics_cfg` +type via `physics_backend_from_cfg`). The backend package lives at +`embodichain/lab/sim/physics/`: + +```text +embodichain/lab/sim/physics/ + __init__.py # registry + make_physics_backend(physics_cfg, manager) + base.py # PhysicsBackend ABC + default.py # DefaultPhysicsBackend (name = "default") + newton.py # NewtonPhysicsBackend (name = "newton") +``` + +`PhysicsBackend` is constructed with a back-reference to its owning +`SimulationManager` (an instance member, not a class singleton — this preserves +EmbodiChain's multiton, which IsaacLab's class-singleton approach would break). +The manager delegates through `self.physics.*` instead of branching on a backend +name: + +- `configure_world(world_config, sim_config)` applies backend-specific + `WorldConfig` fields (default tolerances/GPU flags, or `world_config.newton_cfg`). +- `activate(sim_config)` runs post-world-creation setup (default + `set_physics_config` / GPU-memory config, or `get_newton_manager(self._world)`). +- `prepare()` is the unified "force the backend ready-to-step" entry point. + `SimulationManager.init_gpu_physics()` and `finalize_newton_physics()` both + delegate to it — Newton's "GPU init" is a finalize; the default's "finalize" + is a GPU init. Idempotent; after `invalidate()` it re-prepares (rebuilds). +- `ensure_initialized()` is the lazy `update()`-time wrapper (default: lazy GPU + init; Newton: finalize/rebuild if invalidated). +- `invalidate()` marks the scene dirty after mutation (no-op for default). +- `get_scene()` returns the active physics scene. +- `newton_manager` returns the Newton manager or `None`. + +Capability predicates drive the `add_*` guards (see Parity Matrix below): +`supports_robot`, `supports_soft_bodies`, `supports_cloth`, +`supports_rigid_object_group`, `can_disable_manual_update`. + +Public `SimulationManager` accessors are preserved as thin delegators for +back-compat: `physics_backend`, `is_default_backend`, `is_newton_backend`, +`newton_manager`, `init_gpu_physics()`, `finalize_newton_physics()`, +`get_physics_scene()`. + +Scene mutation invalidates Newton finalization via `_invalidate_newton_physics()` +(delegates to `self.physics.invalidate()`). After finalization, +`_reset_entities_after_finalize()` resets rigid objects, articulations, and +robots so deferred initial state is applied once Newton runtime data is ready. +Rigid object groups are not yet supported on Newton. + +### Object Backend Adapters + +Rigid-body and articulation data access is routed through: + +```text +embodichain/lab/sim/objects/backends/ + base.py # RigidBodyViewBase, ArticulationViewBase (ABCs) + default.py # DefaultRigidBodyView, DefaultArticulationView (PhysX/DexSim-GPU) + newton.py # NewtonRigidBodyView, NewtonArticulationView (Warp) +``` + +`*Data` selects the view at construction via `is_newton_scene(ps)` (a duck-type +check). The views implement lazy body-id resolution and a BUILDER-state +entity-level fallback before the Newton model is finalized. + +EmbodiChain public rigid-body tensor convention is `(x, y, z, qx, qy, qz, qw)`; +the default adapter converts to/from DexSim's `(qx,qy,qz,qw,x,y,z)`, Newton +needs no conversion. + +Newton rigid-object support includes dynamic/kinematic/static creation, local +pose, body state, linear/angular velocity+acceleration, force/torque at COM, +clear dynamics, reset, COM local pose, mass/friction/inertia-diagonal/ +restitution/contact-offset get+set, collision filter (dynamic/kinematic/static/ +pre-finalize), and visual material/visibility/geometry/scale/user-id APIs. +`apply_contact_offset`/`fetch_contact_offset` were added to +`RigidBodyViewBase` and the Newton view. + +Static Newton bodies do not have `RigidBodyData`; static collision-filter writes +use DexSim's per-entity metadata hook when a Newton body ID is not available. + +### Newton-native physics attributes (Phase 3) + +`RigidBodyAttributesCfg` previously flattened to the legacy PhysX-oriented +`PhysicalAttr` via `.attr()`, so on Newton: Newton-native contact/shape params +(`ke`/`kd`/`margin`/`gap`/`mu_torsional`/...) were not representable, PhysX-only +fields were silently ignored, and `density`/`enable_collision` were dropped. +This is now fixed by adopting dexsim's spawn-descriptor pattern at the EmbodiChain +config layer. + +- `cfg.py`: `NewtonCollisionAttributesCfg` (20 fields mirroring + `dexsim.spawn.descs.NewtonCollisionDesc`, all `Optional`, `None` = keep backend + default) + a `newton` sub-config on `RigidBodyAttributesCfg` and + `RigidBodyAttributesOverrideCfg`. `from_dict` parses nested `"newton"`. + `RigidBodyAttributesOverrideCfg.merged_cfg(base)` returns a merged + `RigidBodyAttributesCfg` preserving the `newton` sub-config (override non-None + wins, else base); `merge_with()` keeps its legacy `PhysicalAttr` return for the + default path. `.attr()` is unchanged (no default-backend regression). +- `physics_attrs.py` (new resolver): `ResolvedNewtonShape(NewtonCollisionDesc)` + + `resolve_newton_shape` (projects common `dynamic_friction→mu`, + `restitution`, `enable_collision→has_shape_collision`, `density` — positive so + dexsim computes a positive body mass) + `resolve_newton_body` + (`RigidBodyPhysicsDesc.dynamic/static/kinematic`) + + `resolve_rigid_body_attributes` (dispatch by backend). Re-exports dexsim's + `NEWTON_CONTACT_SOLVER_FIELDS` / `NEWTON_CONTACT_FIELDS` and ports + `warn_ignored_contact_fields` (per-solver) + `warn_backend_mismatched_fields` + (PhysX-only fields on Newton). +- RigidObject spawn (`sim_utils.py`): **opt-in desc-native path** — when + `is_newton and cfg.attrs.newton is not None`, route box/sphere/CONVEX-mesh + through `register_mesh_object_to_newton_patch(newton_shape=, newton_body=)` + (populating the `mgr.dexsim_meta` scaffolding registration/rebuild read), + bypassing legacy `PhysicalAttr` so Newton-native contact/shape params reach the + model. SDF and CoACD keep the legacy path this phase. When `attrs.newton` is + `None`, the legacy `add_rigidbody(attr=)` path is unchanged. +- Articulation: common fields apply via the legacy `set_physical_attr` path on + BUILDER skeletons; `set_dexsim_articulation_cfg` warns when Newton-native + per-link fields are set (dexsim's `NewtonArticulation` has no per-link + contact-material API — see Deferred). + +### Runtime attribute mutation on Newton + +`RigidObject.set_attrs`/`set_damping`/`set_body_type` are no longer warn-and-skip: + +- `set_attrs`: when finalized, applies the Newton-supported subset (mass, + dynamic_friction, restitution, contact_offset) via the batch view and mirrors + all fields to the attr meta; before finalization, mirrors only. +- `set_damping`: documented runtime no-op that mirrors to meta (Newton does not + model per-body damping) so `get_damping`/rebuild stay consistent. +- `set_body_type`: no-op with a clearer message — body type is fixed at + registration on Newton and cannot change at runtime without a rebuild. + +`set_mass`/`set_friction`/`set_inertia` use the batch view when finalized; their +not-ready `else` paths mirror the single field to meta on Newton (the PhysX-bound +`get_physical_body().set_*` are not Newton-patched). `Articulation.set_link_physical_attr` +pushes per-link **mass** live on Newton via `set_link_mass` (mirroring the +dedicated `set_mass`); friction/restitution/contact_offset remain rebuild-time- +only for articulation links. + +### add_robot / add_articulation on Newton + +Robots are URDF articulations; the Newton `load_urdf` patch builds a +`NewtonArticulation`. `add_robot` and `add_articulation` are now **supported** on +Newton (`supports_robot = True`). This required an upstream dexsim fix +(`NewtonArticulation._joint_metas_from_ids`): explicit `joint_ids` were +raw-dict-indexed (including fixed joints) instead of active-joint-indexed, +conflicting with `get_dof()`/`get_actived_joint_names()` and breaking +mimic-jointed robots (dexforce_w1) at spawn. The fix indexes into active joints; +the `joint_ids=None` path is unchanged so existing callers are unaffected. The +dexsim fix lives on dexsim branch `yueci/adapt-embodichain` (commit `d0e86bb02`) +— `add_robot`-on-Newton depends on it being present. + +### Backend capability parity matrix + +`tests/sim/test_backend_parity.py` is the single source of truth for which +features each backend supports (`BACKEND_CAPABILITIES` table). It pins that each +backend's `supports_*`/`can_disable_manual_update` flags match the table, every +`add_robot/add_soft_object/add_cloth_object/add_rigid_object_group` guard raises +`NotImplementedError` iff its flag is False, and the matrix covers every flag and +backend. Current matrix: + +| feature | default | newton | +|--------------------------|---------|--------| +| robot | yes | yes | +| soft_bodies | yes | no | +| cloth | yes | no | +| rigid_object_group | yes | no | +| can_disable_manual_update| yes | no | + +### Currently Unsupported Newton APIs + +`SimulationManager` explicitly rejects these asset types on Newton (per the +parity matrix): + +- `add_soft_object(...)` +- `add_cloth_object(...)` +- `add_rigid_object_group(...)` + +`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. CoACD-decomposed meshes keep the legacy attr +path on Newton (no `attrs.newton` desc-native routing yet). + +Articulation Newton-native **per-link** contact/shape params (`ke`/`kd`/`margin`/ +...) are accepted in config but not applied (dexsim `NewtonArticulation` exposes +no per-link contact-material setter); a warning fires at spawn. Common fields are +applied. + +### Verified Tests + +Newton integration is covered across headless and GPU suites: + +```bash +pytest -q tests/sim/objects/test_rigid_object.py +pytest -q tests/sim/objects/test_articulation.py::TestArticulationNewton +pytest -q tests/sim/objects/test_robot.py::TestRobotNewton +pytest -q tests/sim/test_physics_attrs.py tests/sim/test_backend_parity.py +pytest -q tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_sim_manager_cfg.py +``` + +Recently observed results: Newton rigid (physical_attributes + desc-native +spawn), Newton articulation (incl. per-link mass-live), `TestRobotNewton` +(spawn/finalize/control smoke), 14 headless `physics_attrs` tests, 22 headless +`backend_parity` tests, 5 Newton lifecycle tests, 6 cfg tests — all green. The +default-backend rigid suite (CPU+CUDA) passes with no regression. + +## Improvements To Make + +### API Clarity + +- The `is_newton_scene` sweep is largely complete: backend selection is via the + `PhysicsBackend` ABC and the `add_*` capability guards; the remaining + `is_newton_scene` branches in `rigid_object.py`/`articulation.py` are + legitimate lifecycle fallbacks (BUILDER-state entity dynamics, not-ready meta + reads, static-object paths) that don't map to the batch-oriented view ABC + without extending its semantics. +- `is_use_gpu_physics` still conflates selected tensor/device location, + default-backend GPU API availability, and Newton GPU execution; consider + splitting when a consumer needs to distinguish them. + +### Newton Lifecycle + +- `finalize_newton_physics()` (`self.physics.prepare()`) is the single Newton + preparation 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 force-at-position when DexSim Newton exposes the needed API. +- Validate SDF rigid mesh creation and collision behavior on Newton; route SDF + and CoACD through the desc-native path when `attrs.newton` is set. +- Fix or document kinematic pose-lock semantics. + +### Object Groups, Soft, Cloth + +- Add Newton rigid-object-group support after a design decision (dexsim has no + first-class group API). +- Keep soft and cloth fail-fast until there is an explicit Newton design and + test coverage. dexsim exposes `SoftBodyObject`/`add_softbody`/`add_clothbody` + (requires the VBD solver) — feasible but substantial. + +### Articulation / Robot + +- Apply Newton-native per-link contact/shape params once dexsim exposes a + `NewtonArticulation` per-link shape-material setter. +- Add runtime `Articulation.set_link_physical_attr` Newton live push for + friction/restitution/contact_offset once a live per-link API exists (mass is + already live). + +### 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. + +## Completion Plan + +Done: + +1. Single-rigid-object Newton API stabilized; `test_rigid_object.py` green. +2. Backend capability declarations (`PhysicsBackend.supports_*`) drive `add_*` + guards, pinned by `test_backend_parity.py`. +3. Newton `RigidObject` parity for attributes, damping, body type — implemented + (`set_attrs` live subset + meta-mirror, `set_damping` no-op+meta, + `set_body_type` documented no-op). +4. Tests for Newton lifecycle rebuild and runtime property mutation after + finalization — present (`test_newton_finalize_lifecycle.py`, + `test_rigid_object.py::TestRigidObjectNewton`). +6. Gym env init/reset uses `init_gpu_physics()` / `finalize_newton_physics()` + (already wired via the `base_env.py` pattern). +9. Articulation and robot support on Newton — implemented (incl. upstream + dexsim joint-active-indexing fix); `TestArticulationNewton` and + `TestRobotNewton` green. +13. Multi-env parallel simulation on Newton — already complete via the + spawn-time prototype+clone path (`spawn_rigid_object_entities` / + `spawn_articulation_entities` → dexsim's `clone_actor_to`, + Newton-patched). Newton object views accept multi-entity lists and + resolve one body ID per env. Covered by `TestRigidObjectNewton` + (`NUM_ARENAS=2`, `test_spawn_clones_distinct_entities`), + `TestArticulationNewton` (`num_envs=2`), `TestRobotNewton` + (`num_envs=10`). Implementation plan: + `docs/superpowers/plans/2026-06-22-newton-backend-pr.md`. +14. Differentiable env for APG — implemented. + `embodichain.lab.sim.diff` provides `NewtonStepFunc` + (`torch.autograd.Function`) bridging a `wp.Tape` around + `DifferentiableStepper` into PyTorch autograd, plus `tape_context` + and `differentiable_step` helpers. `SimulationManager` gains + `create_differentiable_stepper` / `create_gradient_rollout` + delegators. `DifferentiableEmbodiedEnv` validates + `NewtonPhysicsCfg(requires_grad=True, solver_type="semi_implicit")` + and overrides `step()` to call `NewtonStepFunc.apply`. The Franka + FR3 reach APG example (`franka_reach_apg.py`) exercises the bridge + end-to-end with a Warp action kernel and a Warp reward kernel + computed inside the tape; `test_franka_apg_smoke_backward` and + `test_franka_apg_one_iter_loss_reduces` are green. Agent context: + `agent_context/topics/differentiable-env/`. + + .. note:: + The Franka task uses an FK-bypass step function + (``newton.eval_fk``) because the ``semi_implicit`` solver does + not propagate gradient through ``joint_target_pos`` to + ``body_q``. The default ``_make_step_fn`` still uses the + differentiable stepper for envs that want the dynamics-grad + path; see the differentiable-env topic for details. + +Remaining: + +5. Implement and test Newton `RigidObjectGroup` (after a design decision). +7. Add rigid-only Newton gym smoke tests. +10. Add soft/cloth support after a dedicated Newton object design and tests. +11. Newton-native per-link contact params for articulations (after dexsim + exposes a per-link shape-material setter). +12. Full migration off legacy `PhysicalAttr` to dexsim's spawn descriptors + (Phase 3 follow-up `3b`) — defer until a third backend appears or dexsim's + attr-path deletion lands. + +## 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. + +PhysicsBackend abstraction: + +- `PhysicsBackend` ABC contract enforced (abstract methods; concrete backends + implement them). `test_backend_parity.py` pins the capability matrix and the + `add_*` guard mapping. +- The Newton finalize/invalidate lifecycle is owned by `NewtonPhysicsBackend` + (`test_newton_finalize_lifecycle.py` — headless, patches the rebuild entry + point). + +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. + +Newton-native attributes (`test_physics_attrs.py`, headless): + +- `from_dict` parses nested `newton`; `resolve_newton_shape` projects common + fields (`friction→mu`, `restitution`, `enable_collision→has_shape_collision`, + `density`); `merged_cfg` propagates `newton`; per-solver warnings + (`xpbd` ignores `ke`/`kd`; `mujoco_warp` ignores `restitution`) and + backend-mismatch warnings fire correctly. + +Rigid object: + +- Dynamic/static/kinematic rigid bodies under Newton. +- Pose, velocity, acceleration, force/torque, reset, COM pose, mass, friction, + inertia, restitution, contact offset, collision filters, geometry APIs behave + consistently with the documented support matrix. +- `attrs.newton` set spawns via the desc-native path; body registers with the + Newton manager after finalize; common fields round-trip via the batch view. +- `set_attrs`/`set_damping`/`set_body_type` produce the documented behavior + (live subset / meta no-op / no-op). + +Articulation / Robot: + +- `TestArticulationNewton`: control API, setters, drive, per-link mass live via + `set_link_physical_attr`, remove. +- `TestRobotNewton`: spawn (URDF assembly), finalize, control-part resolution, + qpos round-trip via the Newton articulation view. + +Gym: + +- Rigid-only Newton env initializes, steps, resets, and reads observations. + +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 + +- The `add_robot`-on-Newton path depends on the upstream dexsim fix + (`_joint_metas_from_ids` active-joint indexing, dexsim + `yueci/adapt-embodichain` `d0e86bb02`). If dexsim is rebuilt from a different + ref, `supports_robot` would need re-gating. +- dexsim's Newton path hardcodes `density=0.0` in its desc resolver; EmbodiChain's + `resolve_newton_shape` sets `density` from the cfg (positive) to avoid the + desc-path mass gap where dynamic bodies without explicit mass+inertia fail to + compute a positive body mass. Watch for dexsim changing this. +- 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; Newton is guarded from those paths. +- Runtime shape/property mutations may require model rebuilds rather than live + updates; Newton-native per-link contact params are build-time only. +- Standalone Newton scripts can segfault during teardown (`sim.destroy()` + + `teardown_newton_physics()`); pytest's `flush_cleanup_queue` teardown path is + stable — use the pytest pattern, not bare scripts. 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/docs/superpowers/plans/2026-06-22-newton-backend-pr.md b/docs/superpowers/plans/2026-06-22-newton-backend-pr.md new file mode 100644 index 00000000..06856d25 --- /dev/null +++ b/docs/superpowers/plans/2026-06-22-newton-backend-pr.md @@ -0,0 +1,1234 @@ +# Newton Backend PR Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Finish the Newton-backend PR. **Target 4 (multi-env) was found +already complete during execution** — EmbodiChain's spawn path already does +prototype+clone across arenas at spawn time, and Newton views already handle +multi-env entity lists. This plan therefore covers only **Target 5 +(differentiable env for APG)** plus branch cleanup and docs. + +**Architecture:** A new `embodichain.lab.sim.diff` package provides a +`torch.autograd.Function` bridge over +`dexsim.engine.newton_physics.DifferentiableStepper`; a new +`DifferentiableEmbodiedEnv` gym subclass wires it into the standard +EmbodiChain env step pipeline. `SimulationManager` gains thin delegators to +dexsim's `create_differentiable_stepper` / `create_gradient_rollout`. + +**Revision history:** Original plan had Tasks 1–4 covering multi-env clone +scaffolding, spawn guards, clone-at-finalize, and body-id resolution. Those +were deleted after code inspection showed the spawn path +(`spawn_rigid_object_entities` → `_spawn_clones_from_prototype`) already +clones prototypes into all arenas at spawn time, Newton views already accept +multi-entity lists, and existing tests (`TestRigidObjectNewton` with +`NUM_ARENAS=2`, `test_spawn_clones_distinct_entities`, +`test_newton_native_attrs_desc_native_spawn` asserting +`obj.num_instances == NUM_ARENAS`) already pass. Task numbers below are +rebased: old Task 5 → Task 1, old Task 6 → Task 2, etc. + +**Tech Stack:** Python 3.10+, PyTorch (autograd), NVIDIA Warp (`wp.Tape`, +`wp.to_torch`/`wp.from_torch`), DexSim Newton physics +(`dexsim.engine.newton_physics`), gymnasium, pytest. + +**Companion spec:** `docs/superpowers/specs/2026-06-21-newton-backend-pr-design.md` + +--- + +## File Map + +**Created:** +- `embodichain/lab/sim/diff/__init__.py` — public re-exports for the diff package +- `embodichain/lab/sim/diff/bridge.py` — `NewtonStepFunc(torch.autograd.Function)`, `tape_context`, `differentiable_step` +- `embodichain/lab/gym/envs/differentiable_env.py` — `DifferentiableEmbodiedEnv` subclass +- `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` — Franka APG example task +- `tests/sim/test_differentiable_stepper.py` +- `tests/gym/envs/test_differentiable_env.py` +- `agent_context/topics/differentiable-env.md` + +**Modified:** +- `embodichain/lab/sim/sim_manager.py` — add `create_differentiable_stepper` / `create_gradient_rollout` delegators +- `agent_context/MAP.yaml` — register new `differentiable-env` topic +- `design/newton-backend-design.md` — mark Target 5 done, link to plan + +**Already complete (Target 4, verified during execution):** +- Multi-env clone-at-spawn: `embodichain/lab/sim/utility/sim_utils.py:spawn_rigid_object_entities` / `spawn_articulation_entities` already prototype-then-clone across all arenas via dexsim's `clone_actor_to` (Newton-patched). +- Newton multi-env views: `embodichain/lab/sim/objects/backends/newton.py:NewtonRigidBodyView` / `NewtonArticulationView` already accept `Sequence[MeshObject]` and resolve one body ID per entity. +- Newton multi-env tests: `tests/sim/objects/test_rigid_object.py::TestRigidObjectNewton` (NUM_ARENAS=2, `test_spawn_clones_distinct_entities`, `test_newton_native_attrs_desc_native_spawn` asserting `obj.num_instances == NUM_ARENAS`), `tests/sim/objects/test_articulation.py::TestArticulationNewton` (num_envs=2), `tests/sim/objects/test_robot.py` (num_envs=10). + +--- + + +## Task 1: Add `create_differentiable_stepper` / `create_gradient_rollout` delegators + +**Files:** +- Modify: `embodichain/lab/sim/sim_manager.py` +- Test: `tests/sim/test_differentiable_stepper.py` + +- [ ] **Step 1: Write the failing test** + +Create `tests/sim/test_differentiable_stepper.py`: + +```python +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Licensed under the Apache License, Version 2.0 (the "License"); +# ---------------------------------------------------------------------------- +"""Tests for the differentiable-stepper delegators on SimulationManager.""" + +from __future__ import annotations + +import pytest + +from embodichain.lab.sim.cfg import DefaultPhysicsCfg, NewtonPhysicsCfg +from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg + + +def test_default_backend_rejects_differentiable_stepper(): + sim = SimulationManager(SimulationManagerCfg( + physics_cfg=DefaultPhysicsCfg(), num_envs=1, headless=True, + )) + with pytest.raises(Exception, match=r"Newton"): + sim.create_differentiable_stepper() + SimulationManager.reset() + + +def test_newton_without_grad_rejects_differentiable_stepper(): + sim = SimulationManager(SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg(requires_grad=False, use_cuda_graph=False), + num_envs=1, headless=True, + )) + sim.finalize_newton_physics() + with pytest.raises(Exception, match=r"grad"): + sim.create_differentiable_stepper() + SimulationManager.reset() + + +def test_newton_with_grad_creates_stepper(): + sim = SimulationManager(SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + requires_grad=True, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=1, headless=True, + )) + sim.finalize_newton_physics() + stepper = sim.create_differentiable_stepper() + from dexsim.engine.newton_physics.differentiable_stepper import ( + DifferentiableStepper, + ) + assert isinstance(stepper, DifferentiableStepper) + SimulationManager.reset() +``` + +- [ ] **Step 2: Run the tests and confirm they fail** + +Run: `pytest -q tests/sim/test_differentiable_stepper.py` +Expected: FAIL — `create_differentiable_stepper` not defined. + +- [ ] **Step 3: Add the delegator methods** + +Edit `embodichain/lab/sim/sim_manager.py` — add near the other Newton +back-compat delegators (search `newton_manager` in the file to locate the +right region): + +```python + def create_differentiable_stepper(self): + """Create a single-step differentiable physics primitive (Newton-only). + + Requires the Newton backend with ``requires_grad=True`` and + ``solver_type="semi_implicit"``. Delegates to + :meth:`dexsim.engine.newton_physics.NewtonManager.create_differentiable_stepper`. + + Raises: + RuntimeError: If the active backend is not Newton or if the + Newton manager is not ready / not in grad mode. + """ + if not self.is_newton_backend: + logger.log_error( + "create_differentiable_stepper requires the Newton backend.") + return self.physics.newton_manager.create_differentiable_stepper() + + def create_gradient_rollout( + self, + record_steps: int, + substeps_per_record: int | None = None, + record_dt: float | None = None, + ): + """Create a gradient rollout buffer (Newton-only). + + Delegates to + :meth:`dexsim.engine.newton_physics.NewtonManager.create_gradient_rollout`. + """ + if not self.is_newton_backend: + logger.log_error( + "create_gradient_rollout requires the Newton backend.") + return self.physics.newton_manager.create_gradient_rollout( + record_steps=record_steps, + substeps_per_record=substeps_per_record, + record_dt=record_dt, + ) +``` + +- [ ] **Step 4: Run the tests and confirm they pass** + +Run: `pytest -q tests/sim/test_differentiable_stepper.py` +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add embodichain/lab/sim/sim_manager.py tests/sim/test_differentiable_stepper.py +git commit -m "feat(sim): SimulationManager delegators for Newton diff stepper + +create_differentiable_stepper and create_gradient_rollout are thin +passthroughs to NewtonManager. Both raise on the default backend. +Backs the new embodichain.lab.sim.diff package (next commit)." +``` + +--- + +## Task 2: Create the `embodichain.lab.sim.diff` package — bridge + +**Files:** +- Create: `embodichain/lab/sim/diff/__init__.py` +- Create: `embodichain/lab/sim/diff/bridge.py` + +The bridge wraps a `wp.Tape()` around one EmbodiChain physics step and +exposes a `torch.autograd.Function` so callers can drive APG with +PyTorch-side action tensors. + +- [ ] **Step 1: Create the package skeleton** + +Create `embodichain/lab/sim/diff/__init__.py`: + +```python +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Licensed under the Apache License, Version 2.0 (the "License"); +# ---------------------------------------------------------------------------- +"""Differentiable Newton stepping for EmbodiChain. + +Bridges DexSim's :class:`~dexsim.engine.newton_physics.DifferentiableStepper` +into PyTorch autograd via a :class:`torch.autograd.Function`, and exposes a +:class:`tape_context` manager for advanced users who want to compose their +own Warp kernels. +""" + +from __future__ import annotations + +from .bridge import ( + NewtonStepFunc, + differentiable_step, + tape_context, +) + +__all__ = [ + "NewtonStepFunc", + "differentiable_step", + "tape_context", +] +``` + +- [ ] **Step 2: Create `bridge.py`** + +Create `embodichain/lab/sim/diff/bridge.py`: + +```python +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Licensed under the Apache License, Version 2.0 (the "License"); +# ---------------------------------------------------------------------------- +"""Warp-tape <-> PyTorch-autograd bridge for Newton physics.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Callable, Iterator + +import torch +import warp as wp + +if TYPE_CHECKING: + from embodichain.lab.sim.sim_manager import SimulationManager + +__all__ = ["NewtonStepFunc", "differentiable_step", "tape_context"] + + +@contextmanager +def tape_context(manager: "SimulationManager") -> Iterator[wp.Tape]: + """Open a Warp tape bound to the manager's Newton state. + + Advanced users compose their own Warp kernels inside this context, then + call ``tape.backward()`` outside the with-block. + """ + if not manager.is_newton_backend: + raise RuntimeError( + "tape_context requires the Newton backend with requires_grad=True.") + tape = wp.Tape() + with tape: + yield tape + + +def differentiable_step( + manager: "SimulationManager", + *, + apply_control_fn: Callable[[wp.Tape], None], + substeps: int, + dt: float | None = None, +) -> dict: + """Run one EmbodiChain-level physics step inside a Warp tape. + + Args: + manager: The owning :class:`SimulationManager` (must be Newton). + apply_control_fn: Callable that writes the joint/body control + targets inside the tape. Invoked once at the start of the + step. Receives the open tape; must launch Warp kernels (or + call dexsim setters that are tape-aware) to populate + ``manager.physics.newton_manager._control``. + substeps: Number of solver substeps to run (typically + ``sim_cfg.sim_steps_per_control``). + dt: Solver dt; defaults to the manager's configured dt. + + Returns: + A dict carrying the tape and the state buffers for the caller to + save in autograd context. + """ + if not manager.is_newton_backend: + raise RuntimeError( + "differentiable_step requires the Newton backend.") + nm = manager.physics.newton_manager + stepper = manager.create_differentiable_stepper() + state_in = nm._state_0 + state_out = nm._model.state() + contacts = stepper.create_contacts() + dt_val = nm.solver_dt if dt is None else float(dt) + + tape = wp.Tape() + with tape: + apply_control_fn(tape) + for _ in range(substeps): + stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) + state_in, state_out = state_out, state_in + + # The final state lives in state_in after the swap. + return { + "tape": tape, + "final_state": state_in, + "stepper": stepper, + } + + +class NewtonStepFunc(torch.autograd.Function): + """torch.autograd.Function bridging Warp tape autodiff to PyTorch. + + Forward: launches the action-to-control Warp kernel, runs + ``substeps`` differentiable solver steps, and reads observation / + reward as torch tensors via ``wp.to_torch`` (zero-copy where + possible). + + Backward: copies upstream grads into the corresponding Warp + ``.grad`` buffers, calls ``tape.backward()``, and returns + ``wp.to_torch(action.grad)`` reshaped to the action's tensor shape. + + Callers must supply a ``sim_state`` dict with the following keys: + manager: SimulationManager (Newton, requires_grad=True) + substeps: int + action_to_control_kernel: callable(action_wp, *kernel_args) + kernel_args: tuple consumed by action_to_control_kernel + obs_reward_fn: callable(final_state) -> dict with torch outputs + """ + + @staticmethod + def forward(ctx, action_torch: torch.Tensor, sim_state: dict): + manager = sim_state["manager"] + substeps = int(sim_state["substeps"]) + kernel = sim_state["action_to_control_kernel"] + kernel_args = sim_state["kernel_args"] + obs_reward_fn = sim_state["obs_reward_fn"] + + nm = manager.physics.newton_manager + stepper = manager.create_differentiable_stepper() + + action_flat = action_torch.detach().clone().reshape(-1).contiguous() + action_wp = wp.from_torch(action_flat, dtype=wp.float32, requires_grad=True) + + state_in = nm._state_0 + state_out = nm._model.state() + contacts = stepper.create_contacts() + dt_val = nm.solver_dt + + tape = wp.Tape() + with tape: + kernel(action_wp, *kernel_args) # writes nm._control inside tape + for _ in range(substeps): + stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) + state_in, state_out = state_out, state_in + + outputs = obs_reward_fn(state_in) + ctx.tape = tape + ctx.action_wp = action_wp + ctx.outputs_wp = outputs.get("_grad_track", {}) + # `outputs` is a dict of torch tensors built from wp.to_torch — the + # caller is responsible for ensuring at least one is grad-tracked. + return tuple(outputs[k] for k in outputs["_order"]) + + @staticmethod + def backward(ctx, *grad_outputs): + # Copy each upstream grad back into the corresponding Warp .grad. + for name, grad_t in zip(ctx.outputs_wp["_order"], grad_outputs): + wp_arr = ctx.outputs_wp[name] + if grad_t is None or wp_arr.grad is None: + continue + wp.copy(wp_arr.grad, + wp.from_torch(grad_t.detach().clone().contiguous(), + dtype=wp.float32)) + ctx.tape.backward() + action_grad = wp.to_torch(ctx.action_wp.grad).clone() + ctx.tape.zero() + # Reshape to the original action layout; second input (sim_state) + # has no gradient. + return action_grad.reshape(ctx.saved_action_shape), None +``` + +> Note: the contract between `obs_reward_fn` and `NewtonStepFunc.backward` +> is intentionally explicit — the caller (the env in Task 3) constructs the +> dict in a way that records which outputs need grad-tracking. The +> `_order` / `_grad_track` plumbing keeps the autograd function fully +> general; the env class hides it from end users. + +- [ ] **Step 3: Lightweight import smoke** + +Run: `python -c "from embodichain.lab.sim.diff import NewtonStepFunc, tape_context, differentiable_step; print('ok')"` +Expected: prints `ok`. + +- [ ] **Step 4: Append a tape-smoke test** + +Append to `tests/sim/test_differentiable_stepper.py`: + +```python +def test_tape_context_records_step(): + import warp as wp + + sim = SimulationManager(SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + requires_grad=True, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=1, headless=True, + )) + sim.finalize_newton_physics() + from embodichain.lab.sim.diff import tape_context + + with tape_context(sim) as tape: + pass # empty tape is valid; tape.backward() on empty is a no-op + + assert isinstance(tape, wp.Tape) + SimulationManager.reset() +``` + +- [ ] **Step 5: Run all diff-stepper tests** + +Run: `pytest -q tests/sim/test_differentiable_stepper.py` +Expected: 4 PASS. + +- [ ] **Step 6: Commit** + +```bash +git add embodichain/lab/sim/diff/__init__.py \ + embodichain/lab/sim/diff/bridge.py \ + tests/sim/test_differentiable_stepper.py +git commit -m "feat(sim/diff): Warp-tape <-> PyTorch-autograd bridge + +New embodichain.lab.sim.diff package: NewtonStepFunc (autograd.Function) +wraps DifferentiableStepper inside a wp.Tape, tape_context is the +low-level context manager for advanced kernels, differentiable_step is +the convenience wrapper. Foundation for DifferentiableEmbodiedEnv." +``` + +--- + +## Task 3: `DifferentiableEmbodiedEnv` gym subclass + +**Files:** +- Create: `embodichain/lab/gym/envs/differentiable_env.py` +- Test: `tests/gym/envs/test_differentiable_env.py` + +- [ ] **Step 1: Inspect `EmbodiedEnv.step` signature** + +Run: `Read embodichain/lab/gym/envs/embodied_env.py` (focus on `step`, +`reset`, `_preprocess_action`, `_step_action`). + +Identify exactly which methods produce the per-step `obs, reward, done, +info`. The override must invoke the same observation/reward managers as +the base class — just inside a tape. + +- [ ] **Step 2: Write the construction-validation test first** + +Create `tests/gym/envs/test_differentiable_env.py`: + +```python +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Licensed under the Apache License, Version 2.0 (the "License"); +# ---------------------------------------------------------------------------- +"""Tests for DifferentiableEmbodiedEnv.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.sim.cfg import DefaultPhysicsCfg, NewtonPhysicsCfg +from embodichain.lab.gym.envs.differentiable_env import ( + DifferentiableEmbodiedEnv, +) +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg + + +def _diff_env_cfg(requires_grad: bool = True, backend: str = "newton") -> EmbodiedEnvCfg: + from embodichain.lab.sim.sim_manager import SimulationManagerCfg + + if backend == "newton": + physics_cfg = NewtonPhysicsCfg( + requires_grad=requires_grad, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ) + else: + physics_cfg = DefaultPhysicsCfg() + sim_cfg = SimulationManagerCfg( + physics_cfg=physics_cfg, num_envs=2, headless=True, + ) + return EmbodiedEnvCfg(sim_cfg=sim_cfg) + + +def test_construct_without_requires_grad_raises(): + with pytest.raises(Exception, match=r"requires_grad"): + DifferentiableEmbodiedEnv(_diff_env_cfg(requires_grad=False)) + + +def test_construct_on_default_backend_raises(): + with pytest.raises(Exception, match=r"Newton"): + DifferentiableEmbodiedEnv(_diff_env_cfg(backend="default")) +``` + +- [ ] **Step 3: Run and confirm failure** + +Run: `pytest -q tests/gym/envs/test_differentiable_env.py` +Expected: FAIL — `DifferentiableEmbodiedEnv` not defined. + +- [ ] **Step 4: Implement `DifferentiableEmbodiedEnv`** + +Create `embodichain/lab/gym/envs/differentiable_env.py`: + +```python +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Licensed under the Apache License, Version 2.0 (the "License"); +# ---------------------------------------------------------------------------- +"""Differentiable Newton-backed EmbodiedEnv for analytic policy gradient. + +Wraps the standard :class:`EmbodiedEnv` step pipeline in a Warp tape and +bridges autograd into PyTorch via +:class:`embodichain.lab.sim.diff.NewtonStepFunc`. Subclasses define how +actions become Newton control writes and how observations/rewards are +read from the post-step state; the bridge handles the tape lifecycle +and the backward pass. + +Usage: + + class MyTask(DifferentiableEmbodiedEnv): + def _apply_action_kernel(self, action_wp, tape): ... + def _read_outputs(self, final_state) -> dict: ... +""" + +from __future__ import annotations + +from abc import abstractmethod +from typing import Any + +import torch + +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg +from embodichain.lab.sim.diff import NewtonStepFunc +from embodichain.utils import logger + +__all__ = ["DifferentiableEmbodiedEnv"] + + +class DifferentiableEmbodiedEnv(EmbodiedEnv): + """EmbodiedEnv variant that exposes APG-ready :py:meth:`step`. + + Subclasses must implement :meth:`_apply_action_kernel` and + :meth:`_read_outputs`; the rest of the EmbodiedEnv contract (reset, + observation managers, reward functors) carries over. + """ + + def __init__(self, cfg: EmbodiedEnvCfg, *args, **kwargs) -> None: + self._validate_diff_cfg(cfg) + super().__init__(cfg, *args, **kwargs) + self._truncate_backward_at: int | None = getattr( + cfg, "truncate_backward_at", None, + ) + + @staticmethod + def _validate_diff_cfg(cfg: EmbodiedEnvCfg) -> None: + physics_cfg = cfg.sim_cfg.physics_cfg + if not isinstance(physics_cfg, NewtonPhysicsCfg): + logger.log_error( + "DifferentiableEmbodiedEnv requires NewtonPhysicsCfg, " + f"got {type(physics_cfg).__name__}.") + if not physics_cfg.requires_grad: + logger.log_error( + "DifferentiableEmbodiedEnv requires requires_grad=True on " + "the NewtonPhysicsCfg.") + + # -- subclass contract ------------------------------------------------ # + + @abstractmethod + def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: + """Inside the open Warp tape, write the action into Newton control. + + Implementations launch a Warp kernel that reads ``action_wp`` + (a ``wp.array(dtype=wp.float32, requires_grad=True)`` of shape + ``[num_envs * action_dim]``) and writes into + ``self.sim.physics.newton_manager._control`` so the next stepper + call uses the new control. + """ + + @abstractmethod + def _read_outputs(self, final_state: Any) -> dict: + """Read the post-step observation and reward as torch tensors. + + Must return a dict with keys ``"obs"``, ``"reward"``, + ``"terminated"``, ``"truncated"``, ``"info"``, plus the + ``_order``/``_grad_track`` metadata expected by + :class:`NewtonStepFunc`. ``obs`` and ``reward`` should be torch + tensors backed by ``wp.to_torch`` of grad-tracked Warp arrays. + """ + + # -- gym surface ------------------------------------------------------ # + + def step(self, action: torch.Tensor): + if not isinstance(action, torch.Tensor): + action = torch.as_tensor(action, dtype=torch.float32) + sim_state = self._build_sim_state_dict(action) + outputs = NewtonStepFunc.apply(action, sim_state) + obs, reward, terminated, truncated = outputs[:4] + info = sim_state["last_info"] + + done_mask = terminated | truncated + if done_mask.any(): + reset_ids = done_mask.nonzero(as_tuple=False).squeeze(-1) + fresh_obs, _ = self.reset(env_ids=reset_ids) + obs = torch.where( + done_mask.unsqueeze(-1).expand_as(obs), + fresh_obs.detach(), obs, + ) + return obs, reward, terminated, truncated, info + + def _build_sim_state_dict(self, action: torch.Tensor) -> dict: + # Pack the args NewtonStepFunc expects. Subclass-supplied kernel + # + output reader; environment-level metadata stays here. + return { + "manager": self.sim, + "substeps": self.sim_cfg.sim_steps_per_control, + "action_to_control_kernel": self._wrap_action_kernel(), + "kernel_args": (), + "obs_reward_fn": self._read_outputs, + "last_info": {}, + } + + def _wrap_action_kernel(self): + env = self + def _inner(action_wp, *_): + env._apply_action_kernel(action_wp, tape=None) + return _inner +``` + +- [ ] **Step 5: Re-run construction tests** + +Run: `pytest -q tests/gym/envs/test_differentiable_env.py::test_construct_without_requires_grad_raises tests/gym/envs/test_differentiable_env.py::test_construct_on_default_backend_raises` +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add embodichain/lab/gym/envs/differentiable_env.py \ + tests/gym/envs/test_differentiable_env.py +git commit -m "feat(gym): DifferentiableEmbodiedEnv for APG + +Newton-only EmbodiedEnv subclass that wraps step() in a Warp tape via +NewtonStepFunc. Subclasses implement _apply_action_kernel and +_read_outputs; the base class handles validation, auto-reset on done, +and the autograd bridge." +``` + +--- + +## Task 4: Franka reach APG example task + +**Files:** +- Create: `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` +- Test: `tests/gym/envs/test_differentiable_env.py` (append) + +Model the example after +`/root/sources/analytic_policy_gradients/envs/franka_reach_env.py`, but +built on EmbodiChain primitives (`add_robot` with the Franka URDF, the +`DifferentiableEmbodiedEnv` base). + +- [ ] **Step 1: Locate Franka URDF in EmbodiChain data** + +Run: `find embodichain/data -iname "fr3*.urdf" -o -iname "*franka*.urdf" | head -5` + +If no URDF is bundled, the example accepts a `urdf_path` override and +falls back to `newton.utils.download_asset("franka_emika_panda")`, +matching the reference env. + +- [ ] **Step 2: Write the example task** + +Create `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` (full +contents below): + +```python +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# Licensed under the Apache License, Version 2.0 (the "License"); +# ---------------------------------------------------------------------------- +"""Franka FR3 reach task with differentiable Newton physics (APG).""" + +from __future__ import annotations + +import math +from typing import Any + +import numpy as np +import torch +import warp as wp + +from embodichain.lab.gym.envs.differentiable_env import DifferentiableEmbodiedEnv +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg +from embodichain.lab.gym.utils.registry import register_env +from embodichain.lab.sim.cfg import ( + NewtonPhysicsCfg, + RobotCfg, + URDFCfg, +) +from embodichain.lab.sim.sim_manager import SimulationManagerCfg + +FRANKA_NUM_ARM_JOINTS = 7 +FRANKA_EE_BODY = "fr3_hand_tcp" +DEFAULT_ACTION_SCALE = 0.2 +DEFAULT_MAX_EPISODE_STEPS = 30 +TARGET_POS_RANGE = { + "x": (0.05, 0.70), + "y": (-0.45, 0.45), + "z": (0.20, 0.95), +} +TARGET_MAX_TILT = math.pi / 3 + + +@wp.kernel +def _set_joint_targets_kernel( + action: wp.array(dtype=wp.float32), + current_q: wp.array(dtype=wp.float32), + target_q: wp.array(dtype=wp.float32), + limit_lo: wp.array(dtype=wp.float32), + limit_hi: wp.array(dtype=wp.float32), + action_scale: wp.float32, + n_joints_per_env: wp.int32, + n_arm: wp.int32, + total: wp.int32, +): + tid = wp.tid() + if tid < total: + env_idx = tid / n_arm + j = tid % n_arm + off = env_idx * n_joints_per_env + j + new_q = current_q[off] + action[tid] * action_scale + target_q[off] = wp.clamp(new_q, limit_lo[j], limit_hi[j]) + + +@register_env("FrankaReachApg-v0") +class FrankaReachApgEnv(DifferentiableEmbodiedEnv): + """Differentiable Franka FR3 reach task. + + Built on EmbodiChain's :class:`DifferentiableEmbodiedEnv`; the + Warp-tape bridge produces ``action.grad`` that flows back through the + semi-implicit Newton solver. + """ + + metadata = {"render_modes": ["human"], "default_num_envs": 4} + + def __init__( + self, + cfg: EmbodiedEnvCfg | None = None, + *, + num_envs: int = 4, + urdf_path: str | None = None, + action_scale: float = DEFAULT_ACTION_SCALE, + max_episode_steps: int = DEFAULT_MAX_EPISODE_STEPS, + device: str = "cuda:0", + ) -> None: + if cfg is None: + cfg = EmbodiedEnvCfg( + sim_cfg=SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + device=device, + requires_grad=True, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=num_envs, + headless=True, + ), + ) + self._urdf_path = urdf_path + self._action_scale = float(action_scale) + self._max_episode_steps = int(max_episode_steps) + super().__init__(cfg) + self._init_franka() + self._init_targets() + + # -- scene setup ----------------------------------------------------- # + + def _init_franka(self) -> None: + urdf = self._urdf_path or self._resolve_default_urdf() + robot_cfg = RobotCfg( + uid="franka", + urdf_cfg=URDFCfg().set_urdf(urdf), + fix_base=True, + ) + self._robot = self.sim.add_robot(robot_cfg) + self.sim.finalize_newton_physics() + + # Cache joint-limit Warp arrays for the action kernel. + model = self.sim.physics.newton_manager._model + lo = np.asarray(model.joint_limit_lower[:FRANKA_NUM_ARM_JOINTS], + dtype=np.float32) + hi = np.asarray(model.joint_limit_upper[:FRANKA_NUM_ARM_JOINTS], + dtype=np.float32) + self._limit_lo_wp = wp.array(lo, dtype=wp.float32, device=model.device) + self._limit_hi_wp = wp.array(hi, dtype=wp.float32, device=model.device) + self._n_joints_per_env = int(len(model.joint_q) // self.sim.num_envs) + + def _resolve_default_urdf(self) -> str: + try: + import newton.utils as nu + + urdf = nu.download_asset("franka_emika_panda") / ( + "urdf/fr3_franka_hand.urdf") + if urdf.exists(): + return str(urdf) + except Exception: + pass + raise FileNotFoundError( + "Franka URDF not available; pass urdf_path explicitly.") + + def _init_targets(self) -> None: + n = self.sim.num_envs + device = self.device + self.target_pos = torch.zeros(n, 3, device=device) + self.target_quat = torch.zeros(n, 4, device=device) + self.last_action = torch.zeros( + n, FRANKA_NUM_ARM_JOINTS, device=device, + ) + self.step_count = torch.zeros(n, dtype=torch.int32, device=device) + self._sample_new_targets(torch.arange(n, device=device)) + + def _sample_new_targets(self, env_ids: torch.Tensor) -> None: + n = env_ids.numel() + d = self.device + self.target_pos[env_ids, 0] = ( + TARGET_POS_RANGE["x"][0] + + torch.rand(n, device=d) + * (TARGET_POS_RANGE["x"][1] - TARGET_POS_RANGE["x"][0])) + self.target_pos[env_ids, 1] = ( + TARGET_POS_RANGE["y"][0] + + torch.rand(n, device=d) + * (TARGET_POS_RANGE["y"][1] - TARGET_POS_RANGE["y"][0])) + self.target_pos[env_ids, 2] = ( + TARGET_POS_RANGE["z"][0] + + torch.rand(n, device=d) + * (TARGET_POS_RANGE["z"][1] - TARGET_POS_RANGE["z"][0])) + # Identity-ish quat, no tilt for the smoke task. + self.target_quat[env_ids] = torch.tensor( + [1.0, 0.0, 0.0, 0.0], device=d).expand(n, -1) + + # -- DifferentiableEmbodiedEnv contract ------------------------------ # + + def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: + nm = self.sim.physics.newton_manager + n_envs = self.sim.num_envs + total = n_envs * FRANKA_NUM_ARM_JOINTS + wp.launch( + _set_joint_targets_kernel, + dim=total, + inputs=[ + action_wp, + nm._state_0.joint_q, + nm._control.joint_target, + self._limit_lo_wp, + self._limit_hi_wp, + wp.float32(self._action_scale), + wp.int32(self._n_joints_per_env), + wp.int32(FRANKA_NUM_ARM_JOINTS), + wp.int32(total), + ], + device=nm._model.device, + ) + + def _read_outputs(self, final_state: Any) -> dict: + nm = self.sim.physics.newton_manager + n = self.sim.num_envs + body_q = wp.to_torch(final_state.body_q).view(n, -1, 7) + ee_idx = self._ee_body_indices() + ee_pose = body_q[torch.arange(n, device=self.device), ee_idx] + eef_pos = ee_pose[:, :3] + eef_quat = ee_pose[:, 3:] + + pos_dist = (eef_pos - self.target_pos).norm(dim=-1) + rot_dist = self._quat_distance(eef_quat, self.target_quat) + reward = ( + -0.2 * pos_dist + + 0.1 * torch.exp(-(pos_dist ** 2) / (2 * 0.1 ** 2)) + - 0.1 * rot_dist + + 0.1 * torch.exp(-(rot_dist ** 2) / (2 * 0.3 ** 2)) + ) + + obs = torch.cat([ + wp.to_torch(final_state.joint_q).view(n, -1)[:, :FRANKA_NUM_ARM_JOINTS], + ee_pose, + self.target_pos, + self.target_quat, + self.last_action, + ], dim=-1) + + terminated = (pos_dist < 0.01) & (rot_dist < 0.3) + self.step_count += 1 + truncated = self.step_count >= self._max_episode_steps + + return { + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": { + "obs": final_state.joint_q, + "reward": None, # reward is torch-built; tape flows via FK + }, + "obs": obs, + "reward": reward, + "terminated": terminated, + "truncated": truncated, + } + + def _ee_body_indices(self) -> torch.Tensor: + if hasattr(self, "_cached_ee_idx"): + return self._cached_ee_idx + model = self.sim.physics.newton_manager._model + idx_per_env = [] + n_per_env = len(model.body_label) // self.sim.num_envs + for i in range(self.sim.num_envs): + for j, label in enumerate(model.body_label): + if FRANKA_EE_BODY in str(label) and (j // n_per_env) == i: + idx_per_env.append(j) + break + self._cached_ee_idx = torch.tensor(idx_per_env, dtype=torch.long, + device=self.device) + return self._cached_ee_idx + + @staticmethod + def _quat_distance(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: + return torch.minimum(((q1 - q2) ** 2).sum(-1), + ((q1 + q2) ** 2).sum(-1)) + + # -- gym overrides --------------------------------------------------- # + + def reset(self, *, seed: int | None = None, env_ids=None): + env_ids = (env_ids if env_ids is not None + else torch.arange(self.sim.num_envs, device=self.device)) + with torch.no_grad(): + self.step_count[env_ids] = 0 + self.last_action[env_ids] = 0.0 + self._sample_new_targets(env_ids) + # Reset Newton joint_q to zero for the touched envs. + jq = wp.to_torch( + self.sim.physics.newton_manager._state_0.joint_q, + ).view(self.sim.num_envs, -1) + jq[env_ids] = 0.0 + obs = self._initial_obs() + return obs, {} + + def _initial_obs(self) -> torch.Tensor: + with torch.no_grad(): + return self._read_outputs( + self.sim.physics.newton_manager._state_0)["obs"] +``` + +- [ ] **Step 3: Smoke import the task** + +Run: `python -c "from embodichain.lab.gym.envs.tasks.special.franka_reach_apg import FrankaReachApgEnv; print('ok')"` +Expected: `ok`. + +- [ ] **Step 4: Append the smoke test** + +Append to `tests/gym/envs/test_differentiable_env.py`: + +```python +@pytest.mark.requires_gpu +def test_franka_apg_smoke_backward(): + try: + from embodichain.lab.gym.envs.tasks.special.franka_reach_apg import ( + FrankaReachApgEnv, + ) + except FileNotFoundError as e: + pytest.skip(f"Franka URDF not available: {e}") + + env = FrankaReachApgEnv(num_envs=2) + env.reset(seed=0) + action = torch.zeros(2, 7, requires_grad=True, device=env.device) + obs, reward, terminated, truncated, info = env.step(action) + assert reward.requires_grad, "Reward must be autograd-tracked." + loss = reward.sum() + loss.backward() + assert action.grad is not None + assert torch.isfinite(action.grad).all() + env.close() + + +@pytest.mark.requires_gpu +def test_franka_apg_one_iter_loss_reduces(): + try: + from embodichain.lab.gym.envs.tasks.special.franka_reach_apg import ( + FrankaReachApgEnv, + ) + except FileNotFoundError as e: + pytest.skip(f"Franka URDF not available: {e}") + + env = FrankaReachApgEnv(num_envs=2) + env.reset(seed=0) + action = torch.zeros(2, 7, requires_grad=True, device=env.device) + opt = torch.optim.SGD([action], lr=0.01) + + losses = [] + for _ in range(3): + env.reset(seed=0) + opt.zero_grad() + _, reward, _, _, _ = env.step(action) + loss = (-reward).sum() + loss.backward() + opt.step() + losses.append(loss.detach().item()) + assert losses[-1] < losses[0], ( + f"APG did not reduce loss: {losses}") + env.close() +``` + +- [ ] **Step 5: Run all differentiable-env tests on a GPU host** + +Run: `pytest -q tests/gym/envs/test_differentiable_env.py` +Expected: 4 PASS (or smoke tests SKIPPED if URDF unavailable / no GPU). + +- [ ] **Step 6: Commit** + +```bash +git add embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py \ + tests/gym/envs/test_differentiable_env.py +git commit -m "feat(gym/tasks): Franka FR3 reach APG example + +End-to-end APG smoke task built on DifferentiableEmbodiedEnv with a +Warp action-to-control kernel and torch-built reward. Verifies the +autograd bridge with one-iteration loss reduction. URDF resolved from +newton.utils.download_asset with explicit override." +``` + +--- + +## Task 5: Documentation — agent_context topic and design doc update + +**Files:** +- Create: `agent_context/topics/differentiable-env.md` +- Modify: `agent_context/MAP.yaml` +- Modify: `design/newton-backend-design.md` + +- [ ] **Step 1: Inspect MAP.yaml format** + +Run: `Read agent_context/MAP.yaml` + +Note the existing entries (`env-framework`, `manager-functor`, ...) and +mirror that structure. + +- [ ] **Step 2: Create the topic file** + +Create `agent_context/topics/differentiable-env.md`: + +```markdown +# Differentiable Env (APG) Context + +EmbodiChain supports analytic policy gradient (APG) via +:class:`embodichain.lab.gym.envs.differentiable_env.DifferentiableEmbodiedEnv`. +The bridge wraps `dexsim.engine.newton_physics.DifferentiableStepper` +inside a `wp.Tape()` and exposes a `torch.autograd.Function` +(`embodichain.lab.sim.diff.NewtonStepFunc`) so PyTorch-side `action` +tensors get a gradient from `tape.backward()`. + +## Required configuration + +- `NewtonPhysicsCfg(requires_grad=True, solver_cfg={"solver_type": "semi_implicit"})` +- `use_cuda_graph=False` (forced by dexsim when grad mode is on) + +The default backend and any other Newton solver are rejected. + +## Subclass contract + +Task authors implement two methods on `DifferentiableEmbodiedEnv`: + +- `_apply_action_kernel(action_wp, tape)` — launch a Warp kernel that + writes joint/body targets into `nm._control` while the tape is open. +- `_read_outputs(final_state)` — build the `obs` / `reward` / `done` + outputs as torch tensors via `wp.to_torch` so the tape can record the + dependency. + +See `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` for the +canonical example. + +## Functor autograd compatibility + +Reward/observation functors that compose torch operations on tensors +obtained via `wp.to_torch` are automatically autograd-compatible. +Functors that detour through CPU / NumPy break the graph; those need +torch-only reimplementations for the differentiable path. + +## Memory + +Each step records `sim_steps_per_control` substeps into the tape. For +long horizons or large `num_envs`, pass `truncate_backward_at=K` on the +env config to split the tape and detach at chunk boundaries. +``` + +- [ ] **Step 3: Register the topic in `MAP.yaml`** + +Edit `agent_context/MAP.yaml` — append: + +```yaml +- id: differentiable-env + aliases: ["apg", "analytic-policy-gradient", "differentiable-rl"] + keywords: [differentiable, gradient, apg, autograd, warp tape] + files: + - topics/differentiable-env.md +``` + +- [ ] **Step 4: Update the Newton design doc** + +Edit `design/newton-backend-design.md`: +- In the "Completion Plan -> Done" list, add items 13/14: + - "13. Multi-env parallel via clone_arena_to (Target 4) — implemented." + - "14. DifferentiableEmbodiedEnv via Warp-tape autograd bridge (Target 5) — implemented." +- In "Remaining", remove items 7 (rigid-only Newton gym smoke tests) and + 8 (gradient rollout wrapper + smoke test) since both are now covered. +- Append a "References" section line: "Implementation plan: + `docs/superpowers/plans/2026-06-22-newton-backend-pr.md`." + +- [ ] **Step 5: Commit** + +```bash +git add agent_context/topics/differentiable-env.md \ + agent_context/MAP.yaml \ + design/newton-backend-design.md +git commit -m "docs: Newton multi-env + DifferentiableEmbodiedEnv + +agent_context routing for the new differentiable-env topic, plus an +update to design/newton-backend-design.md marking Targets 4 and 5 +done with a link to the implementation plan." +``` + +--- + +## Task 6: Branch cleanup and full test run + +**Files:** none (git operations + verification only). + +- [ ] **Step 1: Run the full Newton + diff suite** + +Run: +```bash +pytest -q \ + tests/sim/test_backend_parity.py \ + tests/sim/test_newton_finalize_lifecycle.py \ + tests/sim/test_newton_multi_env.py \ + tests/sim/test_differentiable_stepper.py \ + tests/sim/test_physics_attrs.py \ + tests/sim/test_sim_manager_cfg.py \ + tests/sim/objects/test_rigid_object.py \ + tests/sim/objects/test_articulation.py::TestArticulationNewton \ + tests/sim/objects/test_robot.py::TestRobotNewton \ + tests/gym/envs/test_differentiable_env.py +``` +Expected: all PASS (or GPU-marked tests SKIPPED on a headless host). + +- [ ] **Step 2: Run pre-commit checks** + +Run the `/pre-commit-check` skill — black, headers, type annotations, +exports, docstrings. + +- [ ] **Step 3: Inspect the branch for `wip` commits to squash** + +Run: `git log --oneline main..HEAD | grep -i wip` + +If any remain, plan an interactive cleanup via `git rebase -i main` (the +existing CLAUDE.md disallows `-i`, so do this **manually** outside the +agent or skip squashing if maintainer prefers history-preserving merge). + +- [ ] **Step 4: Create the PR** + +Use the `/pr` skill. Title: `feat(sim): Newton physics backend with +multi-env and differentiable APG`. Body summary: + +- Multi-env on Newton via implicit `clone_arena_to` at finalize. +- New `embodichain.lab.sim.diff` package and `DifferentiableEmbodiedEnv` + for APG on the `semi_implicit` solver. +- Franka FR3 reach APG example task with a one-iter loss-reduction + smoke test. +- Docs: agent_context routing + updated `design/newton-backend-design.md`. + +Reference the design doc and this plan. + +--- + +## Self-Review + +**Spec coverage check (against the revised 6-task plan):** + +- §2 multi-env — **already complete** (verified during execution; existing + `spawn_rigid_object_entities` / `spawn_articulation_entities` prototype-then-clone + at spawn, Newton views accept multi-entity lists, `TestRigidObjectNewton` + with `NUM_ARENAS=2` passes). No task needed. +- §3 module layout (`diff/bridge.py`, `differentiable_env.py`, example) — Tasks 2, 3, 4. +- §3 `NewtonStepFunc` + `tape_context` — Task 2. +- §3 `DifferentiableEmbodiedEnv` validation + step pipeline — Task 3. +- §3 Franka APG example + smoke tests — Task 4. +- §4 manager delegators — Task 1. +- §5 risks: clone re-evaluation under mutation — **N/A** (no clone-at-finalize; + cloning happens once at spawn, before finalize). +- §6 deferred items — out of scope (no tasks, per spec). +- §7 PR shape and commit plan — Task 6. +- §8 test files — Tasks 1, 2, 3, 4. +- §9 acceptance criteria — Task 6. + +No gaps. + +**Placeholder scan:** + +- No "TBD" / "TODO" / "implement later" in step content. + +**Type/signature consistency:** + +- `NewtonStepFunc.apply(action, sim_state)` — Task 2 defines signature; + Task 3 calls with the same args. +- `_apply_action_kernel(action_wp, tape)` — Task 3 abstract method; + Task 4 implements with the same signature. +- `_read_outputs(final_state) -> dict` — Task 3 abstract method; Task 4 + returns the documented `_order` / `_grad_track` shape. + +No drift. diff --git a/docs/superpowers/specs/2026-06-21-newton-backend-pr-design.md b/docs/superpowers/specs/2026-06-21-newton-backend-pr-design.md new file mode 100644 index 00000000..2b93f9bd --- /dev/null +++ b/docs/superpowers/specs/2026-06-21-newton-backend-pr-design.md @@ -0,0 +1,444 @@ +# Newton Physics Backend PR — Design + +Date: 2026-06-21 +Branch: `feature/newton-physics-backend` +Companion: `design/newton-backend-design.md` (current-state record) + +## 1. PR Targets and Scope + +The PR has five targets: + +1. Integrate the Newton physics backend on top of `dexsim` (`/root/sources/dexsim`). +2. Implement `RigidObject`, `Articulation`, and `Robot` on Newton. +3. Refactor `embodichain/lab/sim/cfg.py` to support both backends, including + Newton solver configuration. +4. Support multiple-env parallel simulation on Newton. +5. Support a differentiable env for analytic policy gradient (APG), in the + style of `dexsim/python/dexsim/engine/newton_physics/differentiable_stepper.py`. + +### Status going into this design + +Targets 1, 2, 3 are **already complete on the branch** (see +`design/newton-backend-design.md`): + +- `PhysicsBackend` ABC + registry; `DefaultPhysicsBackend` / `NewtonPhysicsBackend`. +- `DefaultPhysicsCfg` / `NewtonPhysicsCfg` with full solver dispatch + (`mujoco_warp` / `xpbd` / `semi_implicit` / `featherstone` / `vbd`), + `requires_grad`, `broad_phase`, `visualizer_enabled`, + `NewtonCollisionAttributesCfg`. +- Newton `RigidObject`, `Articulation`, `Robot` with batch views, runtime + attribute mutation, per-link mass live push. +- Capability matrix pinned by `tests/sim/test_backend_parity.py`. +- Newton finalize/invalidate lifecycle owned by `NewtonPhysicsBackend`. + +Targets 4 and 5 are **outstanding** and are the focus of this design. +`cfg.py` is otherwise left alone — Phase 3b legacy-`PhysicalAttr` removal is +deferred. + +## 2. Target 4 — Multi-Env Parallel Simulation on Newton + +### Mechanism + +`dexsim` already exposes the primitive we need: +`arena_src.clone_arena_to(arena_i)`. The pattern (see +`/root/sources/dexsim/examples/python/physics/basic/hello_newton.py`) is: + +1. Build a source arena and populate it with rigid bodies / articulations. +2. Add `num_envs - 1` additional empty arenas. +3. Call `arena_src.clone_arena_to(arena_i)` for each. +4. Newton finalize then sees `num_envs` parallel bodies and builds a single + batched model. + +EmbodiChain already builds `num_envs` arenas in +`SimulationManager._build_multiple_arenas` but does not clone — the existing +default-backend pattern is to call `add_*` once per `arena_index`. We add the +clone path on Newton only. + +### User-facing API + +No new public API. The flow is: + +```python +sim_cfg = SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg(...), + num_envs=4, +) +sim = SimulationManager(sim_cfg) +sim.add_rigid_object(cube_cfg) # spawns into arena_0 (source) +sim.add_robot(robot_cfg) # spawns into arena_0 (source) +sim.finalize_newton_physics() # clones arena_0 -> 1..3, then finalizes +``` + +Spawning with `cfg.arena_index > 0` on Newton raises with the message +"Newton spawn must target the source arena (arena_index in {-1, 0}); per-env +clones are produced at finalize." `arena_index == -1` (global) and +`arena_index == 0` both route to arena_0 on Newton. + +### Implementation + +`NewtonPhysicsBackend` gains an `_arenas_cloned: bool` flag (init `False`). +`prepare()` is extended: + +```python +def prepare(self) -> None: + if self._is_finalized and self._lifecycle_state() == "READY": + return + self._clone_source_arena_if_needed() + # ... existing ensure_simulation_prepared_lazy + rebuild_newton_from_scene ... + +def _clone_source_arena_if_needed(self) -> None: + arenas = self._manager._arenas + if len(arenas) <= 1 or self._arenas_cloned: + return + source = arenas[0] + for arena in arenas[1:]: + if self._arena_is_empty(arena): + source.clone_arena_to(arena) + self._arenas_cloned = True +``` + +`invalidate()` resets `_arenas_cloned` **only when scene topology changes**. +The two cases: + +- Topology change (`add_*` / `remove_*`): the corresponding `SimulationManager` + paths already call `self.physics.invalidate()`; we extend that to also + clear `_arenas_cloned` so the next `prepare()` re-clones into the (possibly + new) arenas. Attribute writes (`set_mass`, pose setters) keep + `_arenas_cloned = True`. + +Spawn guards live in `embodichain/lab/sim/utility/sim_utils.py`. Each +`add_rigid_object` / `add_articulation` / `add_robot` Newton path adds a +single guard: + +```python +if _is_newton_backend_active() and cfg.arena_index > 0: + logger.log_error( + "Newton spawn must target the source arena " + "(arena_index in {-1, 0}); per-env clones are produced at finalize." + ) +``` + +### Object Backend Views — Multi-Env Body-ID Resolution + +`NewtonRigidBodyView` and `NewtonArticulationView` (`embodichain/lab/sim/ +objects/backends/newton.py`) currently lazy-resolve a single body ID per +entity. We extend the resolver to return a `[num_envs]` index tensor after +finalize: + +- Each `RigidObject` / `Articulation` records its `entity_name` from arena_0. +- After clone, dexsim produces parallel entities in arenas 1..N-1 with + predictable per-arena names (the `clone_arena_to` namespacing scheme). +- The Newton view queries the finalized model's body/articulation registry + for every name variant and assembles the `[num_envs]` tensor on the + configured device. +- Pre-finalize fallback returns the scalar arena_0 ID, matching today's + BUILDER-state behavior. The scalar path is kept for code that runs before + finalize. + +All batched accessors (`get_body_state`, `set_local_pose`, +`apply_force_torque`, ...) already accept `env_ids` and operate on +`[num_envs, ...]` tensors; the only changes are in the view's +`_resolve_body_ids` and a `_num_envs` field plumbed in from the manager. + +### Default Backend + +Default backend behavior is **unchanged**. The existing pattern (one `add_*` +per `arena_index`) continues to work. Source-arena cloning on the default +backend is deferred. + +### Tests + +`tests/sim/test_newton_multi_env.py` (new): + +- Spawn a dynamic cube and a Franka URDF into arena_0 with `num_envs=4`. +- Finalize; verify rigid-object and articulation `get_body_state` return + shape `[4, ...]` with positions offset by the arena grid spacing. +- Step 10 substeps; verify per-env states diverge under per-env force + application. +- Spawn with `arena_index=1` raises. +- Mutating an attribute (`set_mass`) does not trigger re-clone (assert + `_arenas_cloned` stays `True`); adding a new asset does (assert it + becomes `False`). + +## 3. Target 5 — DifferentiableEmbodiedEnv (APG) + +### Reference Pattern + +`/root/sources/analytic_policy_gradients/envs/franka_reach_env.py` shows +the bridge pattern: a `torch.autograd.Function` (`_NewtonStepFunc`) opens a +`wp.Tape()`, launches Warp kernels in the forward, saves the tape, and runs +`tape.backward()` in the backward to extract `action.grad`. The franka +example bypasses dynamics and takes the gradient through FK only (because +the Featherstone solver does not propagate gradients through control). + +EmbodiChain will take the dynamics-grad path using +`dexsim.engine.newton_physics.DifferentiableStepper` with the +`semi_implicit` solver — this is the configuration `requires_grad=True` +already requires (see `NewtonPhysicsCfg.to_dexsim_cfg`). The FK-only path +is deferred as a future `grad_mode="kinematic"` option. + +### Module Layout + +- `embodichain/lab/sim/diff/__init__.py` — public surface. +- `embodichain/lab/sim/diff/bridge.py` — `_NewtonStepFunc(torch.autograd.Function)`, + `differentiable_step(manager, action, substeps)` helper, `tape_context(manager)` + context manager. +- `embodichain/lab/gym/envs/differentiable_env.py` — `DifferentiableEmbodiedEnv` + subclass. +- `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` — example task. + +`SimulationManager` gains two thin delegators (default backend raises): + +```python +def create_differentiable_stepper(self): + return self.physics.newton_manager.create_differentiable_stepper() + +def create_gradient_rollout(self, *args, **kwargs): + return self.physics.newton_manager.create_gradient_rollout(*args, **kwargs) +``` + +### DifferentiableEmbodiedEnv Contract + +Construction validates the Newton requires-grad config: + +```python +if not isinstance(cfg.sim_cfg.physics_cfg, NewtonPhysicsCfg): + log_error("DifferentiableEmbodiedEnv requires Newton backend.") +if not cfg.sim_cfg.physics_cfg.requires_grad: + log_error("DifferentiableEmbodiedEnv requires requires_grad=True.") +# solver_type=='semi_implicit' is already enforced by NewtonPhysicsCfg.to_dexsim_cfg. +``` + +### Step Pipeline + +`step(action)` is overridden: + +```python +def step(self, action): + if not isinstance(action, torch.Tensor): + action = torch.as_tensor(action, dtype=torch.float32, device=self.device) + obs, reward, terminated, truncated, info = _NewtonStepFunc.apply( + action, self._sim_state_dict() + ) + # auto-reset done envs with torch.where, preserving gradient on live envs + ... + return obs, reward, terminated, truncated, info +``` + +Inside `_NewtonStepFunc.forward`: + +1. Open `wp.Tape()`. +2. Apply the action to drive targets via a Warp kernel (replaces direct + `set_drive_target` calls in `_step_action` where those calls are not + tape-recorded — to be confirmed during implementation; if dexsim's + drive setter already records into the tape, we skip the replacement). +3. Run `DifferentiableStepper.step(state_in, state_out, control, contacts, + dt)` for `sim_steps_per_control` substeps, swapping `state_in`/`state_out`. +4. Evaluate the observation and reward managers, reading `joint_q` / + `body_q` via `wp.to_torch` (zero-copy, autograd-aware). +5. Save `tape`, grad-tracked Warp arrays, and metadata in `ctx`. + +`_NewtonStepFunc.backward`: + +1. Copy upstream `grad_reward` / `grad_obs` into Warp tensors' `.grad`. +2. `ctx.tape.backward()`. +3. Return `wp.to_torch(action_wp.grad)` reshaped to the action shape. +4. `ctx.tape.zero()`. + +### Reward and Observation Functors + +Existing functors that read tensors via `wp.to_torch` or torch operations on +manager-provided state are autograd-compatible by construction (Warp tape +sees the kernel launches; torch ops just compose). Functors that detour +through CPU / NumPy break the graph and will be flagged. The audit is +scoped to the example task's needs; a full functor audit is out of scope. + +The constraint is documented in `agent_context/` (new topic +`differentiable-env`) so future functor authors know the rule. + +### Reset Path + +`reset()` is non-differentiable. Wrap in `torch.no_grad()`, detach any +tensors written into Warp state. Auto-reset on `done` follows the franka +example: compute `obs_after_step`, then where `done_mask`: +`obs = torch.where(done_mask.unsqueeze(-1), fresh_obs.detach(), obs)`. +Live envs keep their gradient connection to the upstream action. + +### Memory and Truncation + +Each tape records all substeps in a single env step. For long +`sim_steps_per_control` or large `num_envs`, GPU memory can grow quickly. +`DifferentiableEmbodiedEnv` accepts an optional `truncate_backward_at` +argument (default `None` = full env step). When set, the tape is split +into chunks of N substeps; chunk boundaries are detached. This is a knob, +not a default behavior change. + +### Example Task + +`embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` mirrors the +APG reference env but is built on EmbodiChain primitives: + +- Franka FR3 URDF spawned via `add_robot` into arena_0. +- `num_envs = 4`, `NewtonPhysicsCfg(requires_grad=True, solver_cfg={"solver_type":"semi_implicit"})`. +- Observation: joint positions + EE pose + target pose + last action. +- Reward: position+orientation tracking matching the reference env. +- `DifferentiableEmbodiedEnv` subclass overriding only task-specific + reward/obs construction. + +### Tests + +`tests/gym/envs/test_differentiable_env.py`: + +1. Constructing `DifferentiableEmbodiedEnv` with `requires_grad=False` + raises. +2. `obs`/`reward` returned from `step(action)` have `requires_grad=True` + and a non-None `grad_fn`. +3. `loss = reward.sum(); loss.backward()` produces `action.grad` of + shape `[num_envs, action_dim]` with finite, non-zero values. +4. Finite-difference parity: per-env autograd gradient matches a + two-sided finite-difference estimate within tolerance on a 2-step + rollout (loose tolerance — `rtol=1e-1, atol=1e-2` — since the + semi-implicit solver is not a smooth function of action). +5. One APG iteration reduces the smoke loss. + +`tests/sim/test_differentiable_stepper.py`: + +1. `manager.create_differentiable_stepper()` raises on default backend. +2. On Newton with `requires_grad=False`, raises with a clear message. +3. On Newton with `requires_grad=True`, one `step()` produces tape-recorded + buffers and `tape.backward()` is callable. + +## 4. SimulationManager and Backend Changes Summary + +``` +embodichain/lab/sim/physics/newton.py + NewtonPhysicsBackend + + _arenas_cloned: bool + + _clone_source_arena_if_needed() + + _arena_is_empty(arena) + ~ prepare() (call clone helper before rebuild) + ~ invalidate() (clear _arenas_cloned only on topology change) + +embodichain/lab/sim/sim_manager.py + + create_differentiable_stepper() (delegates to NewtonManager) + + create_gradient_rollout(*a, **kw) (delegates to NewtonManager) + ~ add_rigid_object / add_articulation / add_robot + also invalidate -> reset _arenas_cloned (already invalidates; + additional flag-reset wired through invalidate()) + +embodichain/lab/sim/objects/backends/newton.py + NewtonRigidBodyView, NewtonArticulationView + + _num_envs + ~ _resolve_body_ids -> returns [num_envs] tensor after finalize + (no signature changes on public methods) + +embodichain/lab/sim/utility/sim_utils.py + + arena_index>0 guard on Newton spawn paths + +embodichain/lab/sim/diff/ (new package) + bridge.py + _NewtonStepFunc(torch.autograd.Function) + differentiable_step(manager, action, substeps) + tape_context(manager) + __init__.py + re-exports + +embodichain/lab/gym/envs/differentiable_env.py (new) + DifferentiableEmbodiedEnv(EmbodiedEnv) + +embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py (new) + FrankaReachApgTask +``` + +`cfg.py` is **unchanged**. + +## 5. Risks + +1. **`clone_arena_to` semantics under post-finalize mutation.** Cloning runs + at finalize. Attribute writes don't trigger re-clone; topology changes + do (via `invalidate()` clearing `_arenas_cloned`). If a user mutates + the source arena's *children list* without going through `add_*` / + `remove_*` (raw dexsim calls), the clone state goes stale. We document + the contract; we do not attempt to detect raw mutations. +2. **Drive-target write inside Tape.** `_step_action` currently writes + joint drive targets via dexsim setters. Verify during implementation + whether those writes are Warp-tape-recorded. If not, the differentiable + path replaces them with a Warp kernel that writes into + `control.joint_target` directly. Decision point at implementation; no + user-facing impact either way. +3. **Tape memory.** Long rollouts × large `num_envs` × full backward can + exhaust GPU memory. `truncate_backward_at` mitigates; we document + recommended values for the example task. +4. **Functor autograd compatibility is opt-in per functor.** No mass + refactor — only the functors needed by the example task are audited. + The contract is documented in `agent_context/` so future authors know + when to use torch ops vs. NumPy detours. +5. **Upstream dexsim.** Continues to depend on `yueci/adapt-embodichain` + for active-joint indexing (existing risk; not introduced here). No new + upstream dependencies — `clone_arena_to`, `DifferentiableStepper`, + and `GradientRollout` are all on dexsim main paths. +6. **Body-ID resolution after clone.** The view extension assumes a + predictable per-arena naming scheme from `clone_arena_to`. We verify + the actual naming pattern during implementation and adjust the + resolver accordingly; if `clone_arena_to` does not namespace bodies + per-arena in a way EmbodiChain can rebuild, fall back to maintaining + parallel entity lists per-`RigidObject` / `Articulation`. + +## 6. Out of Scope (Deferred) + +These targets are intentionally not in this PR: + +- Default-backend cloning via `clone_arena_to`. +- Soft / cloth objects on Newton. +- `RigidObjectGroup` on Newton. +- FK-only differentiable mode (`grad_mode="kinematic"`). +- Per-link Newton-native contact params on articulations (waiting on + dexsim per-link shape-material setter). +- Phase 3b legacy-`PhysicalAttr` removal. +- Functor-wide autograd audit. + +## 7. PR Shape + +Single feature branch `feature/newton-physics-backend`. Commit plan +(after squashing the existing `wip` commits): + +1. `feat(sim/newton): clone source arena at finalize for multi-env` + - `NewtonPhysicsBackend._clone_source_arena_if_needed` + - View multi-env body-id resolution + - Spawn guards for `arena_index>0` on Newton + - `tests/sim/test_newton_multi_env.py` +2. `feat(sim/diff): NewtonStepFunc bridge for Warp tape -> torch autograd` + - `embodichain/lab/sim/diff/` package + - `SimulationManager.create_differentiable_stepper` / + `create_gradient_rollout` delegators + - `tests/sim/test_differentiable_stepper.py` +3. `feat(gym): DifferentiableEmbodiedEnv for APG` + - `embodichain/lab/gym/envs/differentiable_env.py` + - `tests/gym/envs/test_differentiable_env.py` +4. `feat(tasks): Franka reach APG example task` + - `embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py` +5. `docs(newton): update backend design doc + agent_context routing` + - update `design/newton-backend-design.md` Done/Remaining lists + - new `agent_context/` topic `differentiable-env` + +## 8. Tests Summary + +| File | Coverage | +|------|----------| +| `tests/sim/test_newton_multi_env.py` (new) | clone-at-finalize, batched body IDs, attribute mutation does not re-clone, topology change does, arena_index>0 spawn guard | +| `tests/sim/test_differentiable_stepper.py` (new) | manager delegator behavior on each backend; tape recording smoke | +| `tests/gym/envs/test_differentiable_env.py` (new) | construction validation, requires_grad on outputs, backward yields non-zero action.grad, finite-difference parity, one-iter loss reduction | +| `tests/sim/test_backend_parity.py` (existing) | unchanged — multi-env on Newton does not change capability flags | +| `tests/sim/test_newton_finalize_lifecycle.py` (existing) | extend with a multi-env case and an APG-config case | + +## 9. Acceptance Criteria + +This PR is ready to merge when: + +- All targets 1–5 are reflected in code or in this design's deferred list. +- `pytest -q tests/sim tests/gym/envs/test_differentiable_env.py` is green. +- The Franka APG example task runs `python -m embodichain.lab.scripts.run_env + --task=FrankaReachApg-v0 --num-envs=4 --steps=50` and loss decreases. +- `design/newton-backend-design.md` is updated to mark Targets 4 and 5 + Done and to point at this spec for the implementation rationale. +- The `wip` commits on the branch are squashed. 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/differentiable_env.py b/embodichain/lab/gym/envs/differentiable_env.py new file mode 100644 index 00000000..829b1573 --- /dev/null +++ b/embodichain/lab/gym/envs/differentiable_env.py @@ -0,0 +1,176 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Differentiable Newton-backed EmbodiedEnv for analytic policy gradient. + +Wraps the standard :class:`EmbodiedEnv` step pipeline in a Warp tape and +bridges autograd into PyTorch via +:class:`embodichain.lab.sim.diff.NewtonStepFunc`. Subclasses define how +actions become Newton control writes and how observations/rewards are +read from the post-step state; the bridge handles the tape lifecycle +and the backward pass. + +Usage: + + class MyTask(DifferentiableEmbodiedEnv): + def _apply_action_kernel(self, action_wp, tape): ... + def _read_outputs(self, final_state) -> dict: ... +""" + +from __future__ import annotations + +from typing import Any, Callable + +import torch + +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnv, EmbodiedEnvCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg +from embodichain.lab.sim.diff import NewtonStepFunc +from embodichain.utils import logger + +__all__ = ["DifferentiableEmbodiedEnv"] + + +class DifferentiableEmbodiedEnv(EmbodiedEnv): + """EmbodiedEnv variant that exposes APG-ready :py:meth:`step`. + + Subclasses must implement :meth:`_apply_action_kernel` and + :meth:`_read_outputs`; the rest of the EmbodiedEnv contract (reset, + observation managers, reward functors) carries over. + """ + + def __init__(self, cfg: EmbodiedEnvCfg, *args, **kwargs) -> None: + self._validate_diff_cfg(cfg) + super().__init__(cfg, *args, **kwargs) + self._truncate_backward_at: int | None = getattr( + cfg, "truncate_backward_at", None + ) + + @staticmethod + def _validate_diff_cfg(cfg: EmbodiedEnvCfg) -> None: + physics_cfg = cfg.sim_cfg.physics_cfg + if not isinstance(physics_cfg, NewtonPhysicsCfg): + logger.log_error( + "DifferentiableEmbodiedEnv requires NewtonPhysicsCfg, " + f"got {type(physics_cfg).__name__}." + ) + if not physics_cfg.requires_grad: + logger.log_error( + "DifferentiableEmbodiedEnv requires requires_grad=True on " + "the NewtonPhysicsCfg." + ) + + # -- subclass contract ------------------------------------------------ # + + def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: + """Inside the open Warp tape, write the action into Newton control. + + Implementations launch a Warp kernel that reads ``action_wp`` + (a ``wp.array(dtype=wp.float32, requires_grad=True)`` of shape + ``[num_envs * action_dim]``) and writes into + ``self.sim.physics.newton_manager._control`` so the next stepper + call uses the new control. + """ + raise NotImplementedError( + "Subclasses of DifferentiableEmbodiedEnv must implement " + "_apply_action_kernel(action_wp, tape)." + ) + + def _read_outputs(self, final_state: Any) -> dict: + """Read the post-step observation and reward as torch tensors. + + Must return a dict with keys ``"obs"``, ``"reward"``, + ``"terminated"``, ``"truncated"``, plus the ``_order`` and + ``_grad_track`` metadata expected by + :class:`NewtonStepFunc`. ``obs`` and ``reward`` should be torch + tensors backed by ``wp.to_torch`` of grad-tracked Warp arrays. + """ + raise NotImplementedError( + "Subclasses of DifferentiableEmbodiedEnv must implement " + "_read_outputs(final_state)." + ) + + def _make_step_fn(self) -> Callable[[], Any]: + """Return a callable that advances the sim inside the open tape. + + The returned callable takes no arguments and returns the final + Newton :class:`State` after stepping. It is invoked by + :class:`NewtonStepFunc` inside the ``with tape:`` block, so any + Warp kernel launches (or differentiable Newton calls like + ``eval_fk``) are recorded on the tape. + + The default implementation runs the differentiable + :class:`DifferentiableStepper` for ``sim_steps_per_control`` + substeps. Subclasses can override this to swap in an FK-only + differentiable path (bypassing the dynamics solver when it does + not propagate grad through control inputs) or any other + tape-tracked stepping strategy. + """ + manager = self.sim + substeps = self.cfg.sim_steps_per_control + nm = manager.physics.newton_manager + stepper = manager.create_differentiable_stepper() + state_in = nm._state_0 + state_out = nm._model.state() + contacts = stepper.create_contacts() + dt_val = nm.solver_dt + + def _step(): + for _ in range(substeps): + stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) + state_in, state_out = state_out, state_in + return state_in + + return _step + + # -- gym surface ------------------------------------------------------ # + + def step(self, action: torch.Tensor): + if not isinstance(action, torch.Tensor): + action = torch.as_tensor(action, dtype=torch.float32) + sim_state = self._build_sim_state_dict(action) + outputs = NewtonStepFunc.apply(action, sim_state) + obs, reward, terminated, truncated = outputs[:4] + info = sim_state["last_info"] + + done_mask = terminated | truncated + if done_mask.any(): + reset_ids = done_mask.nonzero(as_tuple=False).squeeze(-1) + fresh_obs, _ = self.reset(options={"reset_ids": reset_ids}) + obs = torch.where( + done_mask.unsqueeze(-1).expand_as(obs), + fresh_obs.detach(), + obs, + ) + return obs, reward, terminated, truncated, info + + def _build_sim_state_dict(self, action: torch.Tensor) -> dict: + return { + "manager": self.sim, + "substeps": self.cfg.sim_steps_per_control, + "action_to_control_kernel": self._wrap_action_kernel(), + "kernel_args": (), + "obs_reward_fn": self._read_outputs, + "step_fn": self._make_step_fn(), + "last_info": {}, + } + + def _wrap_action_kernel(self): + env = self + + def _inner(action_wp, *_): + env._apply_action_kernel(action_wp, tape=None) + + return _inner 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/envs/tasks/special/franka_reach_apg.py b/embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py new file mode 100644 index 00000000..71220213 --- /dev/null +++ b/embodichain/lab/gym/envs/tasks/special/franka_reach_apg.py @@ -0,0 +1,494 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Franka FR3 reach task with differentiable Newton physics (APG). + +Built on :class:`DifferentiableEmbodiedEnv`. The Warp-tape bridge +produces ``action.grad`` that flows back through a differentiable +forward-kinematics path (``newton.eval_fk``). The semi_implicit +solver does not propagate grad through ``joint_target_pos`` to +``body_q`` (the grad path is zero), so we bypass the dynamics +solver and run FK directly, matching the reference APG +implementation in +``/root/sources/analytic_policy_gradients/envs/franka_reach_env.py``. +""" + +from __future__ import annotations + +from typing import Any, Callable + +import numpy as np +import torch +import warp as wp +import newton +import newton.utils + +from embodichain.lab.gym.envs.differentiable_env import DifferentiableEmbodiedEnv +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg +from embodichain.lab.gym.utils.registration import register_env +from embodichain.lab.sim.cfg import ( + NewtonPhysicsCfg, + RobotCfg, + URDFCfg, +) +from embodichain.lab.sim.sim_manager import SimulationManagerCfg + +__all__ = ["FrankaReachApgEnv"] + +# Franka FR3 arm has 7 actuated arm joints; the URDF also has 2 finger +# joints (9 dof total). We only control the 7 arm joints. +FRANKA_NUM_ARM_JOINTS = 7 +FRANKA_EE_BODY = "fr3_hand_tcp" +DEFAULT_ACTION_SCALE = 0.2 +DEFAULT_MAX_EPISODE_STEPS = 30 +TARGET_POS_RANGE = { + "x": (0.05, 0.70), + "y": (-0.45, 0.45), + "z": (0.20, 0.95), +} + + +@wp.kernel +def _set_joint_targets_kernel( + action: wp.array(dtype=wp.float32), + current_q: wp.array(dtype=wp.float32), + target_q: wp.array(dtype=wp.float32), + limit_lo: wp.array(dtype=wp.float32), + limit_hi: wp.array(dtype=wp.float32), + action_scale: wp.float32, + n_joints_per_env: wp.int32, + n_arm: wp.int32, + total: wp.int32, +): + """Compute new joint q: target = clamp(current + action * scale, lo, hi).""" + tid = wp.tid() + if tid < total: + env_idx = tid / n_arm + j = tid % n_arm + off = env_idx * n_joints_per_env + j + new_q = current_q[off] + action[tid] * action_scale + target_q[off] = wp.clamp(new_q, limit_lo[j], limit_hi[j]) + + +@wp.kernel +def _reach_reward_kernel( + body_q: wp.array(dtype=wp.transformf), + ee_body_indices: wp.array(dtype=wp.int32), + target_pos: wp.array(dtype=wp.vec3f), + reward_out: wp.array(dtype=wp.float32), +): + """Position-only reach reward (smoke task): -0.2*dist + 0.1*exp(-dist^2/0.02).""" + env_idx = wp.tid() + ee_transform = body_q[ee_body_indices[env_idx]] + eef_pos = wp.transform_get_translation(ee_transform) + diff = eef_pos - target_pos[env_idx] + pos_dist = wp.sqrt(wp.dot(diff, diff) + wp.float32(1e-8)) + reward_out[env_idx] = wp.float32(-0.2) * pos_dist + wp.float32(0.1) * wp.exp( + -pos_dist * pos_dist / wp.float32(0.02) + ) + + +@register_env("FrankaReachApg-v0") +class FrankaReachApgEnv(DifferentiableEmbodiedEnv): + """Differentiable Franka FR3 reach task for analytic policy gradients. + + The environment resolves the Franka FR3 URDF via + ``newton.utils.download_asset("franka_emika_panda")`` (network-dependent) + or an explicit ``urdf_path`` kwarg override. The robot is added through + the standard EmbodiChain ``sim.add_robot(cfg.robot)`` flow driven by + :class:`EmbodiedEnv`/``BaseEnv.__init__``. + + The differentiable path is: + + action -> new_joint_q (action kernel) -> eval_fk -> body_q + -> reward kernel -> reward_wp -> tape.backward -> action.grad + + The dynamics solver (semi_implicit) is bypassed because it does not + propagate gradient through ``joint_target_pos`` to ``body_q`` (the + stiffness-driven grad path evaluates to zero in practice). This + matches the reference APG env's workaround. + """ + + metadata = {"render_modes": ["human"], "default_num_envs": 4} + + def __init__( + self, + cfg: EmbodiedEnvCfg | None = None, + *, + num_envs: int = 4, + urdf_path: str | None = None, + action_scale: float = DEFAULT_ACTION_SCALE, + max_episode_steps: int = DEFAULT_MAX_EPISODE_STEPS, + device: str = "cuda:0", + ) -> None: + self._urdf_path = urdf_path + self._action_scale = float(action_scale) + self._max_episode_steps = int(max_episode_steps) + self._device_str = device + + if cfg is None: + urdf = urdf_path or self._resolve_default_urdf() + robot_cfg = RobotCfg( + uid="franka", + urdf_cfg=URDFCfg().set_urdf(urdf), + fix_base=True, + ) + cfg = EmbodiedEnvCfg( + sim_cfg=SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + device=device, + requires_grad=True, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=num_envs, + headless=True, + ), + robot=robot_cfg, + num_envs=num_envs, + max_episode_steps=max_episode_steps, + ) + # Bug 1 fix: cfg.robot is set BEFORE super().__init__() so that + # EmbodiedEnv._init_sim_state -> BaseEnv._setup_scene -> + # _setup_robot -> sim.add_robot(cfg.robot) has a valid robot to + # add. BaseEnv.__init__ also calls finalize_newton_physics() once + # the scene is built, so we do NOT re-finalize here. + super().__init__(cfg) + # EmbodiedEnv has added the robot and BaseEnv has finalized the + # Newton model. Cache joint-limit Warp arrays and EE body indices. + self._cache_franka_buffers() + self._init_targets() + + # -- scene setup ----------------------------------------------------- # + + def _resolve_default_urdf(self) -> str: + """Resolve the Franka URDF via Newton's asset cache. + + Raises: + FileNotFoundError: If the URDF cannot be downloaded or + located. + """ + try: + urdf = newton.utils.download_asset("franka_emika_panda") / ( + "urdf/fr3_franka_hand.urdf" + ) + if urdf.exists(): + return str(urdf) + except Exception: + pass + raise FileNotFoundError("Franka URDF not available; pass urdf_path explicitly.") + + def _cache_franka_buffers(self) -> None: + """Cache joint-limit Warp arrays, EE body indices, and FK state.""" + nm = self.sim.physics.newton_manager + model = nm._model + # Warp's ``wp.zeros`` / ``wp.launch`` reject ``torch.device`` + # directly (``Invalid device identifier: cuda:0``), so cache the + # Warp-compatible device string up-front. + self._wp_device = model.device + # ``model.joint_limit_lower`` is a ``wp.array``; convert via + # ``.numpy()`` before slicing (``np.asarray`` on a wp.array slice + # raises "Item indexing is not supported on wp.array objects"). + lo = model.joint_limit_lower.numpy()[:FRANKA_NUM_ARM_JOINTS].astype(np.float32) + hi = model.joint_limit_upper.numpy()[:FRANKA_NUM_ARM_JOINTS].astype(np.float32) + self._limit_lo_t = torch.from_numpy(lo).to(self.device) + self._limit_hi_t = torch.from_numpy(hi).to(self.device) + self._limit_lo_wp = wp.array(lo, dtype=wp.float32, device=self._wp_device) + self._limit_hi_wp = wp.array(hi, dtype=wp.float32, device=self._wp_device) + self._n_joints_per_env = int(len(model.joint_q) // self.sim.num_envs) + # Fresh FK state; reused across step calls (eval_fk overwrites it). + self._fk_state = model.state() + self._new_joint_q: wp.array | None = None + # Per-env global EE body indices into the flat body_q array. + self._ee_global_idx = self._compute_ee_body_indices() + self._ee_idx_wp = wp.array( + np.asarray(self._ee_global_idx, dtype=np.int32), + dtype=wp.int32, + device=self._wp_device, + ) + self._ee_idx_t = torch.tensor( + self._ee_global_idx, dtype=torch.long, device=self.device + ) + + def _compute_ee_body_indices(self) -> list[int]: + """Scan model.body_label for the EE body per env. + + Each cloned arena produces a full set of Franka bodies in the + shared Newton model. We pick the ``FRANKA_EE_BODY`` body for + each env block (one global index per env). + """ + nm = self.sim.physics.newton_manager + model = nm._model + n_envs = self.sim.num_envs + n_per_env = len(model.body_label) // n_envs + idx_per_env: list[int] = [] + for i in range(n_envs): + for j in range(n_per_env): + global_idx = i * n_per_env + j + if FRANKA_EE_BODY in str(model.body_label[global_idx]): + idx_per_env.append(global_idx) + break + if len(idx_per_env) != n_envs: + raise RuntimeError( + f"Expected {n_envs} '{FRANKA_EE_BODY}' bodies, " + f"found {len(idx_per_env)}." + ) + return idx_per_env + + def _init_targets(self) -> None: + n = self.sim.num_envs + device = self.device + self.target_pos = torch.zeros(n, 3, device=device) + self.target_quat = torch.zeros(n, 4, device=device) + self.last_action = torch.zeros(n, FRANKA_NUM_ARM_JOINTS, device=device) + self.step_count = torch.zeros(n, dtype=torch.int32, device=device) + self._sample_new_targets(torch.arange(n, device=device)) + + def _sample_new_targets(self, env_ids: torch.Tensor) -> None: + n = env_ids.numel() + d = self.device + self.target_pos[env_ids, 0] = TARGET_POS_RANGE["x"][0] + torch.rand( + n, device=d + ) * (TARGET_POS_RANGE["x"][1] - TARGET_POS_RANGE["x"][0]) + self.target_pos[env_ids, 1] = TARGET_POS_RANGE["y"][0] + torch.rand( + n, device=d + ) * (TARGET_POS_RANGE["y"][1] - TARGET_POS_RANGE["y"][0]) + self.target_pos[env_ids, 2] = TARGET_POS_RANGE["z"][0] + torch.rand( + n, device=d + ) * (TARGET_POS_RANGE["z"][1] - TARGET_POS_RANGE["z"][0]) + # Identity orientation: the smoke task uses position-only reward. + self.target_quat[env_ids] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=d).expand( + n, -1 + ) + + # -- DifferentiableEmbodiedEnv contract ------------------------------ # + + def _make_step_fn(self) -> Callable[[], Any]: + """FK bypass: compute body_q from new_joint_q via newton.eval_fk. + + The semi_implicit solver does not propagate grad through + ``joint_target_pos`` to ``body_q`` (the grad path is zero), so + we bypass the dynamics solver and run forward kinematics + directly inside the tape. ``self._new_joint_q`` is populated by + :meth:`_apply_action_kernel` before this callable runs. + """ + env = self + model = env.sim.physics.newton_manager._model + + def _step(): + newton.eval_fk( + model, + env._new_joint_q, + env._fk_state.joint_qd, + env._fk_state, + ) + return env._fk_state + + return _step + + def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: + """Launch the action-to-control kernel inside the open tape. + + Writes ``new_joint_q = clamp(current_q + action * scale, lo, hi)`` + into a freshly allocated ``self._new_joint_q`` Warp array. The + FK step function then consumes this array via ``newton.eval_fk``. + """ + nm = self.sim.physics.newton_manager + n = self.sim.num_envs + total = n * FRANKA_NUM_ARM_JOINTS + # Allocate a fresh new_joint_q each call so each forward pass + # has its own grad graph (the tape records the kernel writes). + self._new_joint_q = wp.zeros( + n * self._n_joints_per_env, + dtype=wp.float32, + device=self._wp_device, + requires_grad=True, + ) + wp.launch( + _set_joint_targets_kernel, + dim=total, + inputs=[ + action_wp, + nm._state_0.joint_q, + self._new_joint_q, + self._limit_lo_wp, + self._limit_hi_wp, + wp.float32(self._action_scale), + wp.int32(self._n_joints_per_env), + wp.int32(FRANKA_NUM_ARM_JOINTS), + wp.int32(total), + ], + device=self._wp_device, + ) + + def _read_outputs(self, final_state: Any) -> dict: + """Launch the reward kernel and build obs INSIDE the open tape. + + Reward is written into a grad-tracked ``reward_wp`` Warp array, + then exposed as a torch tensor via ``wp.to_torch`` (zero-copy). + The obs is built from ``wp.to_torch(final_state.joint_q)`` and + ``wp.to_torch(final_state.body_q)`` (also tape-tracked). + """ + n = self.sim.num_envs + device = self._wp_device + + # Grad-tracked reward output array. The kernel launches inside + # the open tape so reward_wp carries gradient back through the + # reward kernel -> body_q -> FK -> new_joint_q -> action_wp. + reward_wp = wp.zeros(n, dtype=wp.float32, device=device, requires_grad=True) + target_pos_wp = wp.from_torch( + self.target_pos.detach().clone().contiguous(), dtype=wp.vec3 + ) + wp.launch( + _reach_reward_kernel, + dim=n, + inputs=[final_state.body_q, self._ee_idx_wp, target_pos_wp], + outputs=[reward_wp], + device=device, + ) + + joint_q_t = wp.to_torch(final_state.joint_q).view(n, -1) + body_q_flat = wp.to_torch(final_state.body_q).view(-1, 7) + ee_pose = body_q_flat[self._ee_idx_t] + obs = torch.cat( + [ + joint_q_t[:, :FRANKA_NUM_ARM_JOINTS], + ee_pose, + self.target_pos, + self.target_quat, + self.last_action, + ], + dim=-1, + ) + + reward_t = wp.to_torch(reward_wp) + pos_dist = (ee_pose[:, :3] - self.target_pos).norm(dim=-1).detach() + terminated = pos_dist < 0.01 + truncated = self.step_count >= self._max_episode_steps + + return { + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": { + "obs": None, + "reward": reward_wp, + "terminated": None, + "truncated": None, + }, + "obs": obs, + "reward": reward_t, + "terminated": terminated, + "truncated": truncated, + } + + # -- gym overrides --------------------------------------------------- # + + def step(self, action: torch.Tensor): + """Step the env, then advance the cached joint_q for the next call. + + The parent :meth:`DifferentiableEmbodiedEnv.step` runs the + differentiable bridge. After it returns, we update + ``nm._state_0.joint_q`` (detached) for envs that were not + auto-reset so the next step starts from the new configuration. + """ + if not isinstance(action, torch.Tensor): + action = torch.as_tensor(action, dtype=torch.float32) + clamped_action = torch.clamp(action.to(self.device), -1.0, 1.0) + # Advance step_count BEFORE the bridge runs so _read_outputs + # computes truncated against the post-step value. + self.step_count += 1 + result = super().step(clamped_action) + obs, reward, terminated, truncated, info = result + done_mask = terminated | truncated + live = (~done_mask).nonzero(as_tuple=False).squeeze(-1) + if live.numel() > 0: + with torch.no_grad(): + nm = self.sim.physics.newton_manager + joint_q_t = wp.to_torch(nm._state_0.joint_q).view(self.sim.num_envs, -1) + cur = joint_q_t[live, :FRANKA_NUM_ARM_JOINTS] + delta = clamped_action[live].detach() * self._action_scale + lo = self._limit_lo_t.unsqueeze(0).expand_as(cur) + hi = self._limit_hi_t.unsqueeze(0).expand_as(cur) + joint_q_t[live, :FRANKA_NUM_ARM_JOINTS] = torch.clamp( + cur + delta, lo, hi + ) + self.last_action = clamped_action.detach().clone() + return obs, reward, terminated, truncated, info + + def reset( + self, + *, + seed: int | None = None, + options: dict | None = None, + ): + """Reset joint_q, targets, and step_count for the touched envs. + + Args: + seed: Optional RNG seed for deterministic resets. + options: Optional dict; supports ``{"reset_ids": }`` + for partial resets (used by auto-reset in + :meth:`DifferentiableEmbodiedEnv.step`). + + Returns: + Tuple of ``(obs, info)``. + """ + if seed is not None: + torch.manual_seed(seed) + if options is None: + options = {} + reset_ids = options.get("reset_ids") + if reset_ids is None: + env_ids = torch.arange(self.sim.num_envs, device=self.device) + else: + env_ids = torch.as_tensor(reset_ids, dtype=torch.long, device=self.device) + with torch.no_grad(): + self.step_count[env_ids] = 0 + self.last_action[env_ids] = 0.0 + self._sample_new_targets(env_ids) + nm = self.sim.physics.newton_manager + joint_q_t = wp.to_torch(nm._state_0.joint_q).view(self.sim.num_envs, -1) + joint_q_t[env_ids] = 0.0 + newton.eval_fk( + nm._model, + nm._state_0.joint_q, + nm._state_0.joint_qd, + nm._state_0, + ) + obs = self._initial_obs() + return obs, {} + + def _initial_obs(self) -> torch.Tensor: + """Compute the initial obs from state_0 (no grad, no side effects).""" + with torch.no_grad(): + nm = self.sim.physics.newton_manager + state = nm._state_0 + n = self.sim.num_envs + joint_q_t = wp.to_torch(state.joint_q).view(n, -1) + body_q_flat = wp.to_torch(state.body_q).view(-1, 7) + ee_pose = body_q_flat[self._ee_idx_t] + obs = torch.cat( + [ + joint_q_t[:, :FRANKA_NUM_ARM_JOINTS], + ee_pose, + self.target_pos, + self.target_quat, + self.last_action, + ], + dim=-1, + ) + return obs.detach() + + def close(self) -> None: + """Close the environment and release resources.""" + self.sim.destroy() 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..882fc422 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,161 @@ 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, + ), + sync_to_dexsim=True + ) + 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 +397,165 @@ class WindowRecordCfg: """Video file prefix used when no explicit save path is provided.""" +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." + ) + + +def validate_physics_cfg(physics_cfg: PhysicsCfg) -> None: + """Validate that ``physics_cfg`` is a supported backend configuration.""" + physics_backend_from_cfg(physics_cfg) + + @configclass -class GPUMemoryCfg: - """A gpu memory configuration dataclass that neatly holds all parameters that configure physics GPU memory for simulation""" +class NewtonCollisionAttributesCfg: + """Newton-specific per-shape collision/contact attributes. + + Mirrors :class:`dexsim.spawn.descs.NewtonCollisionDesc` (which in turn + mirrors ``newton.ModelBuilder.ShapeConfig``), so the resolver can overlay + these fields by name. All fields default to ``None`` meaning "keep the + Newton backend default". + + The backend-neutral quantities (sliding friction, restitution, + enable-collision) live on :class:`RigidBodyAttributesCfg` and are projected + onto the Newton ``mu`` / ``restitution`` / ``has_shape_collision`` shape + knobs by the resolver; they are NOT repeated here. + """ - temp_buffer_capacity: int = 2**24 - """Increase this if you get 'PxgPinnedHostLinearMemoryAllocator: overflowing initial allocation size, increase capacity to at least %.' """ + # -- Contact-material fields (per-solver subset, see NEWTON_CONTACT_SOLVER_FIELDS) -- + ke: float | None = None + """Contact stiffness for compliant contacts.""" + kd: float | None = None + """Contact damping for compliant contacts.""" + kf: float | None = None + """Friction stiffness for compliant contacts.""" + ka: float | None = None + """Adhesion stiffness for compliant contacts.""" + kh: float | None = None + """Hydroelastic stiffness scale.""" + mu_torsional: float | None = None + """Torsional friction coefficient.""" + mu_rolling: float | None = None + """Rolling friction coefficient.""" + + # -- Solver-agnostic shape-config fields -- + margin: float | None = None + """Contact margin (shapes within this distance are considered in contact).""" + gap: float | None = None + """Contact gap (rest distance between shapes).""" + is_solid: bool | None = None + """Whether the shape is solid (vs. hollow) for mass computation.""" + collision_group: int | None = None + """Collision group id used by the broad-phase filter.""" + collision_filter_parent: bool | None = None + """Whether to filter collisions with the parent body.""" + has_particle_collision: bool | None = None + """Whether the shape collides with particles.""" + is_visible: bool | None = None + """Whether the shape is visible to the Newton visualizer.""" + is_site: bool | None = None + """Whether the shape is registered as a Newton site.""" + is_hydroelastic: bool | None = None + """Whether to use hydroelastic contact for this shape.""" + + # -- SDF (signed distance field) collision params -- + sdf_narrow_band_range: tuple[float, float] | None = None + """Narrow-band range [inner, outer] for SDF collision.""" + sdf_target_voxel_size: float | None = None + """Target voxel size for SDF generation.""" + sdf_max_resolution: int | None = None + """Maximum grid resolution for SDF generation.""" + sdf_texture_format: str | None = None + """Texture format for SDF collision.""" - max_rigid_contact_count: int = 2**19 - """Increase this if you get 'Contact buffer overflow detected'""" + @classmethod + def from_dict(cls, init_dict: Dict[str, Any]) -> NewtonCollisionAttributesCfg: + """Initialize the configuration from a dictionary.""" + cfg = cls() + for key, value in init_dict.items(): + if hasattr(cfg, key): + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + return cfg - 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 to_newton_collision_desc(self): + """Build a :class:`dexsim.spawn.descs.NewtonCollisionDesc` from this cfg.""" + from dexsim.spawn.descs import NewtonCollisionDesc + + return NewtonCollisionDesc( + **{ + f: getattr(self, f) + for f in ( + "ke", + "kd", + "kf", + "ka", + "kh", + "mu_torsional", + "mu_rolling", + "margin", + "gap", + "is_solid", + "collision_group", + "collision_filter_parent", + "has_particle_collision", + "is_visible", + "is_site", + "is_hydroelastic", + "sdf_narrow_band_range", + "sdf_target_voxel_size", + "sdf_max_resolution", + "sdf_texture_format", + ) + } + ) - 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 _merge_newton_subcfg( + override: NewtonCollisionAttributesCfg | None, + base: NewtonCollisionAttributesCfg | None, +) -> NewtonCollisionAttributesCfg | None: + """Merge a Newton sub-config override onto a base. + + For each Newton field, the override's non-None value wins, else the base's. + Returns ``None`` if neither side sets any field. + """ + if override is None: + return base + if base is None: + return override + merged = NewtonCollisionAttributesCfg() + any_set = False + for field_name in merged.__dataclass_fields__: + if field_name == "newton": + continue + ov = getattr(override, field_name) + val = ov if ov is not None else getattr(base, field_name) + setattr(merged, field_name, val) + if val is not None: + any_set = True + return merged if any_set else None @configclass @@ -227,6 +566,11 @@ class RigidBodyAttributesCfg: 1. The dynamic properties, such as mass, damping, etc. 2. The collision properties. 3. The physics material properties. + + The ``newton`` sub-config carries Newton-specific per-shape contact/shape + knobs (``ke``/``kd``/``margin``/...) that have no PhysX equivalent; it is + ignored on the default backend and applied via the Newton desc-native + registration path when set. """ mass: float = 1.0 @@ -285,8 +629,17 @@ class RigidBodyAttributesCfg: static_friction: float = 0.5 """Static friction coefficient.""" + newton: NewtonCollisionAttributesCfg | None = None + """Newton-specific per-shape contact/shape attributes (ignored on default backend).""" + def attr(self) -> PhysicalAttr: - """Convert to dexsim PhysicalAttr""" + """Convert to dexsim PhysicalAttr. + + This is the legacy PhysX-oriented projection used by the default + backend. Newton-native fields (``self.newton``) are not representable + here; the Newton path uses + :func:`embodichain.lab.sim.physics_attrs.resolve_newton_shape` instead. + """ attr = PhysicalAttr() attr.mass = self.mass attr.contact_offset = self.contact_offset @@ -310,7 +663,9 @@ def from_dict( """Initialize the configuration from a dictionary.""" cfg = cls() for key, value in init_dict.items(): - if hasattr(cfg, key): + if key == "newton" and isinstance(value, dict): + setattr(cfg, key, NewtonCollisionAttributesCfg.from_dict(value)) + elif hasattr(cfg, key): setattr(cfg, key, value) else: logger.log_warning( @@ -345,16 +700,38 @@ class RigidBodyAttributesOverrideCfg: dynamic_friction: float | None = None static_friction: float | None = None + newton: NewtonCollisionAttributesCfg | None = None + """Newton-specific per-shape overrides (None means inherit the base newton sub-config).""" + def merge_with(self, base: RigidBodyAttributesCfg) -> PhysicalAttr: - """Build a :class:`~dexsim.types.PhysicalAttr` from base values and overrides.""" + """Build a :class:`~dexsim.types.PhysicalAttr` from base values and overrides. + + .. note:: + This returns the legacy PhysX projection and therefore drops the + Newton sub-config. For a Newton-aware merge that preserves + ``newton``, use :meth:`merged_cfg` and pass it to the Newton + resolver. + """ + return self.merged_cfg(base).attr() + + def merged_cfg(self, base: RigidBodyAttributesCfg) -> RigidBodyAttributesCfg: + """Merge overrides onto ``base`` into a full :class:`RigidBodyAttributesCfg`. + + Unlike :meth:`merge_with`, this preserves the ``newton`` sub-config + (override's non-None sub-fields win, else base's) so the result can be + fed to the Newton resolver. + """ merged = RigidBodyAttributesCfg() for field_name in merged.__dataclass_fields__: + if field_name == "newton": + continue override_val = getattr(self, field_name) if override_val is not None: setattr(merged, field_name, override_val) else: setattr(merged, field_name, getattr(base, field_name)) - return merged.attr() + merged.newton = _merge_newton_subcfg(self.newton, base.newton) + return merged @classmethod def from_dict( @@ -363,7 +740,9 @@ def from_dict( """Initialize the configuration from a dictionary.""" cfg = cls() for key, value in init_dict.items(): - if hasattr(cfg, key): + if key == "newton" and isinstance(value, dict): + setattr(cfg, key, NewtonCollisionAttributesCfg.from_dict(value)) + elif hasattr(cfg, key): setattr(cfg, key, value) else: logger.log_warning( 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/diff/__init__.py b/embodichain/lab/sim/diff/__init__.py new file mode 100644 index 00000000..45483c72 --- /dev/null +++ b/embodichain/lab/sim/diff/__init__.py @@ -0,0 +1,36 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Differentiable Newton stepping for EmbodiChain. + +Bridges DexSim's :class:`~dexsim.engine.newton_physics.DifferentiableStepper` +into PyTorch autograd via a :class:`torch.autograd.Function`, and exposes a +:class:`tape_context` manager for advanced users who want to compose their +own Warp kernels. +""" + +from __future__ import annotations + +from .bridge import ( + NewtonStepFunc, + differentiable_step, + tape_context, +) + +__all__ = [ + "NewtonStepFunc", + "differentiable_step", + "tape_context", +] diff --git a/embodichain/lab/sim/diff/bridge.py b/embodichain/lab/sim/diff/bridge.py new file mode 100644 index 00000000..ba0ee49b --- /dev/null +++ b/embodichain/lab/sim/diff/bridge.py @@ -0,0 +1,193 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Warp-tape <-> PyTorch-autograd bridge for Newton physics.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Callable, Iterator + +import torch +import warp as wp + +if TYPE_CHECKING: + from embodichain.lab.sim.sim_manager import SimulationManager + +__all__ = ["NewtonStepFunc", "differentiable_step", "tape_context"] + + +@contextmanager +def tape_context(manager: "SimulationManager") -> Iterator[wp.Tape]: + """Open a Warp tape bound to the manager's Newton state. + + Advanced users compose their own Warp kernels inside this context, then + call ``tape.backward()`` outside the with-block. + """ + if not manager.is_newton_backend: + raise RuntimeError( + "tape_context requires the Newton backend with requires_grad=True." + ) + tape = wp.Tape() + with tape: + yield tape + + +def differentiable_step( + manager: "SimulationManager", + *, + apply_control_fn: Callable[[wp.Tape], None], + substeps: int, + dt: float | None = None, +) -> dict: + """Run one EmbodiChain-level physics step inside a Warp tape. + + Args: + manager: The owning :class:`SimulationManager` (must be Newton). + apply_control_fn: Callable that writes the joint/body control + targets inside the tape. Invoked once at the start of the + step. Receives the open tape; must launch Warp kernels (or + call dexsim setters that are tape-aware) to populate + ``manager.physics.newton_manager._control``. + substeps: Number of solver substeps to run (typically + ``sim_cfg.sim_steps_per_control``). + dt: Solver dt; defaults to the manager's configured dt. + + Returns: + A dict carrying the tape and the state buffers for the caller to + save in autograd context. + """ + if not manager.is_newton_backend: + raise RuntimeError("differentiable_step requires the Newton backend.") + nm = manager.physics.newton_manager + stepper = manager.create_differentiable_stepper() + state_in = nm._state_0 + state_out = nm._model.state() + contacts = stepper.create_contacts() + dt_val = nm.solver_dt if dt is None else float(dt) + + tape = wp.Tape() + with tape: + apply_control_fn(tape) + for _ in range(substeps): + stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) + state_in, state_out = state_out, state_in + + # The final state lives in state_in after the swap. + return { + "tape": tape, + "final_state": state_in, + "stepper": stepper, + } + + +class NewtonStepFunc(torch.autograd.Function): + """torch.autograd.Function bridging Warp tape autodiff to PyTorch. + + Forward: launches the action-to-control Warp kernel, runs the + caller-provided ``step_fn`` (differentiable solver loop or FK bypass), + and reads observation / reward as torch tensors via ``wp.to_torch`` + (zero-copy where possible). The obs/reward kernels launched by + ``obs_reward_fn`` run INSIDE the open Warp tape so that their outputs + carry gradient back to ``action_wp``. + + Backward: copies upstream PyTorch grads into the corresponding Warp + ``.grad`` buffers, calls ``tape.backward()``, and returns + ``wp.to_torch(action_wp.grad)`` reshaped to the action's tensor shape. + + Callers must supply a ``sim_state`` dict with the following keys: + manager: SimulationManager (Newton, requires_grad=True) + substeps: int (used by the default solver-based step_fn) + action_to_control_kernel: callable(action_wp, *kernel_args) + kernel_args: tuple consumed by action_to_control_kernel + obs_reward_fn: callable(final_state) -> dict with torch outputs + step_fn: optional callable() -> final Newton state; when omitted + the bridge runs the differentiable stepper for ``substeps`` + iterations (the original solver-based path) + + The ``obs_reward_fn`` must return a dict containing: + _order: tuple of output names (returned in this order) + _grad_track: dict mapping name -> Warp array (or None) whose + ``.grad`` should be seeded from the upstream PyTorch grad + : torch tensor for each name in ``_order`` + """ + + @staticmethod + def forward(ctx, action_torch: torch.Tensor, sim_state: dict): + manager = sim_state["manager"] + substeps = int(sim_state["substeps"]) + kernel = sim_state["action_to_control_kernel"] + kernel_args = sim_state["kernel_args"] + obs_reward_fn = sim_state["obs_reward_fn"] + step_fn = sim_state.get("step_fn") + + # Save the original action shape so backward can reshape the gradient. + ctx.saved_action_shape = action_torch.shape + + nm = manager.physics.newton_manager + + action_flat = action_torch.detach().clone().reshape(-1).contiguous() + action_wp = wp.from_torch(action_flat, dtype=wp.float32, requires_grad=True) + + tape = wp.Tape() + with tape: + kernel(action_wp, *kernel_args) # writes inputs for stepping + if step_fn is not None: + final_state = step_fn() + else: + stepper = manager.create_differentiable_stepper() + state_in = nm._state_0 + state_out = nm._model.state() + contacts = stepper.create_contacts() + dt_val = nm.solver_dt + for _ in range(substeps): + stepper.step(state_in, state_out, contacts=contacts, dt=dt_val) + state_in, state_out = state_out, state_in + final_state = state_in + # Compute obs/reward INSIDE the tape so the reward/obs kernels + # participate in the Warp autodiff graph. The torch tensors + # returned by obs_reward_fn are built via wp.to_torch of + # tape-tracked Warp arrays, so they carry gradient back to + # action_wp when tape.backward() is called. + outputs = obs_reward_fn(final_state) + + ctx.tape = tape + ctx.action_wp = action_wp + ctx.outputs_order = outputs["_order"] + ctx.outputs_grad_track = outputs.get("_grad_track", {}) + return tuple(outputs[k] for k in outputs["_order"]) + + @staticmethod + def backward(ctx, *grad_outputs): + # Copy each upstream grad back into the corresponding Warp .grad. + for name, grad_t in zip(ctx.outputs_order, grad_outputs): + wp_arr = ctx.outputs_grad_track.get(name) + if grad_t is None or wp_arr is None: + continue + # Warp allocates .grad lazily for arrays with requires_grad=True + # that participate in the tape; allocate defensively in case + # the array was created but never written inside the tape. + if wp_arr.grad is None: + wp_arr.grad = wp.zeros_like(wp_arr) + wp.copy( + wp_arr.grad, + wp.from_torch(grad_t.detach().clone().contiguous(), dtype=wp.float32), + ) + ctx.tape.backward() + action_grad = wp.to_torch(ctx.action_wp.grad).clone() + ctx.tape.zero() + # Reshape to the original action layout; second input (sim_state) + # has no gradient. + return action_grad.reshape(ctx.saved_action_shape), None diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index ed5a9e32..06d4e89a 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,16 +38,17 @@ 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 ( - get_dexsim_drive_type, - set_dexsim_articulation_cfg, -) +from embodichain.lab.sim.utility.sim_utils import get_dexsim_drive_type from embodichain.lab.sim.utility.solver_utils import ( create_pk_chain, create_pk_serial_chain, @@ -75,18 +74,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 ) - if self.device.type == "cuda" - else None - ) + else: + self.articulation_view = DefaultArticulationView( + entities=entities, ps=ps, device=device + ) + + # 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 +102,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 +129,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 +151,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 +166,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 +175,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 +184,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 +202,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 +211,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 +220,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 +228,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 +237,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 +246,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 +255,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 +264,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 +410,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 @@ -586,7 +421,7 @@ def __init__( # Store all indices for batch operations self._all_indices = torch.arange(len(entities), dtype=torch.int32) - if device.type == "cuda": + if device.type == "cuda" and not is_newton_scene(self._ps): self._world.update(0.001) self._data = ArticulationData(entities=entities, ps=self._ps, device=device) @@ -600,9 +435,6 @@ def __init__( # Determine if we should use USD properties or cfg properties. if not self.cfg.use_usd_properties: - # Set articulation configuration in DexSim - set_dexsim_articulation_cfg(entities, self.cfg) - num_entities = len(entities) dof = self._data.dof default_cfg = JointDrivePropertiesCfg() @@ -691,7 +523,7 @@ def __init__( self.active_joint_ids = [i for i in range(self.dof) if i not in self.mimic_ids] # TODO: very weird that we must call update here to make sure the GPU indices are valid. - if device.type == "cuda": + if device.type == "cuda" and not is_newton_scene(self._ps): self._world.update(0.001) super().__init__(cfg, entities, device) @@ -807,6 +639,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 +757,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 +913,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 +955,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 +977,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 +1001,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 +1055,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 +1093,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 +1134,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( @@ -1371,6 +1158,14 @@ def set_link_physical_attr( env_ids: Environment indices. If None, all environments are updated. base_attrs: Base config used when ``attrs`` is a partial override. replace_inertial: Recompute inertia when mass changes. + + .. attention:: + On the Newton backend, ``set_physical_attr`` only mirrors attributes + onto link metadata (consumed at the next scene rebuild). Mass is + additionally pushed live via ``set_link_mass`` so runtime per-link + mass overrides take effect immediately (mirroring the dedicated + :meth:`set_mass`). Friction/restitution/contact_offset have no live + per-link API on Newton articulations and are rebuild-time only. """ if link_names is None: matched_link_names = self.link_names @@ -1394,12 +1189,22 @@ def set_link_physical_attr( else: physical_attr = attrs + is_newton = self._data is not None and self._data.is_newton_backend local_env_ids = self._all_indices if env_ids is None else env_ids for env_idx in local_env_ids: for name in matched_link_names: + local_name = self._entity_link_name(env_idx, name) self._entities[env_idx].set_physical_attr( - physical_attr, name, is_replace_inertial=replace_inertial + physical_attr, + local_name, + is_replace_inertial=replace_inertial, ) + # On Newton, set_physical_attr is metadata-only; push mass live + # so runtime per-link mass overrides take effect immediately. + if is_newton: + self._entities[env_idx].set_link_mass( + local_name, physical_attr.mass + ) def set_joint_drive( self, @@ -1578,10 +1383,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 +1468,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 +1850,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..65dd4f06 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/base.py @@ -0,0 +1,360 @@ +# ---------------------------------------------------------------------------- +# 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.""" + ... + + @abstractmethod + def fetch_contact_offset( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch contact offsets into ``data`` as ``(N, 1)``.""" + ... + + @abstractmethod + def apply_contact_offset(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply contact offsets 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..9249e62f --- /dev/null +++ b/embodichain/lab/sim/objects/backends/default.py @@ -0,0 +1,738 @@ +# ---------------------------------------------------------------------------- +# 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_sim_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]) + + def fetch_contact_offset( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + raise NotImplementedError( + "Per-body contact_offset fetch is not exposed by the default backend." + ) + + def apply_contact_offset(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + raise NotImplementedError( + "Per-body contact_offset apply is not exposed by the default backend; " + "set it via RigidBodyAttributesCfg (consumed at build) instead." + ) + + # -- 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_sim_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..63eb1e4b --- /dev/null +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -0,0 +1,825 @@ +# ---------------------------------------------------------------------------- +# 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 + # 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) + + def fetch_contact_offset( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_scalar(self._get_data_type().CONTACT_OFFSET, data, body_ids) + + def apply_contact_offset(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_data(body_ids, self._get_data_type().CONTACT_OFFSET, 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: + 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) + + body_id = int(entity.get_sim_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_sim_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..c1f3357a 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,98 @@ 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 _get_newton_attr_or_none(self, env_idx: int): + """Return the Newton meta PhysicalAttr, or None when not present. + + Unlike :meth:`_get_newton_attr` this does not raise: objects spawned via + the desc-native path (``attrs.newton`` set) carry ``newton_shape``/ + ``newton_body`` descriptors instead of a legacy ``attr``, so they have + no meta ``PhysicalAttr`` to mirror onto. Used by the not-ready setter + paths to tolerate both spawn paths. + """ + 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) + if manager is None: + return None + return getattr(manager, "dexsim_meta", {}).get(entity_handle, {}).get("attr") + + def _set_newton_attr_meta(self, env_idx: int, physical_attr) -> None: + """Mirror a :class:`dexsim.types.PhysicalAttr` onto the stored Newton meta. + + Newton only models a subset of physical attributes at runtime (mass, + friction, restitution, contact_offset, COM, inertia); the remaining + fields (damping, ccd, sleep thresholds, solver iters, ...) are carried + as metadata for rebuild and for getter consistency. This helper keeps + that mirror in sync so :meth:`get_damping` / :meth:`get_mass` and the + next scene rebuild see the user's intent. + """ + attr = self._get_newton_attr(env_idx) + for name in ( + "mass", + "density", + "dynamic_friction", + "static_friction", + "restitution", + "contact_offset", + "rest_offset", + "linear_damping", + "angular_damping", + "sleep_threshold", + "enable_ccd", + "max_depenetration_velocity", + "min_position_iters", + "min_velocity_iters", + "max_linear_velocity", + "max_angular_velocity", + ): + setattr(attr, name, getattr(physical_attr, name)) + + 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 +460,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 +492,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 +551,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 +558,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 +616,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 +684,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,13 +729,61 @@ def set_attrs( f"Length of env_ids {len(local_env_ids)} does not match attrs length {len(attrs)}." ) - # TODO: maybe need to improve the physical attributes setter efficiency. + # Resolve per-env physical attrs into a flat list aligned with local_env_ids. if isinstance(attrs, RigidBodyAttributesCfg): - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_physical_attr(attrs.attr()) + physical_attrs = [attrs.attr() for _ in local_env_ids] else: - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_physical_attr(attrs[i].attr()) + physical_attrs = [a.attr() for a in attrs] + + if is_newton_scene(self._ps): + self._set_newton_attrs(physical_attrs, local_env_ids) + return + + # TODO: maybe need to improve the physical attributes setter efficiency. + for i, env_idx in enumerate(local_env_ids): + self._entities[env_idx].set_physical_attr(physical_attrs[i]) + + def _set_newton_attrs( + self, + physical_attrs: list, + local_env_ids, + ) -> None: + """Apply physical attributes on the Newton backend. + + Newton models only a subset of physical attributes at runtime + (mass, friction, restitution, contact_offset); the rest (damping, ccd, + sleep thresholds, solver iters, rest_offset, static_friction) are + metadata carried for rebuild and getter consistency. When the Newton + model is finalized (READY/STALE) the supported subset is pushed live + via the batch scene API; beforehand (BUILDER) the attributes are only + mirrored onto the meta so the next finalize consumes them. + """ + for i, env_idx in enumerate(local_env_ids): + self._set_newton_attr_meta(env_idx, physical_attrs[i]) + + if self._data is None or not self._data.body_view.is_ready: + logger.log_debug( + "Newton model is not finalized; physical attributes are mirrored " + "to metadata and applied at the next finalize_newton_physics()." + ) + return + + body_ids = self._data.body_ids_for(local_env_ids) + view = self._data.body_view + device = self.device + + def _stack(field: str) -> torch.Tensor: + return torch.as_tensor( + [getattr(a, field) for a in physical_attrs], + dtype=torch.float32, + device=device, + ).unsqueeze(-1) + + # Newton-supported runtime subset. + view.apply_mass(_stack("mass"), body_ids) + view.apply_friction(_stack("dynamic_friction"), body_ids) + view.apply_restitution(_stack("restitution"), body_ids) + view.apply_contact_offset(_stack("contact_offset"), body_ids) def set_mass( self, mass: torch.Tensor, env_ids: Sequence[int] | None = None @@ -677,9 +801,24 @@ 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]) + if is_newton_scene(self._ps): + # Not finalized: mirror to meta (consumed at next finalize). The + # PhysX-bound set_mass is not patched for Newton entities. + attr = self._get_newton_attr_or_none(env_idx) + if attr is not None: + attr.mass = float(mass_np[i]) + else: + 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 +831,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 +863,30 @@ def set_friction( f"Length of env_ids {len(local_env_ids)} does not match friction length {len(friction)}." ) - friction = friction.cpu().numpy() - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].get_physical_body().set_dynamic_friction( - friction[i] + 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, ) - self._entities[env_idx].get_physical_body().set_static_friction(friction[i]) + return + + friction_np = friction.cpu().numpy() + for i, env_idx in enumerate(local_env_ids): + if is_newton_scene(self._ps): + # Not finalized: mirror to meta (Newton has a single mu; consumed + # at next finalize). The PhysX-bound friction setters are not + # patched for Newton entities. + attr = self._get_newton_attr_or_none(env_idx) + if attr is not None: + attr.dynamic_friction = float(friction_np[i]) + else: + self._entities[env_idx].get_physical_body().set_dynamic_friction( + friction_np[i] + ) + self._entities[env_idx].get_physical_body().set_static_friction( + friction_np[i] + ) def get_friction(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """Get friction for the rigid object. @@ -733,11 +899,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) @@ -750,6 +925,12 @@ def set_damping( Args: damping (torch.Tensor): The damping to set with shape (N, 2), where the first column is linear damping and the second column is angular damping. env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. + + .. attention:: + The Newton backend does not simulate per-body linear/angular damping + (its damping is a global solver knob). On Newton this call mirrors + the values onto the attribute metadata so :meth:`get_damping` and + scene rebuilds stay consistent, but has no runtime effect. """ local_env_ids = self._all_indices if env_ids is None else env_ids @@ -758,13 +939,22 @@ def set_damping( f"Length of env_ids {len(local_env_ids)} does not match damping length {len(damping)}." ) - damping = damping.cpu().numpy() + damping = damping.to(dtype=torch.float32, device=self.device) + + if is_newton_scene(self._ps): + for i, env_idx in enumerate(local_env_ids): + attr = self._get_newton_attr(env_idx) + attr.linear_damping = float(damping[i, 0].item()) + attr.angular_damping = float(damping[i, 1].item()) + return + + damping_np = damping.cpu().numpy() for i, env_idx in enumerate(local_env_ids): self._entities[env_idx].get_physical_body().set_linear_damping( - damping[i, 0] + damping_np[i, 0] ) self._entities[env_idx].get_physical_body().set_angular_damping( - damping[i, 1] + damping_np[i, 1] ) def get_damping(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: @@ -780,12 +970,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,11 +1001,26 @@ def set_inertia( f"Length of env_ids {len(local_env_ids)} does not match inertia length {len(inertia)}." ) - inertia = 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] + 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): + if is_newton_scene(self._ps): + # Not finalized: mirror to meta (consumed at next finalize). The + # PhysX-bound inertia setter is not patched for Newton entities. + attr = self._get_newton_attr_or_none(env_idx) + if attr is not None: + attr.inertia = np.asarray(inertia_np[i], dtype=np.float32) + else: + self._entities[ + env_idx + ].get_physical_body().set_mass_space_inertia_tensor(inertia_np[i]) def get_inertia(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """Get inertia tensor for the rigid object. @@ -823,13 +1033,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 +1161,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 +1180,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. @@ -975,9 +1196,22 @@ def set_body_type(self, body_type: str) -> None: Args: body_type (str): The body type to set. Must be one of 'dynamic', or 'kinematic'. + + .. attention:: + On the Newton backend, body type (dynamic/kinematic/static) is fixed + at body registration and cannot be changed at runtime; switching it + would require re-registering the body and rebuilding the model. This + call is therefore a no-op on Newton. """ from dexsim.types import ActorType + if is_newton_scene(self._ps): + logger.log_warning( + "Newton backend does not support changing RigidObject body type at " + "runtime (it is fixed at registration). Skipping set_body_type call." + ) + 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 +1316,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 +1384,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 +1405,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/physics/__init__.py b/embodichain/lab/sim/physics/__init__.py new file mode 100644 index 00000000..deb91a87 --- /dev/null +++ b/embodichain/lab/sim/physics/__init__.py @@ -0,0 +1,73 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Physics backend registry and factory. + +Selects a concrete :class:`PhysicsBackend` from a physics config via +:func:`embodichain.lab.sim.cfg.physics_backend_from_cfg` and instantiates it +with the owning :class:`SimulationManager` as its back-reference. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from embodichain.lab.sim.cfg import physics_backend_from_cfg +from embodichain.utils import logger + +from .base import PhysicsBackend +from .default import DefaultPhysicsBackend +from .newton import NewtonPhysicsBackend + +if TYPE_CHECKING: + from embodichain.lab.sim.sim_manager import SimulationManager + +__all__ = [ + "PhysicsBackend", + "DefaultPhysicsBackend", + "NewtonPhysicsBackend", + "make_physics_backend", +] + +#: Registry of backend name -> backend class. +_BACKENDS: dict[str, type[PhysicsBackend]] = { + "default": DefaultPhysicsBackend, + "newton": NewtonPhysicsBackend, +} + + +def make_physics_backend(physics_cfg, manager: "SimulationManager") -> PhysicsBackend: + """Construct the physics backend for ``physics_cfg``. + + The backend subclass is selected by the *type* of ``physics_cfg`` + (via :func:`physics_backend_from_cfg`), so passing a + :class:`~embodichain.lab.sim.cfg.NewtonPhysicsCfg` activates the Newton + backend and a + :class:`~embodichain.lab.sim.cfg.DefaultPhysicsCfg` activates the default + backend. + + Args: + physics_cfg: The physics backend configuration. + manager: The owning :class:`SimulationManager` (passed as the + backend's back-reference). + + Returns: + The instantiated :class:`PhysicsBackend`. + """ + name = physics_backend_from_cfg(physics_cfg) + cls = _BACKENDS.get(name) + if cls is None: + logger.log_error(f"Unknown physics backend: {name!r}.") + return cls(manager) diff --git a/embodichain/lab/sim/physics/base.py b/embodichain/lab/sim/physics/base.py new file mode 100644 index 00000000..fd6ffab7 --- /dev/null +++ b/embodichain/lab/sim/physics/base.py @@ -0,0 +1,186 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Swappable physics-backend abstraction for :class:`SimulationManager`. + +This module defines the contract that every physics backend (DexSim default, +Newton/Warp, ...) satisfies. The owning :class:`SimulationManager` +holds a single :class:`PhysicsBackend` instance as ``self.physics`` and +delegates the backend-specific lifecycle, scene access, world-config +activation and capability queries to it, instead of branching on a backend +name string throughout the manager. + +The design deliberately mirrors IsaacLab's split of an orchestrator +(``SimulationContext``) from a swappable physics manager (``PhysicsManager``), +with one departure: EmbodiChain keeps the backend as a true *instance* member +rather than a class-singleton, because :class:`SimulationManager` is itself a +multiton (one instance per ``instance_id``) and a class-singleton backend +would break that. + +.. note:: + This ABC covers the *manager-level* backend surface (lifecycle, scene, + capabilities, world-config). The per-asset read/write contract lives in + :mod:`embodichain.lab.sim.objects.backends` (``RigidBodyViewBase`` / + ``ArticulationViewBase``). +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import dexsim + + from embodichain.lab.sim.cfg import SimulationManagerCfg + from embodichain.lab.sim.sim_manager import SimulationManager + +__all__ = ["PhysicsBackend"] + + +class PhysicsBackend(ABC): + """Abstract base class for a swappable physics backend. + + A backend is constructed with a back-reference to its owning + :class:`SimulationManager` (from which it reaches the dexsim world, the + resolved device, the asset registries and the physics config). All + backend-specific behaviour is expressed as overrides of the methods and + properties below; the manager never inspects ``self.physics.name`` to + decide what to do (it only exposes it for backwards-compatible public + properties). + """ + + #: Backend identifier, e.g. ``"default"`` or ``"newton"``. + name: str = "" + + def __init__(self, manager: "SimulationManager") -> None: + self._manager: "SimulationManager" = manager + + # ------------------------------------------------------------------ # + # Construction / world-config activation + # ------------------------------------------------------------------ # + @abstractmethod + def configure_world( + self, + world_config: "dexsim.WorldConfig", + sim_config: "SimulationManagerCfg", + ) -> None: + """Apply backend-specific fields to the dexsim ``WorldConfig``. + + Called from :meth:`SimulationManager._convert_sim_config` after the + shared world-config fields and the resolved device have been set, so + implementations may read ``self._manager.device``. + + Args: + world_config: The dexsim world config to mutate in place. + sim_config: The full simulation manager config. + """ + + @abstractmethod + def activate(self, sim_config: "SimulationManagerCfg") -> None: + """Perform backend setup immediately after the dexsim World is created. + + This is the counterpart of the backend split that used to live in + ``SimulationManager.__init__`` (default ``set_physics_config`` vs + ``get_newton_manager``). + """ + + # ------------------------------------------------------------------ # + # Lifecycle + # ------------------------------------------------------------------ # + @abstractmethod + def ensure_initialized(self) -> None: + """Ensure the backend runtime is ready before a physics step. + + Called at the top of :meth:`SimulationManager.update`. For the default + backend this lazy-initializes GPU physics; for Newton it finalizes the + scene (rebuilding if the scene was mutated). Idempotent. + """ + + @abstractmethod + def invalidate(self) -> None: + """Mark the backend scene as needing re-initialization. + + Called after any scene mutation (adding/removing assets) so that the + next :meth:`ensure_initialized` rebuilds as needed. A no-op for + backends without a dirty/finalize lifecycle. + """ + + @abstractmethod + def prepare(self) -> None: + """Force the backend into a ready-to-step state. + + This unifies what the legacy code exposed as two separate operations - + "GPU physics init" on the default backend and "Newton finalize" - into a + single backend-agnostic entry point. It is idempotent: a backend that is + already ready is a no-op, and after :meth:`invalidate` the next call + re-prepares (re-initializes GPU physics / re-finalizes the Newton scene) + as needed. + + Called both lazily by :meth:`ensure_initialized` before each step and + directly by the public :meth:`SimulationManager.init_gpu_physics` and + :meth:`SimulationManager.finalize_newton_physics` entry points (both of + which delegate here). + """ + + @property + @abstractmethod + def is_initialized(self) -> bool: + """Whether the backend runtime has been initialized/finalized.""" + + # ------------------------------------------------------------------ # + # Scene access + # ------------------------------------------------------------------ # + @abstractmethod + def get_scene(self): + """Return the active physics scene object (default DexSim or Newton).""" + + @property + def newton_manager(self): + """The DexSim Newton manager, or ``None`` if not the Newton backend. + + Returns: + The :class:`dexsim.engine.newton_physics.NewtonManager` for the + Newton backend, otherwise ``None``. + """ + return None + + # ------------------------------------------------------------------ # + # Capabilities (override in subclasses; defaults are conservative) + # ------------------------------------------------------------------ # + @property + def supports_soft_bodies(self) -> bool: + """Whether this backend can simulate soft bodies.""" + return False + + @property + def supports_cloth(self) -> bool: + """Whether this backend can simulate cloth bodies.""" + return False + + @property + def supports_rigid_object_group(self) -> bool: + """Whether this backend supports rigid object groups.""" + return False + + @property + def supports_robot(self) -> bool: + """Whether this backend supports robots (articulated URDF assets).""" + return False + + @property + def can_disable_manual_update(self) -> bool: + """Whether ``set_manual_update(False)`` is permitted on this backend.""" + return True diff --git a/embodichain/lab/sim/physics/default.py b/embodichain/lab/sim/physics/default.py new file mode 100644 index 00000000..fefb96a0 --- /dev/null +++ b/embodichain/lab/sim/physics/default.py @@ -0,0 +1,128 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""DexSim default physics backend.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import dexsim + +from embodichain.lab.sim.cfg import DefaultPhysicsCfg +from embodichain.utils import logger + +from .base import PhysicsBackend + +if TYPE_CHECKING: + import dexsim as _dexsim # noqa: F401 + + from embodichain.lab.sim.cfg import SimulationManagerCfg + +__all__ = ["DefaultPhysicsBackend"] + + +class DefaultPhysicsBackend(PhysicsBackend): + """The legacy DexSim default physics backend (GPU or CPU).""" + + name = "default" + + def __init__(self, manager) -> None: + super().__init__(manager) + self._is_initialized_gpu_physics = False + + # -- construction / world-config activation ------------------------- # + def configure_world(self, world_config, sim_config: "SimulationManagerCfg") -> None: + cfg = sim_config.physics_cfg + assert isinstance(cfg, DefaultPhysicsCfg) + world_config.length_tolerance = cfg.length_tolerance + world_config.speed_tolerance = cfg.speed_tolerance + if self._manager.device.type == "cuda": + world_config.enable_gpu_sim = True + world_config.direct_gpu_api = True + + def activate(self, sim_config: "SimulationManagerCfg") -> None: + cfg = sim_config.physics_cfg + assert isinstance(cfg, DefaultPhysicsCfg) + dexsim.set_physics_config(**cfg.to_dexsim_args()) + dexsim.set_physics_gpu_memory_config(**cfg.gpu_memory.to_dict()) + + # -- lifecycle ------------------------------------------------------ # + def invalidate(self) -> None: + # The default backend has no dirty/finalize lifecycle. + pass + + @property + def is_initialized(self) -> bool: + return self._is_initialized_gpu_physics + + def prepare(self) -> None: + """Initialize GPU physics for the default backend. + + Implements the unified :meth:`PhysicsBackend.prepare` contract. For the + default backend "becoming ready to step" is initializing GPU physics; on + CPU there is nothing to initialize so this is a no-op. + """ + if not self._manager.is_use_gpu_physics: + logger.log_warning( + "The simulation device is not cuda, cannot initialize GPU physics." + ) + return + + if self._is_initialized_gpu_physics: + return + + for art in self._manager._articulations.values(): + art.reallocate_body_data() + for robot in self._manager._robots.values(): + robot.reallocate_body_data() + + # Re-establish rigid object positions after articulation resets, ensuring + # no articulation kinematics step has inadvertently corrupted the broadphase + # state for rigid bodies. + for rigid_obj in self._manager._rigid_objects.values(): + rigid_obj.reset() + + self._is_initialized_gpu_physics = True + + def ensure_initialized(self) -> None: + if self._manager.is_use_gpu_physics and not self._is_initialized_gpu_physics: + logger.log_warning( + "Using GPU physics, but not initialized yet. Forcing initialization." + ) + self.prepare() + + # -- scene ---------------------------------------------------------- # + def get_scene(self): + return self._manager._world.get_physics_scene() + + # -- capabilities --------------------------------------------------- # + # The default backend supports soft/cloth on GPU; the GPU + # precondition itself is enforced separately in SimulationManager. + @property + def supports_soft_bodies(self) -> bool: + return True + + @property + def supports_cloth(self) -> bool: + return True + + @property + def supports_rigid_object_group(self) -> bool: + return True + + @property + def supports_robot(self) -> bool: + return True diff --git a/embodichain/lab/sim/physics/newton.py b/embodichain/lab/sim/physics/newton.py new file mode 100644 index 00000000..86c97639 --- /dev/null +++ b/embodichain/lab/sim/physics/newton.py @@ -0,0 +1,162 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Newton (Warp) physics backend. + +Wraps DexSim's Newton module (``dexsim.engine.newton_physics``), which itself +runs NVIDIA Newton solvers (MuJoCo-Warp / XPBD / Featherstone / VBD / +semi-implicit) on Warp. The backend owns the lazy finalize/invalidate state +machine that rebuilds the Newton model whenever the scene is mutated. +""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING + +from embodichain.utils import logger + +from .base import PhysicsBackend + +if TYPE_CHECKING: + from dexsim.engine.newton_physics import NewtonManager + + from embodichain.lab.sim.cfg import SimulationManagerCfg + +__all__ = ["NewtonPhysicsBackend"] + + +class NewtonPhysicsBackend(PhysicsBackend): + """The DexSim Newton physics backend (Warp-based).""" + + name = "newton" + + def __init__(self, manager) -> None: + super().__init__(manager) + self._newton_manager: "NewtonManager | None" = None + self._is_finalized = False + + # -- construction / world-config activation ------------------------- # + def configure_world(self, world_config, sim_config: "SimulationManagerCfg") -> None: + importlib.import_module("dexsim.engine.newton_physics") + + newton_physics_cfg = sim_config.physics_cfg + world_config.newton_cfg = newton_physics_cfg.to_dexsim_cfg( + gpu_id=sim_config.gpu_id, + ) + + def activate(self, sim_config: "SimulationManagerCfg") -> None: + from dexsim.engine.newton_physics import get_newton_manager + + self._newton_manager = get_newton_manager(self._manager._world) + + # -- lifecycle ------------------------------------------------------ # + def invalidate(self) -> None: + """Mark the Newton scene as needing re-finalization after a mutation.""" + self._is_finalized = False + + @property + def is_initialized(self) -> bool: + return self._is_finalized + + @property + def newton_manager(self) -> "NewtonManager | None": + if self._newton_manager is None: + from dexsim.engine.newton_physics import get_newton_manager + + self._newton_manager = get_newton_manager(self._manager._world) + return self._newton_manager + + def _lifecycle_state(self) -> str: + """Return the Newton manager lifecycle state name, or empty string.""" + mgr = self.newton_manager + return getattr(getattr(mgr, "lifecycle_state", None), "name", "") + + def _reset_entities_after_finalize(self) -> None: + """Apply deferred initial resets once Newton runtime data is ready.""" + for rigid_obj in self._manager._rigid_objects.values(): + rigid_obj.reset() + for articulation in self._manager._articulations.values(): + articulation.reset() + for robot in self._manager._robots.values(): + robot.reset() + # Rigid object groups are not supported on the Newton backend yet. + + def prepare(self) -> None: + """Finalize the Newton scene if it has not been finalized yet. + + Implements the unified :meth:`PhysicsBackend.prepare` contract: this is + both the "finalize" entry point (public + :meth:`SimulationManager.finalize_newton_physics`) and the "GPU init" + entry point (:meth:`SimulationManager.init_gpu_physics`) for the Newton + backend, since Newton's notion of becoming ready to step is finalizing + the model. + """ + if self._is_finalized and self._lifecycle_state() == "READY": + return + + mgr = self.newton_manager + state = self._lifecycle_state() + + if state != "READY": + from dexsim.engine.newton_physics.rebuild import ( + ensure_simulation_prepared_lazy, + rebuild_newton_from_scene, + ) + + safe_to_continue, _ = ensure_simulation_prepared_lazy( + mgr, + self._manager._world, + rebuild_from_scene=rebuild_newton_from_scene, + warn=True, + ) + if not safe_to_continue: + logger.log_error( + "Failed to finalize Newton physics: model is not ready to build " + f"(lifecycle state {state!r})." + ) + return + + state = self._lifecycle_state() + if state != "READY": + logger.log_error( + "Failed to finalize Newton physics: lifecycle state is " + f"{state!r} after simulation preparation." + ) + + self._is_finalized = True + self._reset_entities_after_finalize() + + def ensure_initialized(self) -> None: + self.prepare() + + # -- scene ---------------------------------------------------------- # + def get_scene(self): + return self.newton_manager.scene + + # -- capabilities --------------------------------------------------- # + @property + def supports_robot(self) -> bool: + # Robots are URDF articulations; the Newton ``load_urdf`` patch builds a + # NewtonArticulation, and the shared spawn path (add_robot invalidate + + # _reset_entities_after_finalize) handles the Newton lifecycle. Requires + # the dexsim fix to ``NewtonArticulation._joint_metas_from_ids`` so that + # explicit joint_ids use active-joint indexing (matching get_dof()). + return True + + @property + def can_disable_manual_update(self) -> bool: + # Newton cannot switch between manual and automatic update. + return False diff --git a/embodichain/lab/sim/physics_attrs.py b/embodichain/lab/sim/physics_attrs.py new file mode 100644 index 00000000..7a1d6907 --- /dev/null +++ b/embodichain/lab/sim/physics_attrs.py @@ -0,0 +1,253 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Backend-aware resolution of rigid-body physical attributes. + +This module is the EmbodiChain counterpart of dexsim's spawn-descriptor +resolver (``dexsim.spawn.adapters.newton_adapter``). It decouples the flat +:class:`~embodichain.lab.sim.cfg.RigidBodyAttributesCfg` (backend-neutral common +fields + an optional ``newton`` sub-config) from the backend-specific +descriptors dexsim consumes: + +- On the **default** backend it returns the legacy + :class:`dexsim.types.PhysicalAttr` (unchanged behaviour). +- On the **Newton** backend it builds a resolved Newton shape descriptor + (carrying the backend-neutral ``mu``/``restitution``/``has_shape_collision`` + projected from common fields, plus the Newton-native sub-config fields) and a + :class:`dexsim.spawn.descs.RigidBodyPhysicsDesc` body descriptor, suitable for + dexsim's desc-native ``register_mesh_object_to_newton_patch`` entry point. + +It also emits data-driven warnings (ported from dexsim) when a user sets contact +fields the active Newton solver ignores, or PhysX-only fields on the Newton +backend. + +.. note:: + Newton-native contact/shape params (``ke``/``kd``/``margin``/...) are + **build-time only**: there is no runtime batch API to mutate them. Runtime + mutation (``RigidObject.set_attrs``) still applies the supported live subset + (mass/friction/restitution/contact_offset). +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Any + +import numpy as np + +from dexsim.spawn.descs import ( + NEWTON_CONTACT_FIELDS, + NEWTON_CONTACT_SOLVER_FIELDS, + NewtonCollisionDesc, + RigidBodyPhysicsDesc, +) + +from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.utils import logger + +if TYPE_CHECKING: + from dexsim.types import ActorType, PhysicalAttr + +__all__ = [ + "NEWTON_CONTACT_FIELDS", + "NEWTON_CONTACT_SOLVER_FIELDS", + "ResolvedNewtonShape", + "resolve_newton_shape", + "resolve_newton_body", + "resolve_rigid_body_attributes", + "warn_ignored_contact_fields", + "warn_backend_mismatched_fields", +] + + +# PhysX-only fields (carried on RigidBodyAttributesCfg) that Newton does not +# model per body. Setting them on the Newton backend is a no-op; warn so users +# notice. `static_friction` is folded into Newton's single `mu`; `rest_offset` +# has no Newton per-shape runtime equivalent (only `contact_offset`/`gap`). +_NEWTON_IGNORED_FIELDS: tuple[str, ...] = ( + "angular_damping", + "linear_damping", + "sleep_threshold", + "enable_ccd", + "max_depenetration_velocity", + "min_position_iters", + "min_velocity_iters", + "max_linear_velocity", + "max_angular_velocity", + "rest_offset", + "static_friction", +) + + +@dataclass +class ResolvedNewtonShape(NewtonCollisionDesc): + """Newton shape descriptor after common-field projection. + + Mirrors dexsim's internal ``_ResolvedNewtonCollisionDesc``: a + :class:`dexsim.spawn.descs.NewtonCollisionDesc` extended with the four + ``newton.ModelBuilder.ShapeConfig`` knobs whose values are *projected* from + backend-neutral common fields rather than read from the Newton sub-config. + + Field names mirror ``ShapeConfig`` attributes so dexsim's + ``_newton_shape_cfg_from_desc`` overlays them by name. + """ + + density: float | None = None + mu: float | None = None + restitution: float | None = None + has_shape_collision: bool | None = None + + +def resolve_newton_shape(cfg_attrs: RigidBodyAttributesCfg) -> ResolvedNewtonShape: + """Project a :class:`RigidBodyAttributesCfg` onto a Newton shape descriptor. + + Backend-neutral common fields map to the four projected ``ShapeConfig`` + knobs (``dynamic_friction``→``mu``, ``restitution``, ``enable_collision``→ + ``has_shape_collision``, ``density``); Newton-native sub-config fields are + copied verbatim. ``density`` is always set (positive) so dexsim can compute + a positive body mass from shape density even when only ``mass`` (no + explicit inertia) is given. + + Args: + cfg_attrs: The rigid-body attribute config (with optional ``newton``). + + Returns: + The resolved Newton shape descriptor. + """ + newton_cfg = cfg_attrs.newton + data: dict[str, Any] = {} + if newton_cfg is not None: + for f in fields(NewtonCollisionDesc): + val = getattr(newton_cfg, f.name) + if val is not None: + data[f.name] = val + return ResolvedNewtonShape( + **data, + density=cfg_attrs.density, + mu=cfg_attrs.dynamic_friction, + restitution=cfg_attrs.restitution, + has_shape_collision=cfg_attrs.enable_collision, + ) + + +def resolve_newton_body( + cfg_attrs: RigidBodyAttributesCfg, actor_type: "ActorType" +) -> RigidBodyPhysicsDesc: + """Build a :class:`RigidBodyPhysicsDesc` body descriptor from common fields. + + dexsim reads ``mass``/``inertia``/``com_position``/``com_quaternion`` + duck-typed from the body descriptor (``actor_type`` is passed separately to + the registration). Inertia is forwarded only if set on the cfg; otherwise + dexsim derives it from shape density. + + Args: + cfg_attrs: The rigid-body attribute config. + actor_type: The dexsim :class:`ActorType` for this body. + + Returns: + The body descriptor. + """ + kwargs: dict[str, Any] = {"mass": cfg_attrs.mass} + if cfg_attrs.density is not None: + kwargs["density"] = cfg_attrs.density + # Inertia / COM are not exposed on RigidBodyAttributesCfg today; if a future + # config extension adds them, forward them here. Kept explicit for clarity. + return RigidBodyPhysicsDesc(actor_type=actor_type, **kwargs) + + +def resolve_rigid_body_attributes( + cfg_attrs: RigidBodyAttributesCfg, + backend: str, + solver_type: str | None = None, +) -> "PhysicalAttr | ResolvedNewtonShape": + """Resolve a config into the backend-specific descriptor. + + For the Newton backend this returns the resolved Newton shape descriptor + (and emits per-solver / backend-mismatch warnings); the caller builds the + body descriptor separately via :func:`resolve_newton_body` since it owns the + ``actor_type``. + + Args: + cfg_attrs: The rigid-body attribute config. + backend: ``"default"`` or ``"newton"``. + solver_type: Active Newton solver type (e.g. ``"mujoco_warp"``); only + consulted on the Newton backend for contact-field warnings. May be + ``None`` to skip the per-solver warning. + + Returns: + A :class:`dexsim.types.PhysicalAttr` for the default backend, or a + :class:`ResolvedNewtonShape` for the Newton backend. + """ + if backend == "newton": + shape = resolve_newton_shape(cfg_attrs) + if solver_type is not None: + warn_ignored_contact_fields(shape, solver_type) + warn_backend_mismatched_fields(cfg_attrs, backend) + return shape + return cfg_attrs.attr() + + +def warn_ignored_contact_fields( + newton_shape: NewtonCollisionDesc | ResolvedNewtonShape | None, + solver_type: str, +) -> None: + """Warn for contact-material fields the active Newton solver does not read. + + Ported from dexsim's ``_warn_ignored_contact_fields``. A field the user set + (non-None) that is a contact-material field but not in the active solver's + read set is a harmless no-op; this makes it visible. + """ + if newton_shape is None: + return + read_fields = NEWTON_CONTACT_SOLVER_FIELDS.get(solver_type) + if read_fields is None: + return + ignored = sorted( + f.name + for f in fields(newton_shape) + if getattr(newton_shape, f.name) is not None + and f.name in NEWTON_CONTACT_FIELDS + and f.name not in read_fields + ) + if ignored: + logger.log_warning( + f"Newton solver '{solver_type}' ignores contact field(s) {ignored}; " + "they have no effect for this solver." + ) + + +def warn_backend_mismatched_fields( + cfg_attrs: RigidBodyAttributesCfg, backend: str +) -> None: + """Warn for attribute fields the active backend does not model. + + On the Newton backend, PhysX-only per-body fields (damping, ccd, sleep + thresholds, solver iters, rest_offset, static_friction) are not modelled; + setting them is a no-op. The warning fires only when the user deviated from + the cfg defaults, so it does not spam the common case. + """ + if backend != "newton": + return + defaults = RigidBodyAttributesCfg() + ignored = sorted( + name + for name in _NEWTON_IGNORED_FIELDS + if getattr(cfg_attrs, name) != getattr(defaults, name) + ) + if ignored: + logger.log_warning( + f"Newton backend does not model PhysX-only field(s) {ignored}; " + "they have no runtime effect on Newton." + ) 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 e8ac2916..b6e68d9f 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -47,14 +47,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 +74,10 @@ ) from embodichain.lab.sim.cfg import ( RenderCfg, - PhysicsCfg, + DefaultPhysicsCfg, + NewtonPhysicsCfg, + validate_physics_cfg, MarkerCfg, - GPUMemoryCfg, WindowRecordCfg, LightCfg, RigidObjectCfg, @@ -87,6 +87,7 @@ ArticulationCfg, RobotCfg, ) +from embodichain.lab.sim.physics import make_physics_backend from embodichain.lab.sim import VisualMaterial, VisualMaterialCfg from embodichain.utils import configclass, logger from embodichain.utils.math import look_at_to_pose @@ -94,6 +95,7 @@ __all__ = [ "SimulationManager", "SimulationManagerCfg", + "get_physics_scene", "SIM_CACHE_DIR", "MATERIAL_CACHE_DIR", "CONVEX_DECOMP_DIR", @@ -105,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.""" @@ -139,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: @@ -236,6 +298,12 @@ def __init__( self.sim_config = sim_config self.device = torch.device("cpu") + # Initialize physics backend (selected by the type of physics_cfg). + # The backend is held as an instance member; SimulationManager delegates + # all backend-specific lifecycle/scene/capability logic to it instead of + # branching on a backend name throughout the manager. + self.physics = make_physics_backend(sim_config.physics_cfg, self) + world_config = self._convert_sim_config(sim_config) # Initialize warp runtime context before creating the world. @@ -259,14 +327,11 @@ 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()) - - self._is_initialized_gpu_physics = False - self._ps = self._world.get_physics_scene() + # Activate the physics backend now that the dexsim World exists. + self.physics.activate(sim_config) # activate physics self.enable_physics(True) @@ -410,9 +475,32 @@ 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.name + + @property + def is_default_backend(self) -> bool: + """Whether the existing DexSim default physics backend is active.""" + return self.physics.name == "default" + + @property + def is_newton_backend(self) -> bool: + """Whether the DexSim Newton physics backend is active.""" + return self.physics.name == "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 + return self.physics.newton_manager + @property def is_physics_manually_update(self) -> bool: return self._world.is_physics_manually_update() @@ -451,8 +539,6 @@ 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 sim_config.render_cfg.renderer == "auto": from embodichain.lab.sim.utility.render_utils import ( @@ -466,19 +552,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." @@ -489,6 +572,10 @@ def _convert_sim_config( world_config.gpu_id = sim_config.gpu_id + # Apply backend-specific WorldConfig fields (default tolerances/GPU flags + # or the Newton cfg) via the active backend. + self.physics.configure_world(world_config, sim_config) + return world_config def _init_sim_resources(self) -> None: @@ -497,6 +584,15 @@ def _init_sim_resources(self) -> None: self._default_resources = SimResources() + def _invalidate_newton_physics(self) -> None: + """Mark the active backend scene as needing re-initialization. + + Delegates to the active :class:`PhysicsBackend`; a no-op for backends + without a dirty/finalize lifecycle. Called after every scene mutation + (adding assets, creating the default plane). + """ + self.physics.invalidate() + def enable_physics(self, enable: bool) -> None: """Enable or disable physics simulation. @@ -514,31 +610,81 @@ def set_manual_update(self, enable: bool) -> None: Args: enable (bool): whether to enable manual update. """ + if not self.physics.can_disable_manual_update and enable is False: + logger.log_warning( + "The active 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": - logger.log_warning( - "The simulation device is not cuda, cannot initialize GPU physics." + """Initialize the GPU physics simulation. + + Delegates to the active backend's unified :meth:`PhysicsBackend.prepare` + (for the default backend this performs the real GPU initialization; for + the Newton backend it finalizes the scene). + """ + self.physics.prepare() + + def finalize_newton_physics(self) -> None: + """Finalize the Newton scene if it has not been finalized yet. + + Delegates to the active backend's unified :meth:`PhysicsBackend.prepare` + (for the Newton backend this (re-)finalizes the scene and applies + deferred entity resets; for the default backend it initializes GPU + physics). + """ + self.physics.prepare() + + def create_differentiable_stepper(self): + """Create a single-step differentiable physics primitive (Newton-only). + + Requires the Newton backend with ``requires_grad=True`` and + ``solver_type="semi_implicit"``. Delegates to + :meth:`dexsim.engine.newton_physics.NewtonManager.create_differentiable_stepper`. + + Raises: + RuntimeError: If the active backend is not Newton or if the + Newton manager is not ready / not in grad mode. + """ + if not self.is_newton_backend: + logger.log_error( + "create_differentiable_stepper requires the Newton backend." ) - return + return self.physics.newton_manager.create_differentiable_stepper() - if self._is_initialized_gpu_physics: - return + def create_gradient_rollout( + self, + record_steps: int, + substeps_per_record: int | None = None, + record_dt: float | None = None, + ): + """Create a gradient rollout buffer (Newton-only). - for art in self._articulations.values(): - art.reallocate_body_data() - for robot in self._robots.values(): - robot.reallocate_body_data() + Delegates to + :meth:`dexsim.engine.newton_physics.NewtonManager.create_gradient_rollout`. - # Re-establish rigid object positions after articulation resets, ensuring - # no articulation kinematics step has inadvertently corrupted the broadphase - # state for rigid bodies. - for rigid_obj in self._rigid_objects.values(): - rigid_obj.reset() + Args: + record_steps: Number of record points to capture in the rollout + buffer. + substeps_per_record: Newton substeps between successive record + points. Defaults to the Newton manager's configured + ``num_substeps``. + record_dt: Time interval between successive record points. + Defaults to the Newton manager's configured ``dt``. - self._is_initialized_gpu_physics = True + Raises: + RuntimeError: If the active backend is not Newton or if the + Newton manager is not ready / not in grad mode. + """ + if not self.is_newton_backend: + logger.log_error("create_gradient_rollout requires the Newton backend.") + return self.physics.newton_manager.create_gradient_rollout( + record_steps=record_steps, + substeps_per_record=substeps_per_record, + record_dt=record_dt, + ) def render_camera_group(self, group_ids: list[int]) -> None: """Render all camera group in the simulation. @@ -551,18 +697,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: - logger.log_warning( - f"Using GPU physics, but not initialized yet. Forcing initialization." - ) - self.init_gpu_physics() + # Ensure the active backend runtime is ready (lazy GPU init for the + # default backend, scene finalize for the Newton backend). + self.physics.ensure_initialized() if self.is_physics_manually_update: if physics_dt is None: @@ -577,6 +721,8 @@ def update(self, physics_dt: float | None = None, step: int = 10) -> None: self._window_record_state, physics_dt ) + # TODO: Maybe add newton manager forward kinematics update. + else: logger.log_warning("Physics simulation is not manually updated.") @@ -604,6 +750,10 @@ 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.""" + return self.physics.get_scene() + def open_window(self) -> None: """Open the simulation window.""" self._world.open_window() @@ -687,16 +837,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.""" @@ -884,13 +1027,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 @@ -903,6 +1051,13 @@ def add_soft_object(self, cfg: SoftObjectCfg) -> SoftObject: Returns: SoftObject: The added soft object instance handle. """ + if not self.physics.supports_soft_bodies: + logger.log_error( + f"Soft object support is not enabled for the " + f"{self.physics.name} backend yet.", + error_type=NotImplementedError, + ) + if not self.is_use_gpu_physics: logger.log_error("Soft object requires GPU physics to be enabled.") @@ -933,6 +1088,13 @@ def add_cloth_object(self, cfg: ClothObjectCfg) -> ClothObject: Returns: ClothObject: The added cloth object instance handle. """ + if not self.physics.supports_cloth: + logger.log_error( + f"Cloth object support is not enabled for the " + f"{self.physics.name} backend yet.", + error_type=NotImplementedError, + ) + if not self.is_use_gpu_physics: logger.log_error("Cloth object requires GPU physics to be enabled.") @@ -1026,6 +1188,13 @@ def add_rigid_object_group(self, cfg: RigidObjectGroupCfg) -> RigidObjectGroup: Args: cfg (RigidObjectGroupCfg): Configuration for the rigid object group. """ + if not self.physics.supports_rigid_object_group: + logger.log_error( + f"Rigid object group support is not enabled for the " + f"{self.physics.name} backend yet.", + error_type=NotImplementedError, + ) + from embodichain.lab.sim.utility.sim_utils import ( load_mesh_objects_from_cfg, ) @@ -1055,10 +1224,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 @@ -1129,7 +1301,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] @@ -1142,39 +1313,27 @@ def add_articulation( is_usd = cfg.fpath.endswith((".usd", ".usda", ".usdc")) if is_usd: - # TODO: Currently add checking for num_envs when file is USD. After we support spawn via cloning, we can remove this. - if len(env_list) > 1: - logger.log_error(f"Currently not supporting multiple arenas for USD.") - env = self._env - results = env.import_from_usd_file( - cfg.fpath, return_object=True, cache_dir=self._convex_decomp_dir + from embodichain.lab.sim.utility.sim_utils import ( + spawn_usd_articulation_entities, ) - # print("USD import results:", results) - - articulations_found = [] - for key, value in results.items(): - if isinstance(value, dexsim.engine.Articulation): - articulations_found.append(value) - if len(articulations_found) == 0: - logger.log_error(f"No articulation found in USD file {cfg.fpath}.") - elif len(articulations_found) > 1: - logger.log_error( - f"Multiple articulations found in USD file {cfg.fpath}. " - ) - elif len(articulations_found) == 1: - obj_list.append(articulations_found[0]) + obj_list = spawn_usd_articulation_entities( + cfg, env_list, cache_dir=self._convex_decomp_dir + ) else: # non-usd file does not support this option, will be forced set False to avoid potential issues. cfg.use_usd_properties = False - for env in env_list: - art = env.load_urdf(cfg.fpath) - obj_list.append(art) + from embodichain.lab.sim.utility.sim_utils import ( + spawn_articulation_entities, + ) + + obj_list = spawn_articulation_entities(cfg, env_list) articulation = Articulation(cfg=cfg, entities=obj_list, device=self.device) self._articulations[uid] = articulation + self._invalidate_newton_physics() return articulation @@ -1209,6 +1368,12 @@ def add_robot(self, cfg: RobotCfg) -> Robot | None: Returns: Robot | None: The added robot instance handle, or None if failed. """ + if not self.physics.supports_robot: + logger.log_error( + f"Robot support is not enabled for the " + f"{self.physics.name} backend yet.", + error_type=NotImplementedError, + ) uid = cfg.uid if cfg.fpath is None: @@ -1238,37 +1403,25 @@ def add_robot(self, cfg: RobotCfg) -> Robot | None: is_usd = cfg.fpath.endswith((".usd", ".usda", ".usdc")) if is_usd: - # TODO: Currently add checking for num_envs when file is USD. After we support spawn via cloning, we can remove this. - if len(env_list) > 1: - logger.log_error(f"Currently not supporting multiple arenas for USD.") - env = self._env - results = env.import_from_usd_file(cfg.fpath, return_object=True) - # print("USD import results:", results) - - articulations_found = [] - for key, value in results.items(): - if isinstance(value, dexsim.engine.Articulation): - articulations_found.append(value) - - if len(articulations_found) == 0: - logger.log_error(f"No articulation found in USD file {cfg.fpath}.") - elif len(articulations_found) > 1: - logger.log_error( - f"Multiple articulations found in USD file {cfg.fpath}. " - ) - elif len(articulations_found) == 1: - obj_list.append(articulations_found[0]) + from embodichain.lab.sim.utility.sim_utils import ( + spawn_usd_articulation_entities, + ) + + obj_list = spawn_usd_articulation_entities(cfg, env_list) else: # non-usd file does not support this option, will be forced set False to avoid potential issues. cfg.use_usd_properties = False - for env in env_list: - art = env.load_urdf(cfg.fpath) - obj_list.append(art) + from embodichain.lab.sim.utility.sim_utils import ( + spawn_articulation_entities, + ) + + obj_list = spawn_articulation_entities(cfg, env_list) robot = Robot(cfg=cfg, entities=obj_list, device=self.device) self._robots[uid] = robot + self._invalidate_newton_physics() return robot @@ -2280,7 +2433,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 @@ -2322,3 +2474,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..9f0257d9 100644 --- a/embodichain/lab/sim/utility/sim_utils.py +++ b/embodichain/lab/sim/utility/sim_utils.py @@ -14,6 +14,8 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os import dexsim import open3d as o3d @@ -21,12 +23,14 @@ from typing import List, Union from dexsim.types import ( + CloneStrategy, DriveType, ArticulationFlag, LoadOption, + ObjectCloneOptions, RigidBodyShape, SDFConfig, - PhysicalAttr, + ActorType, ) from dexsim.engine import Articulation from dexsim.environment import Env, Arena @@ -46,6 +50,146 @@ 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 _newton_solver_type() -> str | None: + """Return the active Newton solver type, or None if unavailable.""" + try: + from embodichain.lab.sim.sim_manager import get_physics_scene + + mgr = getattr(get_physics_scene(), "manager", None) + if mgr is None: + return None + return getattr(getattr(mgr, "cfg", None), "solver_cfg", None).solver_type + except Exception: + return None + + +def _attach_newton_rigidbody_desc( + obj: MeshObject, + cfg: RigidObjectCfg, + body_type: ActorType, + shape_type: RigidBodyShape, +) -> None: + """Attach rigid-body physics via dexsim's Newton desc-native path. + + Used when ``cfg.attrs.newton`` is set on the Newton backend: builds the + resolved Newton shape descriptor (common fields projected + Newton-native + sub-config) and a ``RigidBodyPhysicsDesc`` body descriptor, populates the + ``mgr.dexsim_meta`` scaffolding that dexsim's registration/rebuild reads + (mirroring ``NewtonSpawnAdapter._attach_newton``), and registers via + ``register_mesh_object_to_newton_patch`` — fully bypassing the legacy + ``PhysicalAttr`` path so Newton-native contact/shape params reach the model. + Emits per-solver / backend-mismatch warnings. + """ + from embodichain.lab.sim.sim_manager import get_physics_scene + from dexsim.engine.newton_physics.rigid_body.registration import ( + register_mesh_object_to_newton_patch, + ) + from dexsim.engine.newton_physics.registry import _get_entity_native_handle + from embodichain.lab.sim.physics_attrs import ( + resolve_newton_body, + resolve_newton_shape, + warn_ignored_contact_fields, + warn_backend_mismatched_fields, + ) + + mgr = getattr(get_physics_scene(), "manager", None) + if mgr is None: + logger.log_error( + "Newton manager is unavailable; cannot attach rigid body via the " + "desc-native path." + ) + shape = resolve_newton_shape(cfg.attrs) + solver_type = _newton_solver_type() + if solver_type is not None: + warn_ignored_contact_fields(shape, solver_type) + warn_backend_mismatched_fields(cfg.attrs, "newton") + body = resolve_newton_body(cfg.attrs, body_type) + + # Populate the dexsim_meta scaffolding registration/rebuild read. This + # mirrors dexsim's NewtonSpawnAdapter._attach_newton meta dict so the body + # rebuilds correctly on the next finalize. + entity_handle = _get_entity_native_handle(obj) + arena = obj.get_arena() if hasattr(obj, "get_arena") else None + arena_handle = arena.get_native_handle() if arena is not None else -1 + mgr.dexsim_meta[entity_handle] = { + "actor_type": body_type, + "shape_type": shape_type, + "node_scale": np.asarray(obj.get_scale(), dtype=np.float32).reshape(-1)[:3], + "body_scale": np.asarray(obj.get_body_scale(), dtype=np.float32).reshape(-1)[ + :3 + ], + "arena_native_handle": arena_handle, + "newton_world_index": -1, + "newton_shape": shape, + "newton_body": body, + } + + register_mesh_object_to_newton_patch( + mgr, + obj, + body_type, + shape_type, + attr=None, + mesh_source_obj=obj, + newton_shape=shape, + newton_body=body, + ) + # Newton requires body scale after rigid-body creation. + _set_body_scale_after_rigidbody(obj, cfg.body_scale) + + +def _use_newton_desc_path(cfg: RigidObjectCfg) -> bool: + """Whether to route rigid-body spawn through the Newton desc-native path.""" + return _is_newton_backend_active() and cfg.attrs.newton is not None + + +def _newton_subcfg_has_fields(newton_cfg) -> bool: + """Return True if a Newton sub-config sets any field.""" + if newton_cfg is None: + return False + return any( + getattr(newton_cfg, f.name, None) is not None + for f in newton_cfg.__dataclass_fields__ + if f.name != "newton" + ) + + +def _warn_newton_articulation_native_attrs(cfg: "ArticulationCfg") -> None: + """Warn that Newton-native per-link contact params are not applied to articulations. + + dexsim's ``NewtonArticulation`` exposes no per-link contact-material setter + (ke/kd/margin/...), so the ``attrs.newton`` sub-config on an articulation is + accepted for config symmetry but cannot be applied per-link on Newton today. + Common fields (mass/friction/restitution/contact_offset) are still applied + via the legacy ``set_physical_attr`` path. + """ + sources = [] + if _newton_subcfg_has_fields(getattr(cfg.attrs, "newton", None)): + sources.append("attrs.newton") + for group_name, group_cfg in (cfg.link_attrs or {}).items(): + if _newton_subcfg_has_fields(getattr(group_cfg.attrs, "newton", None)): + sources.append(f"link_attrs['{group_name}'].attrs.newton") + if sources: + logger.log_warning( + "Newton-native per-link contact/shape params (" + ", ".join(sources) + ") " + "are not yet applied to articulation links on the Newton backend " + "(no dexsim per-link contact-material API). Common fields are applied." + ) + + def get_dexsim_arenas() -> List[dexsim.environment.Arena]: """Get all arenas in the default dexsim world. @@ -135,12 +279,155 @@ def _apply_link_physics_overrides( art.set_physical_attr(physical_attr, name, is_replace_inertial=replace_inertial) -def set_dexsim_articulation_cfg(arts: List[Articulation], cfg: ArticulationCfg) -> None: - """Set articulation configuration for a list of dexsim articulations. +def default_articulation_clone_options() -> ObjectCloneOptions: + """Return clone options used when duplicating articulations across arenas.""" + options = ObjectCloneOptions() + options.render.material = CloneStrategy.DEEP_COPY + return options + + +def default_rigid_object_clone_options() -> ObjectCloneOptions: + """Return clone options used when duplicating rigid actors across arenas.""" + options = ObjectCloneOptions() + options.render.material = CloneStrategy.DEEP_COPY + return options + + +def _clone_actor_between_arenas( + source_arena: Arena | Env, + source_name: str, + target_arena: Arena | Env, + target_name: str, + clone_options: ObjectCloneOptions, +) -> MeshObject: + """Clone a mesh actor from one arena/env to another.""" + return source_arena.clone_actor_to( + source_name, target_arena, target_name, clone_options + ) + + +def _clone_articulation_between_arenas( + source_arena: Arena | Env, + source_name: str, + target_arena: Arena | Env, + target_name: str, + clone_options: ObjectCloneOptions, +) -> Articulation: + """Clone an articulation from one arena/env to another.""" + if _is_newton_backend_active(): + return source_arena.clone_skeleton_to( + source_name, target_arena, target_name, clone_options + ) + return source_arena.clone_articulation_to( + source_name, target_arena, target_name, clone_options + ) + + +def spawn_articulation_entities( + cfg: ArticulationCfg, + env_list: list[Arena | Env], + *, + clone_options: ObjectCloneOptions | None = None, +) -> list[Articulation]: + """Load one articulation prototype and clone it into additional arenas. + + DexSim configuration is applied once on the prototype before cloning. + """ + if cfg.uid is None: + logger.log_error("Articulation uid must be set before spawning entities.") + + if clone_options is None: + clone_options = default_articulation_clone_options() + + source_env = env_list[0] + prototype_name = f"{cfg.uid}_0" + prototype = source_env.load_urdf(cfg.fpath) + prototype.set_name(prototype_name) + + if not cfg.use_usd_properties: + set_dexsim_articulation_cfg(prototype, cfg) + + entities = [prototype] + for env_idx in range(1, len(env_list)): + target_name = f"{cfg.uid}_{env_idx}" + clone = _clone_articulation_between_arenas( + source_env, + prototype_name, + env_list[env_idx], + target_name, + clone_options, + ) + if clone is None: + logger.log_error( + f"Failed to clone articulation '{prototype_name}' into env {env_idx}." + ) + entities.append(clone) + return entities + + +def _find_single_articulation_in_usd_import(results: dict, fpath: str) -> Articulation: + """Return the sole articulation imported from a USD file.""" + articulations_found = [ + value for value in results.values() if isinstance(value, Articulation) + ] + if len(articulations_found) == 0: + logger.log_error(f"No articulation found in USD file {fpath}.") + if len(articulations_found) > 1: + logger.log_error(f"Multiple articulations found in USD file {fpath}.") + return articulations_found[0] + + +def spawn_usd_articulation_entities( + cfg: ArticulationCfg, + env_list: list[Arena | Env], + *, + cache_dir: str | None = None, + clone_options: ObjectCloneOptions | None = None, +) -> list[Articulation]: + """Import one USD articulation prototype and clone it into additional arenas.""" + if cfg.uid is None: + logger.log_error("Articulation uid must be set before spawning entities.") + if len(env_list) == 0: + return [] + + if clone_options is None: + clone_options = default_articulation_clone_options() + + source_env = env_list[0] + prototype_name = f"{cfg.uid}_0" + results = source_env.import_from_usd_file( + cfg.fpath, return_object=True, cache_dir=cache_dir + ) + prototype = _find_single_articulation_in_usd_import(results, cfg.fpath) + prototype.set_name(prototype_name) + + if not cfg.use_usd_properties: + set_dexsim_articulation_cfg(prototype, cfg) + + entities = [prototype] + for env_idx in range(1, len(env_list)): + target_name = f"{cfg.uid}_{env_idx}" + clone = _clone_articulation_between_arenas( + source_env, + prototype_name, + env_list[env_idx], + target_name, + clone_options, + ) + if clone is None: + logger.log_error( + f"Failed to clone articulation '{prototype_name}' into env {env_idx}." + ) + entities.append(clone) + return entities + + +def set_dexsim_articulation_cfg(art: Articulation, cfg: ArticulationCfg) -> None: + """Apply EmbodiChain articulation cfg to a single DexSim articulation entity. Args: - arts (List[Articulation]): List of dexsim articulations to configure. - cfg (ArticulationCfg): Configuration object containing articulation settings. + art: DexSim articulation (or Newton skeleton carrier) to configure. + cfg: EmbodiChain articulation configuration. """ def get_drive_type(drive_pros): @@ -160,35 +447,46 @@ def get_drive_type(drive_pros): else: logger.log_error(f"Unknow drive type {drive_type}") - for i, art in enumerate(arts): + 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() + if is_newton_art: + for name in link_names: + art.set_physical_attr(cfg.attrs.attr(), name) + _warn_newton_articulation_native_attrs(cfg) + else: art.set_physical_attr(cfg.attrs.attr()) - 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 - ) + _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 + ) + 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: - physical_body = art.get_physical_body(name) - inertia = physical_body.get_mass_space_inertia_tensor() - inertia = np.maximum(inertia, 1e-4) - physical_body.set_mass_space_inertia_tensor(inertia) + 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) + physical_body.set_mass_space_inertia_tensor(inertia) - if i == 0 and cfg.compute_uv: - render_body = art.get_render_body(name) - if render_body: - render_body.set_projective_uv() + if cfg.compute_uv: + render_body = art.get_render_body(name) + if render_body: + render_body.set_projective_uv() - # TODO: will crash when exit if not explicitly delete. - # This may due to the destruction of render body order when exiting. - del render_body + # TODO: will crash when exit if not explicitly delete. + # This may due to the destruction of render body order when exiting. + del render_body def is_rt_enabled() -> bool: @@ -252,121 +550,241 @@ def create_sphere( return spheres -def load_mesh_objects_from_cfg( - cfg: RigidObjectCfg, env_list: List[Arena], cache_dir: str | None = None -) -> List[MeshObject]: - """Load mesh objects from configuration. - - Args: - cfg (RigidObjectCfg): Configuration for the rigid object. - env_list (List[Arena]): List of arenas to load the objects into. +def _mesh_load_option_from_cfg(cfg: RigidObjectCfg) -> LoadOption: + """Build DexSim mesh load options from a rigid-object configuration.""" + option = LoadOption() + option.rebuild_normals = cfg.shape.load_option.rebuild_normals + option.rebuild_tangent = cfg.shape.load_option.rebuild_tangent + option.rebuild_3rdnormal = cfg.shape.load_option.rebuild_3rdnormal + option.rebuild_3rdtangent = cfg.shape.load_option.rebuild_3rdtangent + option.smooth = cfg.shape.load_option.smooth + return option - cache_dir (str | None, optional): Directory for caching convex decomposition files. Defaults to None - Returns: - List[MeshObject]: List of loaded mesh objects. - """ - obj_list = [] - body_type = cfg.to_dexsim_body_type() - if isinstance(cfg.shape, MeshCfg): - option = LoadOption() - option.rebuild_normals = cfg.shape.load_option.rebuild_normals - option.rebuild_tangent = cfg.shape.load_option.rebuild_tangent - option.rebuild_3rdnormal = cfg.shape.load_option.rebuild_3rdnormal - option.rebuild_3rdtangent = cfg.shape.load_option.rebuild_3rdtangent - option.smooth = cfg.shape.load_option.smooth +def _apply_mesh_uv_mapping(obj: MeshObject, cfg: RigidObjectCfg) -> None: + """Compute and apply UV mapping for a mesh rigid-object prototype.""" + if not cfg.shape.compute_uv: + return - cfg: RigidObjectCfg - max_convex_hull_num = cfg.max_convex_hull_num - fpath = cfg.shape.fpath + vertices = obj.get_vertices() + triangles = obj.get_triangles() + o3d_mesh = o3d.t.geometry.TriangleMesh(vertices, triangles) + _, uvs = get_mesh_auto_uv(o3d_mesh, np.array(cfg.shape.project_direction)) + obj.set_uv_mapping(uvs) - compute_uv = cfg.shape.compute_uv - is_usd = fpath.endswith((".usd", ".usda", ".usdc")) - if is_usd: - # TODO: Currently add checking for num_envs when file is USD. After we support spawn via cloning, we can remove this. - if len(env_list) > 1: - logger.log_error(f"Currently not supporting multiple arenas for USD.") - _env: dexsim.environment.Env = dexsim.default_world().get_env() - results = _env.import_from_usd_file(fpath, return_object=True) - # print(f"import usd result: {results}") - - rigidbodys_found = [] - for key, value in results.items(): - if isinstance(value, MeshObject): - rigidbodys_found.append(value) - if len(rigidbodys_found) == 0: - logger.log_error(f"No rigid body found in USD file: {fpath}") - elif len(rigidbodys_found) > 1: - logger.log_error(f"Multiple rigid bodies found in USD file: {fpath}.") - elif len(rigidbodys_found) == 1: - obj_list.append(rigidbodys_found[0]) - return obj_list +def _configure_primitive_rigidbody( + obj: MeshObject, + cfg: RigidObjectCfg, + body_type, + *, + is_newton_backend: bool, + shape_type: RigidBodyShape, +) -> None: + """Attach primitive rigid-body physics to a cube or sphere prototype.""" + if is_newton_backend and cfg.attrs.newton is not None: + _attach_newton_rigidbody_desc(obj, cfg, body_type, shape_type) + return + if not is_newton_backend: + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody(body_type, shape_type, cfg.attrs.attr()) + if is_newton_backend: + _set_body_scale_after_rigidbody(obj, cfg.body_scale) + + +def _import_usd_rigid_prototype( + env: Arena | Env, + fpath: str, + prototype_name: str, +) -> MeshObject: + """Import a single rigid mesh actor from USD as the spawn prototype.""" + results = env.import_from_usd_file(fpath, return_object=True) + rigidbodys_found = [ + value for value in results.values() if isinstance(value, MeshObject) + ] + if len(rigidbodys_found) == 0: + logger.log_error(f"No rigid body found in USD file: {fpath}") + if len(rigidbodys_found) > 1: + logger.log_error(f"Multiple rigid bodies found in USD file: {fpath}.") + prototype = rigidbodys_found[0] + prototype.set_name(prototype_name) + return prototype + + +def _load_rigid_mesh_prototype( + env: Arena | Env, + cfg: RigidObjectCfg, + *, + cache_dir: str | None, + body_type, + is_newton_backend: bool, +) -> MeshObject: + """Load and configure one mesh rigid-object prototype in the source arena.""" + option = _mesh_load_option_from_cfg(cfg) + fpath = cfg.shape.fpath + max_convex_hull_num = cfg.max_convex_hull_num + + if max_convex_hull_num > 1: + obj = env.load_actor_with_coacd( + fpath, + duplicate=True, + attach_scene=True, + option=option, + cache_path=cache_dir, + actor_type=body_type, + 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(resolution=cfg.sdf_resolution) + obj.add_physical_body( + body_type, + RigidBodyShape.SDF, + config=sdf_cfg, + attr=cfg.attrs.attr(), + ) + else: + obj = env.load_actor(fpath, duplicate=True, attach_scene=True, option=option) + if is_newton_backend and cfg.attrs.newton is not None: + _attach_newton_rigidbody_desc(obj, cfg, body_type, RigidBodyShape.CONVEX) else: - # non-usd file does not support this option, will be forced set False to avoid potential issues. - cfg.use_usd_properties = False + obj.add_rigidbody(body_type, RigidBodyShape.CONVEX, cfg.attrs.attr()) + + _apply_mesh_uv_mapping(obj, cfg) + return obj + + +def _spawn_clones_from_prototype( + source_env: Arena | Env, + prototype_name: str, + env_list: list[Arena | Env], + uid: str, + clone_options: ObjectCloneOptions, +) -> list[MeshObject]: + """Return the prototype plus clones for all remaining arenas.""" + prototype = source_env.get_actor(prototype_name) + if prototype is None: + logger.log_error( + f"Rigid object prototype '{prototype_name}' was not found in the source arena." + ) - for i, env in enumerate(env_list): - if max_convex_hull_num > 1: - obj = env.load_actor_with_coacd( - fpath, - duplicate=True, - attach_scene=True, - option=option, - cache_path=cache_dir, - actor_type=body_type, - max_convex_hull_num=max_convex_hull_num, - ) - elif cfg.sdf_resolution > 0: - obj = env.load_actor( - fpath, duplicate=True, attach_scene=True, option=option - ) - sdf_cfg = SDFConfig() - sdf_cfg.resolution = cfg.sdf_resolution - obj.add_physical_body( - body_type, - RigidBodyShape.SDF, - config=sdf_cfg, - attr=PhysicalAttr(), - ) - else: - obj = env.load_actor( - fpath, duplicate=True, attach_scene=True, option=option - ) - obj.add_rigidbody(body_type, RigidBodyShape.CONVEX) - obj.set_name(f"{cfg.uid}_{i}") - obj_list.append(obj) + entities = [prototype] + for env_idx in range(1, len(env_list)): + target_name = f"{uid}_{env_idx}" + clone = _clone_actor_between_arenas( + source_env, + prototype_name, + env_list[env_idx], + target_name, + clone_options, + ) + if clone is None: + logger.log_error( + f"Failed to clone rigid object '{prototype_name}' into env {env_idx}." + ) + entities.append(clone) + return entities + + +def spawn_rigid_object_entities( + cfg: RigidObjectCfg, + env_list: list[Arena | Env], + *, + cache_dir: str | None = None, + clone_options: ObjectCloneOptions | None = None, +) -> list[MeshObject]: + """Load one rigid-object prototype and clone it into additional arenas. + + Mesh loading, convex decomposition, and physics setup run once on the + prototype in ``env_list[0]`` before cloning. + """ + if cfg.uid is None: + logger.log_error("Rigid object uid must be set before spawning entities.") + if len(env_list) == 0: + return [] - if compute_uv: - vertices = obj.get_vertices() - triangles = obj.get_triangles() + if clone_options is None: + clone_options = default_rigid_object_clone_options() - o3d_mesh = o3d.t.geometry.TriangleMesh(vertices, triangles) - _, uvs = get_mesh_auto_uv( - o3d_mesh, np.array(cfg.shape.project_direction) - ) - obj.set_uv_mapping(uvs) + body_type = cfg.to_dexsim_body_type() + is_newton_backend = _is_newton_backend_active() + source_env = env_list[0] + prototype_name = f"{cfg.uid}_0" + if isinstance(cfg.shape, MeshCfg): + fpath = cfg.shape.fpath + is_usd = fpath.endswith((".usd", ".usda", ".usdc")) + if is_usd: + prototype = _import_usd_rigid_prototype(source_env, fpath, prototype_name) + else: + cfg.use_usd_properties = False + prototype = _load_rigid_mesh_prototype( + source_env, + cfg, + cache_dir=cache_dir, + body_type=body_type, + is_newton_backend=is_newton_backend, + ) + prototype.set_name(prototype_name) elif isinstance(cfg.shape, CubeCfg): - from embodichain.lab.sim.utility.sim_utils import create_cube - - obj_list = create_cube(env_list, cfg.shape.size, uid=cfg.uid) - for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.BOX) - + prototype = source_env.create_cube( + cfg.shape.size[0], cfg.shape.size[1], cfg.shape.size[2] + ) + prototype.set_name(prototype_name) + _configure_primitive_rigidbody( + prototype, + cfg, + body_type, + is_newton_backend=is_newton_backend, + shape_type=RigidBodyShape.BOX, + ) elif isinstance(cfg.shape, SphereCfg): - from embodichain.lab.sim.utility.sim_utils import create_sphere - - obj_list = create_sphere( - env_list, cfg.shape.radius, cfg.shape.resolution, uid=cfg.uid + prototype = source_env.create_sphere(cfg.shape.radius, cfg.shape.resolution) + prototype.set_name(prototype_name) + _configure_primitive_rigidbody( + prototype, + cfg, + body_type, + is_newton_backend=is_newton_backend, + shape_type=RigidBodyShape.SPHERE, ) - for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.SPHERE) else: logger.log_error( - f"Unsupported rigid object shape type: {type(cfg.shape)}. Supported types: MeshCfg, CubeCfg, SphereCfg." + f"Unsupported rigid object shape type: {type(cfg.shape)}. " + "Supported types: MeshCfg, CubeCfg, SphereCfg." ) - return obj_list + return [] + + if len(env_list) == 1: + return [prototype] + return _spawn_clones_from_prototype( + source_env, prototype_name, env_list, cfg.uid, clone_options + ) + + +def load_mesh_objects_from_cfg( + cfg: RigidObjectCfg, env_list: List[Arena], cache_dir: str | None = None +) -> List[MeshObject]: + """Load mesh objects from configuration. + + Args: + cfg (RigidObjectCfg): Configuration for the rigid object. + env_list (List[Arena]): List of arenas to load the objects into. + + cache_dir (str | None, optional): Directory for caching convex decomposition files. Defaults to None + Returns: + List[MeshObject]: List of loaded mesh objects. + """ + return spawn_rigid_object_entities(cfg, env_list, cache_dir=cache_dir) def load_soft_object_from_cfg( 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 6b8b9eff..a940ebcc 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -26,6 +26,7 @@ from embodichain.lab.sim.solvers import PytorchSolverCfg 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/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 6c06c467..f3063b5c 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, @@ -52,8 +53,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 ddea8b77..3b2b680d 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 d8601703..af62f5ff 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -25,7 +25,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 @@ -41,7 +45,7 @@ def main(): ) add_env_launcher_args_to_parser(parser) parser.add_argument( - "--record-steps", + "--max_steps", type=int, default=1000, help="Number of simulation steps to record before exiting in headless mode.", @@ -66,7 +70,8 @@ def main(): height=1080, headless=args.headless, 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, ), @@ -83,8 +88,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, @@ -101,11 +107,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, ) ) @@ -130,19 +137,13 @@ def main(): print( "[INFO]: The output path is reported by `SimulationManager.start_window_record()`." ) - print(f"[INFO]: Running {args.record_steps} steps before exporting the video") + print(f"[INFO]: Running {args.max_steps} steps before exporting the video") # Run the simulation - run_simulation( - sim, - max_steps=args.record_steps if args.headless else None, - ) + run_simulation(sim, max_steps=args.max_steps) -def run_simulation( - sim: SimulationManager, - max_steps: int | None = None, -): +def run_simulation(sim: SimulationManager, max_steps: int | None = None): """Run the simulation loop. Args: @@ -164,6 +165,9 @@ def run_simulation( 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_differentiable_env.py b/tests/gym/envs/test_differentiable_env.py new file mode 100644 index 00000000..13104597 --- /dev/null +++ b/tests/gym/envs/test_differentiable_env.py @@ -0,0 +1,113 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Tests for DifferentiableEmbodiedEnv.""" + +from __future__ import annotations + +import pytest +import torch + +from embodichain.lab.gym.envs.differentiable_env import ( + DifferentiableEmbodiedEnv, +) +from embodichain.lab.gym.envs.embodied_env import EmbodiedEnvCfg +from embodichain.lab.sim.cfg import DefaultPhysicsCfg, NewtonPhysicsCfg +from embodichain.lab.sim.sim_manager import SimulationManagerCfg + + +def _diff_env_cfg( + requires_grad: bool = True, backend: str = "newton" +) -> EmbodiedEnvCfg: + if backend == "newton": + physics_cfg = NewtonPhysicsCfg( + requires_grad=requires_grad, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ) + else: + physics_cfg = DefaultPhysicsCfg() + sim_cfg = SimulationManagerCfg( + physics_cfg=physics_cfg, + num_envs=2, + headless=True, + ) + return EmbodiedEnvCfg(sim_cfg=sim_cfg) + + +def test_construct_without_requires_grad_raises(): + with pytest.raises(Exception, match=r"requires_grad"): + DifferentiableEmbodiedEnv(_diff_env_cfg(requires_grad=False)) + + +def test_construct_on_default_backend_raises(): + with pytest.raises(Exception, match=r"Newton"): + DifferentiableEmbodiedEnv(_diff_env_cfg(backend="default")) + + +def _import_franka_env(): + """Import the Franka APG env, skipping if the URDF is unavailable. + + The URDF resolves through ``newton.utils.download_asset`` which + requires network access on first run. Tests skip cleanly when the + asset cannot be fetched. + """ + from embodichain.lab.gym.envs.tasks.special.franka_reach_apg import ( + FrankaReachApgEnv, + ) + + return FrankaReachApgEnv + + +def test_franka_apg_smoke_backward(): + """Verify reward is autograd-tracked and action.grad flows back.""" + try: + FrankaReachApgEnv = _import_franka_env() + except FileNotFoundError as e: + pytest.skip(f"Franka URDF not available: {e}") + + env = FrankaReachApgEnv(num_envs=2) + env.reset(seed=0) + action = torch.zeros(2, 7, requires_grad=True, device=env.device) + obs, reward, terminated, truncated, info = env.step(action) + assert reward.requires_grad, "Reward must be autograd-tracked." + loss = reward.sum() + loss.backward() + assert action.grad is not None + assert torch.isfinite(action.grad).all() + + +def test_franka_apg_one_iter_loss_reduces(): + """Verify a single SGD step reduces the APG loss.""" + try: + FrankaReachApgEnv = _import_franka_env() + except FileNotFoundError as e: + pytest.skip(f"Franka URDF not available: {e}") + + env = FrankaReachApgEnv(num_envs=2) + env.reset(seed=0) + action = torch.zeros(2, 7, requires_grad=True, device=env.device) + opt = torch.optim.SGD([action], lr=0.01) + + losses = [] + for _ in range(3): + env.reset(seed=0) + opt.zero_grad() + _, reward, _, _, _ = env.step(action) + loss = (-reward).sum() + loss.backward() + opt.step() + losses.append(loss.detach().item()) + assert losses[-1] < losses[0], f"APG did not reduce loss: {losses}" 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..8c85f067 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,22 @@ 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"): + physics_cfg = physics_cfg_for_backend(physics) + if physics == "newton": + physics_cfg.solver_cfg = { + "solver_type": "mujoco_warp", + "njmax": 8192, + "nconmax": 8192, + } config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + headless=True, + device=device, + num_envs=NUM_ARENAS, + physics_cfg=physics_cfg, ) self.sim = SimulationManager(config) + self.physics = physics art_path = get_data_path(ART_PATH) assert os.path.isfile(art_path) @@ -91,8 +109,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 +333,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 +446,88 @@ 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() + + def test_set_link_physical_attr_mass_live_on_newton(self): + """Per-link mass set via set_link_physical_attr takes effect live on Newton. + + On Newton, ``set_physical_attr`` is metadata-only; the fix pushes mass + live via ``set_link_mass`` (mirroring the dedicated set_mass). Verify a + runtime per-link mass override round-trips through get_mass. + """ + link_name = self.art.link_names[0] + original = self.art.get_mass(link_names=[link_name])[0, 0].item() + new_mass = original + 1.5 + self.art.set_link_physical_attr( + RigidBodyAttributesOverrideCfg(mass=new_mass), + link_names=[link_name], + ) + live_mass = self.art.get_mass(link_names=[link_name])[0, 0].item() + assert ( + abs(live_mass - new_mass) < 1e-3 + ), f"per-link mass {new_mass} not applied live on Newton (got {live_mass})" + + 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..991b84c9 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -13,23 +13,24 @@ # 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.cfg import NewtonCollisionAttributesCfg 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 +39,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 +103,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.""" @@ -93,6 +122,14 @@ def test_is_static(self): not self.chair.is_static ), "Chair should be kinematic but is marked static" + def test_spawn_clones_distinct_entities(self): + """Multi-env rigid objects are spawned via prototype + clone_actor_to.""" + assert len(self.duck._entities) == NUM_ARENAS + handles = {entity.get_native_handle() for entity in self.duck._entities} + assert len(handles) == NUM_ARENAS, "Each arena clone must be a distinct actor" + assert self.duck._entities[0].get_name() == "duck_0" + assert self.duck._entities[1].get_name() == "duck_1" + def test_local_pose_behavior(self): """Test set_local_pose and get_local_pose: - duck pose is correctly set @@ -158,9 +195,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 +443,86 @@ 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 applies the Newton-supported subset (mass, friction, + # restitution, contact_offset) at runtime and mirrors the rest. + self.duck.set_attrs( + RigidBodyAttributesCfg(mass=2.5, dynamic_friction=0.7, restitution=0.4) + ) + assert torch.allclose( + self.duck.get_mass(), + torch.full((NUM_ARENAS,), 2.5, device=self.sim.device), + atol=1e-5, + ), "Newton set_attrs(mass) did not apply via batch API" + assert torch.allclose( + self.duck.get_friction(), + torch.full((NUM_ARENAS,), 0.7, device=self.sim.device), + atol=1e-5, + ), "Newton set_attrs(dynamic_friction) did not apply via batch API" + + # set_body_type is a runtime no-op on Newton (body type is fixed at + # registration); the call must not change body_type. + 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 is a runtime no-op on Newton (not modelled per body) but + # mirrors onto metadata so get_damping stays consistent. + new_damping = torch.full((NUM_ARENAS, 2), 0.2, device=self.sim.device) + self.duck.set_damping(new_damping) + assert torch.allclose( + self.duck.get_damping(), new_damping, atol=1e-5 + ), "Newton set_damping should mirror onto metadata for get_damping" + + 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 +600,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 = _make_test_com_pose(self.sim.device) - 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 == ( + 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 +725,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 +970,63 @@ 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() + + def test_newton_native_attrs_desc_native_spawn(self): + """RigidObject with attrs.newton spawns via the desc-native path on Newton. + + Setting ``attrs.newton`` routes spawn through + ``register_mesh_object_to_newton_patch`` (bypassing legacy PhysicalAttr), + so Newton-native contact/shape params reach the model. Verifies the + body is registered with the Newton manager after finalize. + """ + duck_path = get_data_path(DUCK_PATH) + cfg = RigidObjectCfg( + uid="duck_newton_native", + shape=MeshCfg(fpath=duck_path), + body_type="dynamic", + attrs=RigidBodyAttributesCfg( + mass=1.0, + dynamic_friction=0.5, + restitution=0.1, + newton=NewtonCollisionAttributesCfg(ke=1e3, kd=50.0, margin=0.01), + ), + ) + obj: RigidObject = self.sim.add_rigid_object(cfg=cfg) + self.sim.finalize_newton_physics() + + assert obj.num_instances == NUM_ARENAS + assert obj.body_type == "dynamic" + # The body must be registered with the Newton manager post-finalize. + mgr = self.sim.newton_manager + assert mgr is not None + assert mgr.registered_body_count() > 0 + # Common fields round-trip via the batch view (mass applied live). + assert torch.allclose( + obj.get_mass(), + torch.full((NUM_ARENAS,), 1.0, device=self.sim.device), + atol=1e-5, + ) + + @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..3cfb48e4 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -22,6 +22,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.dexforce_w1 import DexforceW1Cfg +from embodichain.lab.sim.cfg import physics_cfg_for_backend from embodichain.data import get_data_path # Define control parts @@ -50,11 +51,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 +69,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): @@ -331,6 +332,68 @@ def setup_method(self): self.setup_simulation("cuda") +def _teardown_newton_physics() -> None: + from dexsim.engine.newton_physics import teardown_newton_physics + + teardown_newton_physics() + + +class TestRobotNewton: + """Focused Robot-on-Newton coverage (spawn, finalize, control surface). + + A robot is a URDF articulation; the Newton ``load_urdf`` patch builds a + NewtonArticulation. This exercises the add_robot -> finalize_newton_physics + -> control-part / qpos path end-to-end on Newton. It does NOT inherit the + full BaseRobotTest suite because rebuilding the (complex, mimic-jointed) + dexforce_w1 Newton model per test method is prohibitively slow; the + default/CUDA classes already cover the shared control-part/FK/IK logic. + """ + + def setup_method(self): + physics_cfg = physics_cfg_for_backend("newton") + physics_cfg.solver_cfg = { + "solver_type": "mujoco_warp", + "njmax": 8192, + "nconmax": 8192, + } + config = SimulationManagerCfg( + headless=True, device="cuda", num_envs=1, physics_cfg=physics_cfg + ) + self.sim = SimulationManager(config) + cfg = DexforceW1Cfg.from_dict( + {"uid": "dexforce_w1", "version": "v021", "arm_kind": "anthropomorphic"} + ) + self.robot: Robot = self.sim.add_robot(cfg=cfg) + self.sim.finalize_newton_physics() + + 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_newton_robot_spawn_and_control(self): + """Robot spawns on Newton, finalizes, and exposes a working control surface.""" + assert self.sim.is_newton_backend + assert self.sim.physics._lifecycle_state() == "READY" + assert self.robot.dof > 0 + + left_ids = self.robot.get_joint_ids("left_arm") + right_ids = self.robot.get_joint_ids("right_arm") + assert len(left_ids) > 0 and len(right_ids) > 0 + + # State round-trip via the Newton articulation view. + qpos = torch.zeros( + (1, self.robot.dof), dtype=torch.float32, device=self.sim.device + ) + self.robot.set_qpos(qpos, env_ids=None, target=False) + assert torch.allclose(self.robot.body_data.qpos, qpos, atol=1e-5) + + if __name__ == "__main__": # Run tests directly test_cpu = TestRobotCUDA() 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..8d8f95d1 100644 --- a/tests/sim/objects/test_usd.py +++ b/tests/sim/objects/test_usd.py @@ -32,21 +32,21 @@ from embodichain.lab.sim.shapes import MeshCfg from embodichain.data import get_data_path -NUM_ARENAS = 1 +NUM_ARENAS = 2 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): @@ -73,6 +73,9 @@ def test_import_rigid(self): default_attr.min_position_iters, default_attr.min_velocity_iters, ) + assert len(sugar_box._entities) == NUM_ARENAS + handles = {entity.get_native_handle() for entity in sugar_box._entities} + assert len(handles) == NUM_ARENAS def test_import_articulation(self): default_drive = JointDrivePropertiesCfg() @@ -178,13 +181,13 @@ def teardown_method(self): gc.collect() -@pytest.mark.skip(reason="Skipping CUDA tests temporarily") +# @pytest.mark.skip(reason="Skipping CUDA tests temporarily") class TestUsdCPU(BaseUsdTest): def setup_method(self): self.setup_simulation("cpu") -@pytest.mark.skip(reason="Skipping CUDA tests temporarily") +# @pytest.mark.skip(reason="Skipping CUDA tests temporarily") class TestUsdCUDA(BaseUsdTest): def setup_method(self): self.setup_simulation("cuda") 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_backend_parity.py b/tests/sim/test_backend_parity.py new file mode 100644 index 00000000..0d550cec --- /dev/null +++ b/tests/sim/test_backend_parity.py @@ -0,0 +1,182 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Backend capability parity matrix. + +This is the single source of truth for which simulation features each physics +backend supports. It pins the capability contract so that: + +- flipping a ``supports_*`` flag (or adding a backend) fails loudly, and +- every ``SimulationManager.add_*`` capability guard maps 1:1 to its flag. + +Headless (no GPU / no dexsim world): backends are constructed with a minimal +fake owning-manager back-ref, and the ``add_*`` guard mapping is exercised by +binding a fake ``physics`` onto a bare ``SimulationManager`` via +``object.__new__`` (mirroring the lifecycle-test pattern). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from embodichain.lab.sim.physics import ( + DefaultPhysicsBackend, + NewtonPhysicsBackend, + PhysicsBackend, +) +from embodichain.lab.sim.sim_manager import SimulationManager + +# --------------------------------------------------------------------------- +# The parity matrix — edit this table when a backend gains/loses a feature. +# --------------------------------------------------------------------------- +# feature -> {backend -> supported} +BACKEND_CAPABILITIES: dict[str, dict[str, bool]] = { + "robot": {"default": True, "newton": True}, + "soft_bodies": {"default": True, "newton": False}, + "cloth": {"default": True, "newton": False}, + "rigid_object_group": {"default": True, "newton": False}, + "can_disable_manual_update": {"default": True, "newton": False}, +} + +BACKENDS: dict[str, type[PhysicsBackend]] = { + "default": DefaultPhysicsBackend, + "newton": NewtonPhysicsBackend, +} + +# Map each capability flag to the SimulationManager.add_* method whose +# NotImplementedError guard consults it. ``None`` means the flag is consulted +# elsewhere (e.g. set_manual_update) rather than an add_* guard. +CAPABILITY_TO_ADD_METHOD: dict[str, str | None] = { + "robot": "add_robot", + "soft_bodies": "add_soft_object", + "cloth": "add_cloth_object", + "rigid_object_group": "add_rigid_object_group", + "can_disable_manual_update": None, +} + + +def _make_backend(name: str) -> PhysicsBackend: + """Construct a backend with a minimal fake owning-manager back-ref.""" + return BACKENDS[name](SimpleNamespace()) + + +@pytest.mark.parametrize("backend_name", list(BACKENDS)) +def test_backend_name_matches(backend_name: str) -> None: + backend = _make_backend(backend_name) + assert backend.name == backend_name + + +@pytest.mark.parametrize("backend_name", list(BACKENDS)) +@pytest.mark.parametrize( + "feature", [f for f in BACKEND_CAPABILITIES if f != "can_disable_manual_update"] +) +def test_supports_flags_match_matrix(backend_name: str, feature: str) -> None: + """Each backend's supports_* property matches the parity matrix.""" + backend = _make_backend(backend_name) + expected = BACKEND_CAPABILITIES[feature][backend_name] + actual = getattr(backend, f"supports_{feature}") + assert ( + actual is expected + ), f"{backend_name}.supports_{feature} = {actual}, matrix says {expected}" + + +@pytest.mark.parametrize("backend_name", list(BACKENDS)) +def test_can_disable_manual_update_matches_matrix(backend_name: str) -> None: + backend = _make_backend(backend_name) + expected = BACKEND_CAPABILITIES["can_disable_manual_update"][backend_name] + assert backend.can_disable_manual_update is expected + + +def _make_sim_with_backend(backend: PhysicsBackend) -> SimulationManager: + """Build a bare SimulationManager whose ``physics`` is the given backend. + + The add_* capability guards consult only ``self.physics.supports_*`` (plus a + few uid/existence checks that run after the guard), so a bare instance with + ``physics`` + the registries set is enough to assert the guard fires. + """ + sim = object.__new__(SimulationManager) + sim.physics = backend + sim._soft_objects = {} + sim._cloth_objects = {} + sim._rigid_object_groups = {} + sim._robots = {} + sim._rigid_objects = {} + sim._articulations = {} + return sim + + +@pytest.mark.parametrize( + "feature,add_method", + [(f, m) for f, m in CAPABILITY_TO_ADD_METHOD.items() if m is not None], +) +@pytest.mark.parametrize("backend_name", list(BACKENDS)) +def test_add_method_guard_maps_to_capability( + backend_name: str, feature: str, add_method: str +) -> None: + """add_ raises NotImplementedError iff the backend lacks the flag. + + For unsupported features the guard must fire before any world access; for + supported features the method proceeds past the guard (and is expected to + fail later on the missing world — we only assert it does NOT raise + NotImplementedError at the guard). + """ + backend = _make_backend(backend_name) + sim = _make_sim_with_backend(backend) + supported = BACKEND_CAPABILITIES[feature][backend_name] + method = getattr(sim, add_method) + + # Minimal cfg stub: add_* only reads .uid before/after the guard. + cfg = SimpleNamespace(uid=None) + + if supported: + # Past the guard it will hit missing-world attrs; assert the failure is + # NOT the capability NotImplementedError. + with pytest.raises(Exception) as exc_info: + method(cfg=cfg) + assert not isinstance(exc_info.value, NotImplementedError), ( + f"{add_method} raised NotImplementedError on the {backend_name} " + f"backend despite supports_{feature}=True" + ) + assert "not enabled" not in str(exc_info.value) + else: + with pytest.raises(NotImplementedError, match="not enabled"): + method(cfg=cfg) + + +def test_matrix_covers_all_capability_flags() -> None: + """Every supports_* / can_disable_manual_update flag is in the matrix.""" + flag_names = { + name[len("supports_") :] if name.startswith("supports_") else name + for name in dir(PhysicsBackend) + if name.startswith("supports_") or name == "can_disable_manual_update" + } + matrix_features = set(BACKEND_CAPABILITIES) + assert ( + flag_names == matrix_features + ), f"capability flags {flag_names} != matrix features {matrix_features}" + + +def test_matrix_covers_all_backends() -> None: + """Every concrete backend class is in the matrix.""" + # Discover concrete (non-abstract) backends by instantiation. + concrete = set(BACKENDS) + matrix_backends = {b for feats in BACKEND_CAPABILITIES.values() for b in feats} + assert concrete == matrix_backends + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) 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_differentiable_stepper.py b/tests/sim/test_differentiable_stepper.py new file mode 100644 index 00000000..a628e87d --- /dev/null +++ b/tests/sim/test_differentiable_stepper.py @@ -0,0 +1,100 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Tests for the differentiable-stepper delegators on SimulationManager.""" + +from __future__ import annotations + +import pytest + +from embodichain.lab.sim.cfg import DefaultPhysicsCfg, NewtonPhysicsCfg +from embodichain.lab.sim.sim_manager import SimulationManager, SimulationManagerCfg + + +def test_default_backend_rejects_differentiable_stepper(): + sim = SimulationManager( + SimulationManagerCfg( + physics_cfg=DefaultPhysicsCfg(), + num_envs=1, + headless=True, + ) + ) + with pytest.raises(Exception, match=r"Newton"): + sim.create_differentiable_stepper() + SimulationManager.reset() + + +def test_newton_without_grad_rejects_differentiable_stepper(): + sim = SimulationManager( + SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + requires_grad=False, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=1, + headless=True, + ) + ) + sim.finalize_newton_physics() + with pytest.raises(Exception, match=r"grad"): + sim.create_differentiable_stepper() + SimulationManager.reset() + + +def test_newton_with_grad_creates_stepper(): + sim = SimulationManager( + SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + requires_grad=True, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=1, + headless=True, + ) + ) + sim.finalize_newton_physics() + stepper = sim.create_differentiable_stepper() + from dexsim.engine.newton_physics.differentiable_stepper import ( + DifferentiableStepper, + ) + + assert isinstance(stepper, DifferentiableStepper) + SimulationManager.reset() + + +def test_tape_context_records_step(): + import warp as wp + + sim = SimulationManager( + SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + requires_grad=True, + solver_cfg={"solver_type": "semi_implicit"}, + use_cuda_graph=False, + ), + num_envs=1, + headless=True, + ) + ) + sim.finalize_newton_physics() + from embodichain.lab.sim.diff import tape_context + + with tape_context(sim) as tape: + pass # empty tape is valid; tape.backward() on empty is a no-op + + assert isinstance(tape, wp.Tape) + SimulationManager.reset() diff --git a/tests/sim/test_newton_finalize_lifecycle.py b/tests/sim/test_newton_finalize_lifecycle.py new file mode 100644 index 00000000..3b2adefd --- /dev/null +++ b/tests/sim/test_newton_finalize_lifecycle.py @@ -0,0 +1,198 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Unit tests for the Newton physics backend finalize/invalidate lifecycle. + +These tests exercise :class:`NewtonPhysicsBackend` in isolation (no GPU and no +live dexsim world required) by injecting a fake Newton manager and patching the +``ensure_simulation_prepared_lazy`` rebuild entry point. They verify the +backend owns the dirty/finalize state machine that used to live inline in +:class:`SimulationManager`. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import patch + +from embodichain.lab.sim.physics import NewtonPhysicsBackend + + +class _Resettable: + """Stand-in for a RigidObject/Articulation with a reset() call counter.""" + + def __init__(self) -> None: + self.reset_calls = 0 + + def reset(self) -> None: + self.reset_calls += 1 + + +class _FakeNewtonManager: + """Stand-in for dexsim's NewtonManager exposing only the lifecycle state.""" + + def __init__(self) -> None: + self.lifecycle_state = SimpleNamespace(name="BUILDER") + + +def _make_backend() -> tuple[ + NewtonPhysicsBackend, + _FakeNewtonManager, + _Resettable, + _Resettable, + _Resettable, + _Resettable, +]: + rigid_obj = _Resettable() + rigid_group = _Resettable() # groups must NOT be reset by the Newton backend. + articulation = _Resettable() + robot = _Resettable() # a robot is an articulation and is reset like one. + newton_mgr = _FakeNewtonManager() + + # Minimal owning-SimulationManager stand-in: only the attributes the backend + # touches during finalize / reset are needed. + manager = SimpleNamespace( + _world=object(), + _rigid_objects={"rigid": rigid_obj}, + _rigid_object_groups={"rigid_group": rigid_group}, + _articulations={"art": articulation}, + _robots={"robot": robot}, + ) + + backend = NewtonPhysicsBackend(manager) + # Inject the fake manager so finalize() does not call get_newton_manager. + backend._newton_manager = newton_mgr + return backend, newton_mgr, rigid_obj, rigid_group, articulation, robot + + +def _fake_ensure_prepared_lazy(mgr, world, *, rebuild_from_scene, warn): + """Mimic the real rebuild: bring the Newton model to the READY state.""" + mgr.lifecycle_state.name = "READY" + return True, None + + +@patch( + "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", + new=_fake_ensure_prepared_lazy, +) +def test_finalize_resets_entities_after_ready() -> None: + ( + backend, + newton_mgr, + rigid_obj, + rigid_group, + articulation, + robot, + ) = _make_backend() + + assert not backend.is_initialized + backend.prepare() + + assert newton_mgr.lifecycle_state.name == "READY" + assert backend.is_initialized + assert rigid_obj.reset_calls == 1 + assert articulation.reset_calls == 1 + assert robot.reset_calls == 1 + # Rigid object groups are not supported on the Newton backend: not reset. + assert rigid_group.reset_calls == 0 + + +@patch( + "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", + new=_fake_ensure_prepared_lazy, +) +def test_finalize_does_not_repeat_deferred_reset() -> None: + ( + backend, + _newton_mgr, + rigid_obj, + _rigid_group, + articulation, + robot, + ) = _make_backend() + + backend.prepare() + backend.prepare() + + assert rigid_obj.reset_calls == 1 + assert articulation.reset_calls == 1 + assert robot.reset_calls == 1 + + +@patch( + "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", + new=_fake_ensure_prepared_lazy, +) +def test_invalidation_allows_next_finalize_to_reset_again() -> None: + ( + backend, + _newton_mgr, + rigid_obj, + _rigid_group, + articulation, + robot, + ) = _make_backend() + + backend.prepare() + backend.invalidate() + assert not backend.is_initialized + backend.prepare() + + assert rigid_obj.reset_calls == 2 + assert articulation.reset_calls == 2 + assert robot.reset_calls == 2 + + +@patch( + "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", + new=_fake_ensure_prepared_lazy, +) +def test_finalize_raises_when_rebuild_unsafe() -> None: + backend, _newton_mgr, rigid_obj, _rigid_group, _articulation, _robot = ( + _make_backend() + ) + + # An unsafe rebuild makes finalize() raise (logger.log_error raises by + # default). It must not mark itself initialized nor reset entities. + with patch( + "dexsim.engine.newton_physics.rebuild.ensure_simulation_prepared_lazy", + new=lambda mgr, world, *, rebuild_from_scene, warn: (False, None), + ): + try: + backend.prepare() + except RuntimeError: + pass + else: # pragma: no cover - defensive + raise AssertionError("finalize() should raise on an unsafe rebuild") + + assert not backend.is_initialized + assert rigid_obj.reset_calls == 0 + + +def test_invalidate_is_idempotent_and_only_clears_finalized_flag() -> None: + ( + backend, + _newton_mgr, + _rigid_obj, + _rigid_group, + _articulation, + _robot, + ) = _make_backend() + backend._is_finalized = True + + backend.invalidate() + backend.invalidate() + + assert not backend.is_initialized diff --git a/tests/sim/test_physics_attrs.py b/tests/sim/test_physics_attrs.py new file mode 100644 index 00000000..77652538 --- /dev/null +++ b/tests/sim/test_physics_attrs.py @@ -0,0 +1,202 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Headless unit tests for the backend-aware rigid-body attribute resolver. + +No GPU / dexsim world required — these exercise the config layer and the +``physics_attrs`` resolver/warning logic in isolation. +""" + +from __future__ import annotations + +import logging + +import pytest + +from embodichain.lab.sim.cfg import ( + NewtonCollisionAttributesCfg, + RigidBodyAttributesCfg, + RigidBodyAttributesOverrideCfg, +) +from embodichain.lab.sim.physics_attrs import ( + NEWTON_CONTACT_SOLVER_FIELDS, + ResolvedNewtonShape, + resolve_newton_body, + resolve_newton_shape, + resolve_rigid_body_attributes, + warn_backend_mismatched_fields, + warn_ignored_contact_fields, +) + + +def test_from_dict_parses_nested_newton() -> None: + cfg = RigidBodyAttributesCfg.from_dict( + {"mass": 2.0, "restitution": 0.3, "newton": {"ke": 1e3, "margin": 0.01}} + ) + assert cfg.mass == 2.0 + assert cfg.restitution == 0.3 + assert isinstance(cfg.newton, NewtonCollisionAttributesCfg) + assert cfg.newton.ke == 1e3 + assert cfg.newton.margin == 0.01 + # unset newton fields stay None + assert cfg.newton.kd is None + + +def test_override_from_dict_parses_nested_newton() -> None: + ov = RigidBodyAttributesOverrideCfg.from_dict({"newton": {"kd": 50.0}}) + assert isinstance(ov.newton, NewtonCollisionAttributesCfg) + assert ov.newton.kd == 50.0 + assert ov.newton.ke is None + + +def test_resolve_newton_shape_projects_common_fields() -> None: + cfg = RigidBodyAttributesCfg( + mass=2.0, + dynamic_friction=0.4, + restitution=0.2, + enable_collision=False, + density=800.0, + newton=NewtonCollisionAttributesCfg(ke=1e3, margin=0.01), + ) + shape = resolve_newton_shape(cfg) + assert isinstance(shape, ResolvedNewtonShape) + # common fields projected onto Newton ShapeConfig knobs + assert shape.mu == 0.4 # dynamic_friction -> mu + assert shape.restitution == 0.2 + assert shape.has_shape_collision is False # enable_collision -> has_shape_collision + assert shape.density == 800.0 # positive, so dexsim computes a positive body mass + # newton-native sub-config fields copied verbatim + assert shape.ke == 1e3 + assert shape.margin == 0.01 + # unset newton-native fields stay None + assert shape.kd is None + + +def test_resolve_newton_shape_without_subconfig() -> None: + cfg = RigidBodyAttributesCfg(dynamic_friction=0.5, restitution=0.1) + shape = resolve_newton_shape(cfg) + assert shape.mu == 0.5 + assert shape.restitution == 0.1 + assert shape.has_shape_collision is True # default enable_collision + assert shape.ke is None # no newton sub-config + + +def test_resolve_newton_body_carries_mass_and_density() -> None: + from dexsim.types import ActorType + + cfg = RigidBodyAttributesCfg(mass=2.0, density=800.0) + body = resolve_newton_body(cfg, ActorType.DYNAMIC) + assert body.actor_type == ActorType.DYNAMIC + assert body.mass == 2.0 + assert body.density == 800.0 + + +def test_resolve_rigid_body_attributes_dispatches_by_backend() -> None: + cfg = RigidBodyAttributesCfg(mass=2.0, newton=NewtonCollisionAttributesCfg(ke=1e3)) + # default backend -> legacy PhysicalAttr + pa = resolve_rigid_body_attributes(cfg, "default") + assert pa.mass == 2.0 + # newton backend -> resolved shape + shape = resolve_rigid_body_attributes(cfg, "newton", solver_type=None) + assert isinstance(shape, ResolvedNewtonShape) + assert shape.ke == 1e3 + + +def test_merge_with_propagates_newton_via_merged_cfg() -> None: + base = RigidBodyAttributesCfg( + mass=1.0, newton=NewtonCollisionAttributesCfg(ke=1e3, margin=0.01) + ) + override = RigidBodyAttributesOverrideCfg( + mass=3.0, newton=NewtonCollisionAttributesCfg(kd=50.0) + ) + merged = override.merged_cfg(base) + # override wins for mass + assert merged.mass == 3.0 + # newton sub-config: override non-None wins, else base + assert merged.newton.ke == 1e3 # from base (override None) + assert merged.newton.kd == 50.0 # from override + assert merged.newton.margin == 0.01 # from base + # legacy merge_with still returns a PhysicalAttr (drops newton) + pa = override.merge_with(base) + assert pa.mass == 3.0 + + +def test_warn_ignored_contact_fields_xpbd(caplog) -> None: + shape = ResolvedNewtonShape(ke=1e3, kd=50.0, mu=0.5, restitution=0.2) + with caplog.at_level(logging.WARNING): + warn_ignored_contact_fields(shape, "xpbd") + # xpbd reads {mu, restitution, mu_torsional, mu_rolling}; ke/kd ignored + msg = caplog.text + assert "xpbd" in msg + assert "ke" in msg and "kd" in msg + + +def test_warn_ignored_contact_fields_mujoco_warp_no_ke_kd_warning( + caplog, +) -> None: + shape = ResolvedNewtonShape(ke=1e3, kd=50.0, mu=0.5) + with caplog.at_level(logging.WARNING): + warn_ignored_contact_fields(shape, "mujoco_warp") + # mujoco_warp reads {ke, kd, mu, kh, mu_torsional, mu_rolling}; ke/kd NOT ignored + assert "ke" not in caplog.text or "ignores" not in caplog.text + + +def test_warn_ignored_contact_fields_restitution_on_mujoco_warp( + caplog, +) -> None: + # mujoco_warp does NOT read restitution -> should warn + shape = ResolvedNewtonShape(restitution=0.3, mu=0.5) + with caplog.at_level(logging.WARNING): + warn_ignored_contact_fields(shape, "mujoco_warp") + assert "restitution" in caplog.text + + +def test_warn_backend_mismatched_fields_newton(caplog) -> None: + # PhysX-only fields deviating from defaults on Newton -> warn + cfg = RigidBodyAttributesCfg(enable_ccd=True, linear_damping=0.9) + with caplog.at_level(logging.WARNING): + warn_backend_mismatched_fields(cfg, "newton") + msg = caplog.text + assert "enable_ccd" in msg + assert "linear_damping" in msg + + +def test_warn_backend_mismatched_fields_no_warn_for_defaults(caplog) -> None: + # all defaults -> no warning + cfg = RigidBodyAttributesCfg() + with caplog.at_level(logging.WARNING): + warn_backend_mismatched_fields(cfg, "newton") + assert caplog.text == "" + + +def test_warn_backend_mismatched_fields_no_warn_on_default(caplog) -> None: + cfg = RigidBodyAttributesCfg(enable_ccd=True) + with caplog.at_level(logging.WARNING): + warn_backend_mismatched_fields(cfg, "default") + assert caplog.text == "" + + +def test_newton_contact_solver_fields_table_sanity() -> None: + # union of per-solver read sets == NEWTON_CONTACT_FIELDS + from embodichain.lab.sim.physics_attrs import NEWTON_CONTACT_FIELDS + + union = set() + for fields_set in NEWTON_CONTACT_SOLVER_FIELDS.values(): + union |= set(fields_set) + assert union == set(NEWTON_CONTACT_FIELDS) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) 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)