diff --git a/.agents/skills/add-robot/SKILL.md b/.agents/skills/add-robot/SKILL.md index ee7236ddb..67bce0f31 100644 --- a/.agents/skills/add-robot/SKILL.md +++ b/.agents/skills/add-robot/SKILL.md @@ -17,7 +17,7 @@ Every robot config subclasses `RobotCfg` and overrides two hooks: - `_build_defaults(self, init_dict=None)` — read variant fields from `init_dict`, set them on `self`, then populate `urdf_cfg` / `control_parts` / `solver_cfg` / - `drive_pros` / `attrs`. + `joint_drive_props` / `attrs`. - `build_pk_serial_chain(self, device=...)` — return `{control_part: pk.SerialChain}`, reading the PK URDF from a single `_pk_urdf_path` source. @@ -46,7 +46,7 @@ A cfg's `_build_defaults` must populate: - `urdf_cfg` (URDFCfg) or `fpath` - `control_parts` (Dict[str, List[str]]; joint names support regex) - `solver_cfg` (Dict[str, SolverCfg]; keys match `control_parts`) -- `drive_pros` (JointDrivePropertiesCfg) +- `joint_drive_props` (JointDrivePropertiesCfg) - `attrs` (RigidBodyAttributesCfg) `build_pk_serial_chain` must read from `_pk_urdf_path` (a property for @@ -69,7 +69,7 @@ must match the matching `control_parts` entry (the test stub asserts this). self.urdf_cfg = URDFCfg(components=[...]) self.control_parts = {"arm": ["JOINT[1-6]"]} self.solver_cfg = {"arm": OPWSolverCfg(end_link_name="link6", root_link_name="base_link")} - self.drive_pros = JointDrivePropertiesCfg(stiffness={"JOINT[1-6]": 1e4}) + self.joint_drive_props = JointDrivePropertiesCfg(stiffness={"JOINT[1-6]": 1e4}) ``` Variant-aware template (reads version / arm_kind): @@ -79,7 +79,7 @@ must match the matching `control_parts` entry (the test stub asserts this). init_dict = init_dict or {} self.version = MyRobotVersion(init_dict.get("version", "v1")) self.arm_kind = MyRobotArmKind(init_dict.get("arm_kind", "default")) - ... # then urdf_cfg / control_parts / solver_cfg / drive_pros / attrs + ... # then urdf_cfg / control_parts / solver_cfg / joint_drive_props / attrs ``` 4. **Implement `build_pk_serial_chain`** reading from `_pk_urdf_path`: @@ -137,7 +137,7 @@ must match the matching `control_parts` entry (the test stub asserts this). | `urdf_cfg` | URDFCfg | URDF file and components | | `control_parts` | Dict[str, List[str]] | Joint groups for control | | `solver_cfg` | Dict[str, SolverCfg] | IK solver configurations | -| `drive_pros` | JointDrivePropertiesCfg | Joint stiffness, damping, force | +| `joint_drive_props` | JointDrivePropertiesCfg | Joint drive, limits, friction, and armature | | `attrs` | RigidBodyAttributesCfg | Rigid-body physics attributes | | variant fields | enum / str / bool | Optional subclass fields | | `_pk_urdf_path` | property or method → str | URDF for the FK/IK serial chain | diff --git a/.agents/skills/add-solver/SKILL.md b/.agents/skills/add-solver/SKILL.md index 115f46f36..9507d1038 100644 --- a/.agents/skills/add-solver/SKILL.md +++ b/.agents/skills/add-solver/SKILL.md @@ -219,7 +219,7 @@ exactly: `test_ur_solver.py`) to sample joint configs within limits with a safety margin. - A `BaseSolverTest` class with: - - `setup_simulation(self, sim_device)` — builds a `SimulationManagerCfg`, + - `setup_simulation(self, device)` — builds a `SimulationManagerCfg`, a `RobotCfg` whose `solver_cfg={"arm": SolverCfg(...)}` uses the new solver, and adds the robot via `self.sim.add_robot(cfg=cfg)`. - `test_ik(self)` — the round-trip contract: diff --git a/.agents/skills/add-test/SKILL.md b/.agents/skills/add-test/SKILL.md index 72bdd27e2..0e48aa3dc 100644 --- a/.agents/skills/add-test/SKILL.md +++ b/.agents/skills/add-test/SKILL.md @@ -83,7 +83,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg class TestMySimComponent: def setup_method(self): - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) # ... setup ... diff --git a/agent_context/MAP.yaml b/agent_context/MAP.yaml index ff8aa5405..0e60fe285 100644 --- a/agent_context/MAP.yaml +++ b/agent_context/MAP.yaml @@ -21,7 +21,15 @@ topics: - simulation - SimulationManager - SimulationManagerCfg + - PhysicsBackendCfg + - DefaultPhysicsCfg - DexSim + - AutoSolverCfg + - Newton AutoSolver + - solver_cfg + - enable_multiccd + - cone + - impratio - world - arena - physics step @@ -31,19 +39,45 @@ topics: - manual update - GPU physics - simulation lifecycle + - collision_policy + - collision isolation + - arena isolation + - JointDrivePropertiesCfg + - joint_drive_props + - root_props + - AssetPhysicsMode + - asset_physics_mode + - RigidBodyPhysicsCfg + - recompute_inertia + - MeshCollisionCfg + - mesh collision approximation + - shape.collision + - target_mode + - drive_type + - mimic joint - ArticulationJointKinematics - get_parent_joint_chain + - quaternion + - xyzw + - pose convention paths: - topics/simulation-system/simulation-system.md source_of_truth: - embodichain/lab/sim/__init__.py - embodichain/lab/sim/sim_manager.py - - embodichain/lab/sim/cfg.py + - embodichain/lab/sim/cfg/ + - embodichain/lab/sim/physics/ - embodichain/lab/sim/common.py - embodichain/lab/sim/material.py - embodichain/lab/sim/profiler.py + - embodichain/lab/sim/spawn/descriptors.py + - embodichain/lab/sim/spawn/scene.py - embodichain/lab/sim/objects/__init__.py - embodichain/lab/sim/objects/articulation.py + - embodichain/lab/sim/objects/rigid_object.py + - embodichain/lab/sim/objects/backends/ + - embodichain/lab/sim/objects/deformable/ + - embodichain/utils/math.py - embodichain/lab/sim/sensors/__init__.py - embodichain/lab/sim/solvers/__init__.py - embodichain/lab/sim/planners/__init__.py @@ -234,21 +268,34 @@ topics: keywords: - robot - RobotCfg + - RobotPresetCfg - Robot - control - drive + - JointDrivePropertiesCfg + - joint_drive_props + - root_props + - AssetPhysicsMode + - asset_physics_mode + - target_mode + - drive_type - joint + - mimic joint - urdf - cobotmagic - dexforce_w1 - gripper - arm + - quaternion + - xyzw paths: - topics/robot-system/robot-system.md source_of_truth: + - embodichain/lab/gym/envs/embodied_env.py - embodichain/lab/sim/objects/robot.py + - embodichain/lab/sim/sim_manager.py - embodichain/lab/sim/robots/ - - embodichain/lab/sim/cfg.py + - embodichain/lab/sim/cfg/ related_topics: - simulation-system - ik-solvers @@ -279,6 +326,8 @@ topics: - rgb - depth - pointcloud + - quaternion + - xyzw paths: - topics/sensor-system/sensor-system.md source_of_truth: @@ -325,6 +374,8 @@ topics: - deformable - soft body - cloth + - quaternion + - xyzw paths: - topics/sim-visualization/sim-visualization.md source_of_truth: @@ -337,6 +388,7 @@ topics: - embodichain/lab/sim/objects/rigid_object_group.py - embodichain/lab/sim/objects/soft_object.py - embodichain/lab/sim/objects/cloth_object.py + - embodichain/lab/sim/objects/deformable/ related_topics: - simulation-system - env-framework @@ -386,6 +438,9 @@ topics: - collision_world_batch_mode - collision_geometry_by_id - make_planning_scene_provider + - quaternion + - xyzw + - wxyz paths: - topics/motion-planning/motion-planning.md source_of_truth: @@ -548,6 +603,38 @@ 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 + - quaternion + - xyzw + paths: + - topics/differentiable-env/differentiable-env.md + source_of_truth: + - embodichain/lab/gym/envs/differentiable_env.py + - embodichain/lab/sim/diff/ + - embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py + related_topics: + - env-framework + - rl-learning + status: active - id: atomic-actions title: Atomic Actions @@ -888,6 +975,8 @@ topics: - pour water - 配置生成环境 - 运行时环境注册 + - SemanticPose + - quaternion_xyzw paths: - topics/expert-programs/expert-programs.md source_of_truth: diff --git a/agent_context/topics/atomic-actions/atomic-actions.md b/agent_context/topics/atomic-actions/atomic-actions.md index 83a084269..db7a3e380 100644 --- a/agent_context/topics/atomic-actions/atomic-actions.md +++ b/agent_context/topics/atomic-actions/atomic-actions.md @@ -805,6 +805,10 @@ structured invalidation/replan events, and requires terminal completion. The dynamic-obstacle example additionally uses dense `morphit` robot collision spheres, proves that the moved cuboid intersects the original TCP path, and requires the replanned TCP path to retain a positive minimum clearance. +The `place.py` tutorial authors matching Newton contact stiffness and damping +on the cube and gripper collision links before `prepare()`. MuJoCo-Warp's +default response is too compliant for this force-closure replay and otherwise +lets the cube slip near its pickup pose instead of reaching the place target. Semantic integration tutorials live under `scripts/tutorials/semantic_skill/`. Both examples separate `create_*_application()` (scene/profile/runtime and 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 000000000..dc58718b3 --- /dev/null +++ b/agent_context/topics/differentiable-env/differentiable-env.md @@ -0,0 +1,116 @@ +# 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`. + +Newton/Warp `body_q` transforms contain position followed by a native `xyzw` +quaternion. This already matches EmbodiChain's quaternion convention, so the +differentiable bridge and FK reward path must not reorder those four +components. The Franka target pose likewise uses `xyz + xyzw`, with identity +orientation `(0, 0, 0, 1)`. + +## 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_tasks.special.franka_reach_apg` 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_tasks/embodichain_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-learning diff --git a/agent_context/topics/expert-programs/expert-programs.md b/agent_context/topics/expert-programs/expert-programs.md index 457eedcd7..4f76a04d4 100644 --- a/agent_context/topics/expert-programs/expert-programs.md +++ b/agent_context/topics/expert-programs/expert-programs.md @@ -55,6 +55,12 @@ to its `ParallelCfg`; it is not a standalone program node. Nested parallel blocks are rejected. Built-in call configs are `PickCfg`, `PlaceCfg`, and `HandOverCfg`; `RegisteredSemanticCallCfg` is the explicit catalog extension. +Expert Program poses follow the EmbodiChain quaternion contract. `PoseCfg` and +serialized target poses require the key `quaternion_xyzw`; `SemanticPose` +stores and reports the same order. The configured hand-over service uses +`final_quaternion_xyzw`. Identity is `[0, 0, 0, 1]`; legacy `*_wxyz` keys are +unknown fields and are rejected rather than silently reinterpreted. + Both programmatic config construction and untrusted decoding enforce exact types, discriminators, references, finite numeric values, and bounded depth, node count, repeat count, and expanded call count. The loader additionally diff --git a/agent_context/topics/motion-planning/motion-planning.md b/agent_context/topics/motion-planning/motion-planning.md index 9d167a58a..3ecdc3015 100644 --- a/agent_context/topics/motion-planning/motion-planning.md +++ b/agent_context/topics/motion-planning/motion-planning.md @@ -97,6 +97,12 @@ Learning-based EEF waypoint planner. Franka Panda only. ### CuroboPlanner collision worlds +EmbodiChain planner inputs and robot FK results use `xyz + xyzw`. CuRobo's +native pose representation uses `xyz + wxyz`; `curobo_planner.py` and +`curobo_yaml.py` perform that conversion exactly once when constructing CuRobo +goals and obstacle YAML. Dynamic obstacle inputs expressed as homogeneous +matrices do not need a quaternion-order convention until that boundary. + `CuroboWorldCfg.rigid_objects` accepts either a mapping or a sequence. Use `Mapping[registry_id, RigidObject]` for a registry-backed integration. The mapping key is the authoritative logical/source obstacle ID used by the diff --git a/agent_context/topics/randomization/randomization.md b/agent_context/topics/randomization/randomization.md index 057206007..7f2ec8b91 100644 --- a/agent_context/topics/randomization/randomization.md +++ b/agent_context/topics/randomization/randomization.md @@ -25,12 +25,20 @@ The `__init__.py` of the randomization package re-exports everything via `from . | Function | Target | Key params | |---|---|---| -| `randomize_rigid_object_mass` | `RigidObject` mass | `mass_range`, `relative` | +| `randomize_rigid_object_mass` | Dynamic `RigidObject` mass/inertia | `mass_range`, `relative`, `recompute_inertia`, `min_mass` | | `randomize_rigid_object_center_of_mass` | `RigidObject` CoM offset | `com_pos_offset_range` | -| `randomize_articulation_mass` | `Articulation` link masses | `mass_range` (uniform or per-link dict), `link_names` (regex), `relative` | - -- `relative=True` adds sampled value to the initial/default mass instead of replacing. -- `randomize_articulation_mass` supports a `dict[str, tuple]` for per-link ranges; when used, `link_names` is ignored. +| `randomize_articulation_mass` | `Articulation` link mass/inertia | `mass_range` (uniform or per-link dict), `link_names` (regex), `relative`, `recompute_inertia`, `min_mass` | + +- `relative=True` adds the sampled value to the backend-resolved initial mass + stored in the target object's `default_mass` snapshot; repeated calls + therefore do not accumulate for either rigid objects or articulations. +- Rigid-object and articulation mass samples are clamped to positive + `min_mass`. By default, inertia is recomputed from the corresponding + initialization snapshot using the mass ratio; set `recompute_inertia=False` + only when inertia is managed separately. +- Non-dynamic rigid objects are skipped with a warning. +- `randomize_articulation_mass` supports a `dict[str, tuple]` for per-link + ranges; when used, `link_names` is ignored. - Link names are resolved via `resolve_matching_names` (regex matching). ### Visual (`visual.py`) @@ -171,13 +179,16 @@ Used in `params` to reference simulation objects by `uid`. The manager resolves ### Sampling -All randomizers use `embodichain.utils.math.sample_uniform(lower, upper, size)` for uniform sampling. +Randomizers use `embodichain.utils.math.sample_uniform(...)` for uniform +sampling where applicable. Physics samples are allocated on the target object's +device, not assumed to share `env.device`. ## Common Failure Modes | Symptom | Likely cause | |---|---| | Randomizer silently does nothing | `entity_cfg.uid` not found in `sim.get_rigid_object_uid_list()` — all randomizers early-return on UID mismatch | +| Rigid-object mass is clamped | The sampled absolute mass or relative result was below positive `min_mass` | | `ValueError` on link name | `mass_range` dict key doesn't match any `articulation.link_names` | | Camera randomization error | Extrinsics config has neither `parent` nor `eye` set — unsupported mode | | Light randomization not per-env | By design: `randomize_light` applies same values across all envs | diff --git a/agent_context/topics/robot-system/robot-system.md b/agent_context/topics/robot-system/robot-system.md index bf7e4ce9a..faeb78d5e 100644 --- a/agent_context/topics/robot-system/robot-system.md +++ b/agent_context/topics/robot-system/robot-system.md @@ -5,10 +5,13 @@ | What | Path | |---|---| | Robot runtime class | `embodichain/lab/sim/objects/robot.py` → `Robot` | -| RobotCfg base config | `embodichain/lab/sim/cfg.py` → `RobotCfg` (line ~1455) | -| ArticulationCfg parent | `embodichain/lab/sim/cfg.py` → `ArticulationCfg` (line ~1345) | -| JointDrivePropertiesCfg | `embodichain/lab/sim/cfg.py` → `JointDrivePropertiesCfg` (line ~654) | +| RobotCfg base config | `embodichain/lab/sim/cfg/robot.py` → `RobotCfg` | +| Replace-only backend preset | `embodichain/lab/sim/cfg/robot.py` → `RobotPresetCfg` | +| Environment robot declaration | `embodichain/lab/gym/envs/embodied_env.py` → `EmbodiedEnvCfg.robot` | +| ArticulationCfg parent | `embodichain/lab/sim/cfg/articulation.py` → `ArticulationCfg` | +| Joint drive/dynamics config | `embodichain/lab/sim/cfg/articulation.py` → `JointDrivePropertiesCfg` | | Robot registry (all robots) | `embodichain/lab/sim/robots/__init__.py` | +| Robot executable smoke entry points | Each specified robot module's ``__main__`` block | | DexforceW1 config package | `embodichain/lab/sim/robots/dexforce_w1/` | | CobotMagic config | `embodichain/lab/sim/robots/cobotmagic.py` | | Add-robot tutorial | `docs/source/tutorial/add_robot.rst` | @@ -23,16 +26,20 @@ A `Robot` is instantiated with a `RobotCfg` and a list of DexSim `Articulation` entities. +Robot FK/IK and end-effector pose APIs use the EmbodiChain convention: +quaternions are `xyzw`, and 7D poses are `xyz + xyzw`. Solver or planner +adapters convert only when their external library uses another order. + ## RobotCfg Pattern Inheritance chain: ``` ObjectBaseCfg uid, init_pos, init_rot, init_local_pose - └─ ArticulationCfg fpath, drive_pros, attrs, link_attrs, fix_base, - │ disable_self_collision, init_qpos, body_scale, - │ build_pk_chain, use_usd_properties - └─ RobotCfg control_parts, urdf_cfg, solver_cfg, drive_pros (override default to "force") + └─ ArticulationCfg fpath, joint_drive_props, attrs, link_attrs, root_props, + │ init_qpos, qpos_limits, body_scale, build_pk_chain, + │ asset_physics_mode + └─ RobotCfg control_parts, urdf_cfg, solver_cfg, joint_drive_props (position+velocity force default) ├─ DexforceW1Cfg version, hand_versions, with_default_eef └─ CobotMagicCfg (dual-arm defaults) ``` @@ -44,8 +51,10 @@ Key fields on `RobotCfg`: | `control_parts` | `Dict[str, List[str]] \| None` | Part name → joint names (supports regex like `JOINT[1-6]`) | | `urdf_cfg` | `URDFCfg \| None` | Multi-component URDF assembly (e.g. left_arm + right_arm) | | `solver_cfg` | `SolverCfg \| Dict[str, SolverCfg] \| None` | IK solver config; dict keys must match `control_parts` keys | -| `drive_pros` | `JointDrivePropertiesCfg` | Default drive type is `"force"` (overrides Articulation's `"none"`) | -| `attrs` | `RigidBodyAttributesCfg` | Rigid-body physics attributes (mass, friction, damping, ...) | +| `joint_drive_props` | `JointDrivePropertiesCfg` | Single joint-property entry point for target mode, gains, effort/velocity limits, passive friction, and armature. Robot supplies the established `drive_type="force"`; unspecified fields remain source-owned | +| `asset_physics_mode` | `AssetPhysicsMode` | Robot defaults to `overlay`; generic articulations default to `preserve` | +| `attrs` | `RigidBodyPhysicsCfg` | Grouped rigid-body physics. Flat attribute keys are rejected; COM quaternions use `xyzw`. | +| `root_props` | `ArticulationRootPropertiesCfg` | Sole root-property interface. Fixed-base/self-collision are portable; root sleep and paired solver-iteration fields are Default-only | | variant fields | `enum \| str \| bool` | Optional subclass fields (e.g. `version`, `with_default_eef`) | | `_pk_urdf_path` | `property \| method → str` | URDF for the FK/IK serial chain (one source, so it can't drift from sim) | @@ -64,7 +73,8 @@ def from_dict(cls, init_dict): - **`_build_defaults(self, init_dict=None)`** — read variant fields from `init_dict`, set them on `self`, then populate `urdf_cfg`, `control_parts`, `solver_cfg`, - `drive_pros` and `attrs`. (Base `RobotCfg._build_defaults` is a no-op.) + `joint_drive_props` and `attrs`. (Base + `RobotCfg._build_defaults` is a no-op.) - **`build_pk_serial_chain(self, device=...)`** — return `{control_part: pk.SerialChain}`, reading the PK URDF from a single `_pk_urdf_path` source (a property for constant-path robots, a method when the path depends on a variant). @@ -78,6 +88,35 @@ restores it. Every config, including this exception, must satisfy `type(cfg).from_dict(cfg.to_dict())` without changing the selected components or applying a derived transform twice. +### Physics backend portability + +Keep backend-neutral intent in one ordinary `RobotCfg`. In particular, +`CollisionPropertiesCfg.contact_offset/rest_offset` compile directly to +Default and to Newton's `margin=rest_offset`, +`gap=contact_offset-rest_offset`. Use `DefaultCollisionPropertiesCfg` only as a +Default-native extension point; those two inherited fields are portable. +Default-only articulation sleep and solver iterations belong directly in +`ArticulationRootPropertiesCfg` under `root_props`; `sleep_threshold`, +`min_position_iters`, and `min_velocity_iters` no longer exist as flat +`ArticulationCfg` fields. EmbodiChain applies these values to the Default-native +articulation root before the first reset, while Newton ignores them. Use +`DefaultRigidBodyPropertiesCfg` under `attrs` or +`link_attrs` only when the intended target is an individual rigid body/link. +Keep portable rigid-body values and one selected backend subtype in the single +matching `RigidBodyPhysicsCfg` slot. Backend-specific whole-robot alternatives +belong in `RobotPresetCfg`; there are no coexisting per-property backend blocks. + +When a backend truly needs a different asset or complete actuator/physics +definition, subclass `RobotPresetCfg` and declare complete alternatives. The +required `default` field selects the Default backend and is the Newton fallback; +optional names include `newton`, `newton_mujoco_warp`/`newton_mjwarp`, and other +`newton_` profiles. `SimulationManager.add_robot()` selects from its +existing `physics_cfg` and active Newton solver, returns a deep-copied complete +`RobotCfg`, and never merges fields across alternatives. `EmbodiedEnvCfg.robot` +accepts either form and delegates selection to that same boundary. Prefer a +single portable `RobotCfg`; use a preset only for irreducible backend +differences. + W1 robot and hand releases use separate types and registries: - `DexforceW1Version` selects body/arm assets, kinematics, and flange calibration @@ -115,28 +154,88 @@ control_parts = { - `Robot.get_link_names(name)` returns child link names for a part. - Internal `ControlGroup` dataclass stores `joint_names`, `joint_ids`, `link_names` per part. +For Spawn-bound robots, control-part IDs are resolved by name against the +final batch `qpos` layout. Newton may use a different source-articulation +traversal order, so do not derive control-part IDs by enumerating native joint +names. `init_qpos` keeps its source-articulation order and is remapped by name +when the robot resets. + +### Mimic joints across physics backends + +`Articulation.mimic_ids` and `mimic_parents` use the final batch-state joint +order, just like `qpos`, `qvel`, and control-part IDs. Spawn source metadata is +normalized by joint name before these properties are exposed. Newton initial +positions are also projected onto each URDF relation +`child = multiplier * parent + offset` before the first simulation step. + +Newton's MuJoCo-Warp solver currently represents URDF mimic joints as rigid +equality constraints without Default's compliance control. For mimic parents +that have a position drive, the Spawn-bound articulation installs the internal +Newton soft-mimic adapter from `objects/backends/newton.py`. It disables only +the articulation's native mimic rows, copies each parent's gains to the child, +and mirrors parent position/velocity targets into the child. This avoids +corrective-impulse instability on light finger links while retaining the +authored multiplier and offset. Passive/velocity-only mechanisms, gradient +mode, other Newton solvers, and the Default backend keep their native mimic +constraints. + ## Drive Properties -`JointDrivePropertiesCfg` controls the physics drive for joints: +`JointDrivePropertiesCfg` is the single joint-property config: | Field | Type | Default | Notes | |---|---|---|---| -| `drive_type` | `"force" \| "acceleration" \| "none"` | `"force"` (on RobotCfg) | `"none"` means no applied force | +| `drive_type` | `"force" \| "acceleration" \| "none"` | `"force"` (on RobotCfg) | Original drive response; active `"acceleration"` is Default-only | +| `target_mode` | `"none" \| "position" \| "velocity" \| "position_velocity" \| "effort"` or per-joint mapping | Derived from `drive_type` | Portable actuator intent; integer values 0–4 are accepted. `force` defaults to `position_velocity` | | `stiffness` | `float \| Dict[str, float]` | `1e4` | Per-joint via dict; keys support regex | | `damping` | `float \| Dict[str, float]` | `1e3` | Same | -| `max_effort` | `float \| Dict[str, float]` | `1e10` | Max torque/force | -| `max_velocity` | `float \| Dict[str, float]` | `1e10` | rad/s or m/s | -| `friction` | `float \| Dict[str, float]` | `0.0` | Joint friction | +| `max_effort` | `float \| Dict[str, float]` | `None` | Max torque/force | +| `max_velocity` | `float \| Dict[str, float]` | `None` | rad/s or m/s | +| `friction` | `float \| Dict[str, float]` | `None` | Passive joint friction | +| `armature` | `float \| Dict[str, float]` | `None` | Added joint-space inertia | When using a dict, keys are joint names or regex patterns matching joint names. Control-part names can also be used as keys (resolved via `ArticulationCfg` logic). +Target mode is backend-neutral and belongs directly on +`JointDrivePropertiesCfg`. Default emulates the target selection with its drive +mode and effective gains; Newton authors `JointTargetMode` values for +`"none"`, `"position"`, `"velocity"`, `"position_velocity"`, and +`"effort"` (integer values 0–4). `NewtonJointDrivePropertiesCfg` remains only +to round-trip older `joint_drive_props.backend: newton` dictionaries; do not use it in +new specified robots. + +`drive_type` retains its original meaning. With no explicit `target_mode`, +`force` and `acceleration` select `position_velocity`, while `none` selects a +passive target. An explicit target mode overrides that target default. Active +acceleration drives are rejected on Newton because Newton has no equivalent +mass-independent response. + +For solver-independent safety, `none` and `effort` clear Kp/Kd, while +`velocity` clears Kp. MuJoCo Warp consumes the target-mode enum natively. Other +Newton solvers use the gain fallback; their position-only fallback assumes the +velocity target remains zero. Direct generalized effort continues through +`Articulation.set_qf()` and can also act as feed-forward effort with an active +PD drive. + +These rules are resolved to exact joint names after URDF/USD source resolution +and before Spawn finalization. Common effort/velocity/armature values are +authored on `JointDesc`; the portable target intent lowers to Default drive +mode/gains and Newton's integer target mode. The dual-arm builder preserves the +config type and mirrors regex-keyed values to the generated `left_`/`right_` +names. + +`qpos_limits` accepts either joint-name/regex rules or a flattened +`(num_dofs, 2)` array. Both forms are resolved into common `JointDesc` +limits before the Default or Newton model is built; do not add a post-bind +Newton rebuild for initial limits. + ## Adding a New Robot Full guide: `docs/source/tutorial/add_robot.rst` · Quick reference: `docs/source/guides/add_robot.rst` Minimal checklist: 1. Create a `@configclass` inheriting `RobotCfg`. -2. Override `_build_defaults(self, init_dict=None)` — read variant fields from `init_dict`, then populate `urdf_cfg`, `control_parts`, `solver_cfg`, `drive_pros` and `attrs`. +2. Override `_build_defaults(self, init_dict=None)` — read variant fields from `init_dict`, then populate `urdf_cfg`, `control_parts`, `solver_cfg`, `joint_drive_props` and `attrs`. 3. Keep `from_dict` as the 3-line template (`cls()` → `_build_defaults` → `merge_robot_cfg`) unless version-derived state requires an explicitly documented post-merge step. 4. Define `control_parts` mapping part names to joint name lists. 5. Configure `solver_cfg` (one `SolverCfg` per control part). @@ -155,16 +254,25 @@ custom-transform, component-version, and public-builder round-trips. | Robot | Config Class | Module | Structure | Notes | |---|---|---|---|---| | DexForce W1 | `DexforceW1Cfg` | `embodichain/lab/sim/robots/dexforce_w1/` | Package (`cfg.py`, `types.py`, `specs.py`, `hand_specs.py`, `params.py`, `utils.py`) | Humanoid; robot and hand versions are independently registered | -| CobotMagic | `CobotMagicCfg` | `embodichain/lab/sim/robots/cobotmagic.py` | Single file | Dual-arm; 6-DOF arms + 2-DOF grippers; uses OPW solver | +| CobotMagic | `CobotMagicCfg` | `embodichain/lab/sim/robots/cobotmagic.py` | Single file | Dual-arm; 6-DOF arms + 2-DOF grippers; portable collision envelope, Default-native root iterations, OPW solver | + +## Executable smoke programs + +Every specified robot module accepts ``--physics {default,newton}`` in its +``__main__`` smoke program and resolves the selection through +``physics_cfg_for_backend()``. CobotMagic, Franka, UR, and DualArm retain the +Default backend as their command-line default; DexforceW1 retains Newton as its +default. These entry points exercise the same ordinary ``RobotCfg`` definitions +on either backend rather than maintaining backend-specific demo configs. ## Common Failure Modes - **`solver_cfg` keys don't match `control_parts` keys** — solver init silently uses wrong part or errors at IK time. - **Regex joint names not expanded** — if robot is not properly initialized, regex patterns like `JOINT[1-6]` remain unexpanded. Always construct via `from_dict()` or let `Robot.__init__` handle expansion. -- **`drive_type="none"` inherited from ArticulationCfg** — if you inherit `ArticulationCfg` directly instead of `RobotCfg`, the default drive type is `"none"` (no forces applied). Override to `"force"`. +- **No drive config on generic `ArticulationCfg`** — its `joint_drive_props=None` keeps source drives. Use `RobotCfg` for the standard position+velocity force-drive defaults, or provide an explicit sparse drive overlay. - **Missing `urdf_cfg` for multi-component robots** — single-file robots use `fpath`; multi-component robots (e.g. dual-arm) require `urdf_cfg` with component transforms. - **Mimic joints not excluded** — `get_joint_ids(remove_mimic=False)` includes mimic joints by default. Pass `remove_mimic=True` for active-only joints. -- **`init_qpos` shape mismatch** — must be `(num_joints,)`. A wrong-length array causes silent truncation or index errors at sim start. +- **`init_qpos` shape mismatch** — must match active DOFs. A wrong-length array causes initialization errors. - **`all` instead of `__all__`** — lowercase `all` does not work with `from module import *`; use `__all__`. - **`solver_cfg` set in multiple places** — set it once in `_build_defaults` only; setting it elsewhere (e.g. a build helper) gets overwritten and is dead code. - **PK URDF drifts from the sim URDF** — route `build_pk_serial_chain` through `_pk_urdf_path` and keep the DOF drift-guard test so silent drift is caught. diff --git a/agent_context/topics/sensor-system/sensor-system.md b/agent_context/topics/sensor-system/sensor-system.md index 23a25c244..bc8e16fd2 100644 --- a/agent_context/topics/sensor-system/sensor-system.md +++ b/agent_context/topics/sensor-system/sensor-system.md @@ -51,7 +51,7 @@ Defines the sensor pose relative to its parent frame: | Field | Type | Default | Notes | |---|---|---|---| | `pos` | `Tuple[float, float, float]` | `(0, 0, 0)` | Position in parent frame | -| `quat` | `Tuple[float, float, float, float]` | `(1, 0, 0, 0)` | Orientation as `(w, x, y, z)` quaternion | +| `quat` | `Tuple[float, float, float, float]` | `(0, 0, 0, 1)` | Orientation as `(x, y, z, w)` quaternion | | `parent` | `str \| None` | `None` | Parent frame name (e.g. robot link); `None` = arena frame | The `transformation` property returns a `4×4 torch.Tensor` homogeneous matrix. @@ -62,6 +62,14 @@ The `transformation` property returns a `4×4 torch.Tensor` homogeneous matrix. ## Camera System +`Camera` and `StereoCamera` are created through +`SimulationManager.add_sensor()`. The owning manager is passed explicitly so +each camera resolves its World and ordered per-environment Arenas through that +manager even when multiple simulation managers are active. The manager also +owns semantic parent resolution and deferred attachment; cameras only attach +to concrete per-environment render nodes and report attachment after that +operation succeeds. + ### CameraCfg | Field | Type | Default | Notes | @@ -117,7 +125,7 @@ Properties `left_to_right` and `right_to_left` return `4×4` transform tensors. - **`sensor_type` string mismatch** — `SensorCfg.from_dict()` looks up `sensor_type + "Cfg"` in the sensors module. A typo (e.g. `"camera"` instead of `"Camera"`) causes `AttributeError`. - **Depth not enabled** — `enable_depth` defaults to `False`. Accessing depth data without enabling it returns empty tensors. -- **Parent frame not found** — `OffsetCfg.parent` must exactly match a link name in the scene. A wrong name silently places the sensor at the arena origin. +- **Parent frame not found** — `OffsetCfg.parent` must match a link name in a Spawn-bound robot or articulation. Missing or ambiguous names raise during immediate attachment or `SimulationManager.prepare()`. - **Stereo baseline sign** — `left_to_right_pos` defines translation from left to right camera. Flipping the sign inverts the disparity. - **Contact sensor buffer overflow** — `max_contacts_per_env` caps the contact count. Exceeding it silently drops contacts; increase if the scene has dense collisions. - **View attribute flags** — `Camera.get_view_attrib()` computes `dr.ViewFlags` from enabled booleans. Adding a new data type requires both the `enable_*` flag and the corresponding `ViewFlags` bit. diff --git a/agent_context/topics/sim-visualization/sim-visualization.md b/agent_context/topics/sim-visualization/sim-visualization.md index 763e25979..58ae1dfa8 100644 --- a/agent_context/topics/sim-visualization/sim-visualization.md +++ b/agent_context/topics/sim-visualization/sim-visualization.md @@ -185,12 +185,15 @@ frames. Mesh geometry is identified by a SHA-256 hash of local vertices and faces. Static nodes sharing geometry are sent through one Viser batched-mesh handle. -Normal frames update only positions, `wxyz` quaternions, and visibility. +Normal frames update only positions, Viser-native `wxyz` quaternions, and visibility. Identifiers are URL-escaped before becoming Viser path components. -EmbodiChain pose vectors use `(x, y, z, qw, qx, qy, qz)`. The protocol uses -normalized `wxyz` quaternions. `pose_to_position_wxyz()` is the conversion -boundary and also accepts homogeneous `(..., 4, 4)` matrices. +EmbodiChain pose vectors use `(x, y, z, qx, qy, qz, qw)`. The visualization +protocol follows Viser and stores normalized `wxyz` quaternions. +`pose_to_position_wxyz()` converts EmbodiChain `xyz + xyzw` pose vectors at +that boundary and also accepts homogeneous `(..., 4, 4)` matrices. Protocol +dataclass fields already named `wxyz` remain protocol-native and must not be +interpreted as EmbodiChain pose vectors. Arena offsets are added to rigid, robot, articulation, and camera poses. Deformable vertices are stored relative to the corresponding arena node. @@ -203,8 +206,8 @@ Deformable vertices are stored relative to the corresponding arena node. | `RigidObjectGroup` | One node and pose per constituent object | | `Robot` | One mesh node per non-empty link | | `Articulation` | One mesh node per non-empty link | -| `SoftObject` | Live collision vertices with a cached convex-hull surface | -| `ClothObject` | Live physical vertices with render triangles mapped onto the welded physical vertex buffer | +| Volume `DeformableObject` (`SoftObject`) | Live collision vertices with a cached convex-hull surface | +| Surface `DeformableObject` (`ClothObject`) | Live physical vertices with render triangles mapped onto the welded physical vertex buffer | | `Camera` | Frustum plus optional low-frequency RGB preview | | Default ground | 1000 m × 1000 m XY grid, 1 m cells, 10 m sections | | `SceneOverlays` | Frames, targets, trajectories, and point clouds | @@ -232,11 +235,16 @@ slow rendering or clients cannot accumulate an image backlog. ## Deformables -Soft bodies and cloth require GPU physics. Their live vertices are sampled at -`soft_body_fps`, independently from `scene_fps`. +Volume and surface deformables currently require Default-backend GPU physics. +Their live vertices are sampled at `soft_body_fps`, independently from +`scene_fps`. `SceneExporter` enumerates the manager's single deformable +registry and reads both topologies through `get_surface_vertices()` and +`get_surface_triangles()`; it does not branch on legacy buffer APIs. The +`deformable_type` discriminator only selects the existing soft/cloth browser +node kind, path, and color. - DexSim does not expose soft-body collision triangle connectivity. - `SoftBodyData.collision_surface_triangles` therefore caches a SciPy + `VolumeDeformableData.collision_surface_triangles` therefore caches a SciPy `ConvexHull` over rest collision vertices. The preview follows deformation but cannot preserve concave render detail. - Cloth maps all render-mesh triangles onto DexSim's welded rest-vertex buffer diff --git a/agent_context/topics/simulation-system/simulation-system.md b/agent_context/topics/simulation-system/simulation-system.md index f76ded88c..dc4fd1023 100644 --- a/agent_context/topics/simulation-system/simulation-system.md +++ b/agent_context/topics/simulation-system/simulation-system.md @@ -7,7 +7,9 @@ | Public simulation package | `embodichain/lab/sim/__init__.py` | | World and scene owner | `embodichain/lab/sim/sim_manager.py` → `SimulationManager` | | Global simulation config | `embodichain/lab/sim/sim_manager.py` → `SimulationManagerCfg` | -| Object and physics configs | `embodichain/lab/sim/cfg.py` | +| Spawn lifecycle coordinator | `embodichain/lab/sim/spawn/scene.py` → `SpawnScene` | +| EmbodiChain-to-Spawn translation | `embodichain/lab/sim/spawn/descriptors.py` | +| Object and physics configs | `embodichain/lab/sim/cfg/` (public facade: `cfg/__init__.py`) | | Gym lifecycle integration | `embodichain/lab/gym/envs/base_env.py` | | Task scene construction | `embodichain/lab/gym/envs/embodied_env.py` | @@ -17,11 +19,18 @@ object, sensor, solver, planner, or atomic-action API from its own subpackage. ## Ownership -`SimulationManager` owns one DexSim `World`, its global environment, -parallel arenas, and the Python registries for scene resources: +`SimulationManager` owns one DexSim `World`, a `SpawnScene`, and the Python +registries for scene resources. DexSim's `SceneBuilder` and `SpawnResult` own +descriptor revisions, native materialization, replicated arenas, and backend +handles. `SimulationManager` owns the readiness boundary for each committed +Spawn topology revision. EmbodiChain registry objects are stable facades: +`add_*()` returns a declared facade and `prepare()` binds that same object in +place. + +The registries cover: - rigid objects and rigid-object groups; -- soft and cloth objects; +- volume and surface deformables in one deformable-object registry; - articulations and robots; - rigid constraints, sensors, lights, gizmos, and markers; - visual materials and texture caches; @@ -40,9 +49,20 @@ The environment-owned lifecycle is: EnvCfg.sim_cfg → BaseEnv._setup_scene() → SimulationManager(SimulationManagerCfg) - → create World, global environment, defaults, and N arenas - → EmbodiedEnv adds robot, objects, lights, and sensors - → initialize GPU physics after scene construction when using CUDA + → create World and a replicated Spawn scene declaration + → EmbodiedEnv declares robot, objects, lights, and physical sensors + → Default may materialize native handles eagerly + → Newton keeps physical descriptors deferred + → SimulationManager.prepare() + → for Newton, resolve source metadata and configure exact-name overlays + → finalize/rebuild pending Spawn descriptors once + → for Default, apply pending source overlays to materialized handles + → apply Default articulation-root runtime properties to native handles + → prepare manager-owned runtime buffers for the committed revision + → bind declared EmbodiChain facades in place + → publish bound state through the backend render-sync hook + → attach sensors whose parents are now materialized + → initialize metadata-dependent robot, action, and render-only resources → BaseEnv.step() → preprocess/apply action → SimulationManager.update(physics_dt, sim_steps_per_control) @@ -55,28 +75,126 @@ EnvCfg.sim_cfg → SimulationManager.destroy() ``` +`destroy(exit_process=False)` queues native cleanup; callers flush that queue +only after their scene/object locals have unwound. During +`SimulationManager._deferred_destroy()`, the manager stops recording and the +native window, invokes `PhysicsBackend.prepare_for_teardown()`, then runs GC +before closing the Spawn result, environment, and World. Default backends use +the no-op hook. Newton synchronizes its resolved Warp CUDA device and clears +its render bridge while Spawn still owns the parent skeletons, so cached link +views cannot be destructed after their native parents. + +After backend materialization, dynamic `RigidObject`, `Articulation`, and +`RigidObjectGroup` facades capture their resolved mass, inertia diagonal, and +local center-of-mass pose in their data objects. The layouts are `[env]` in +`RigidBodyData`, `[env, link]` in `ArticulationData`, and `[env, object]` in +`RigidBodyGroupData`. Each data object exposes current `mass`, `inertia`, and +`com_pose` values plus immutable `default_*` initialization snapshots. Runtime +property writes do not change these snapshots. During reset, only the selected +environment rows are restored before dynamics are cleared and the configured +pose is reapplied; reset-mode event functors then run from this clean physical +baseline in the episode-initialization hook. + +## Quaternion and pose convention + +All EmbodiChain-owned public and runtime quaternion tensors use +`(x, y, z, w)` (`xyzw`). A 7D pose or state therefore uses +`(px, py, pz, qx, qy, qz, qw)` (`xyz + xyzw`), and the identity quaternion is +`(0, 0, 0, 1)`. This includes object/root/link/COM state, robot FK and IK, +sensor offsets, manager observations/actions, semantic poses, and task +configuration. `embodichain.utils.math` follows the same convention. + +Backend and library adapters must preserve the external API's native order and +convert exactly once at that boundary. DexSim/Spawn rigid and articulation pose +buffers are native `xyzw + xyz`, so their adapters only permute pose layout. +DexSim mass-property and COM descriptors are native `wxyz`, so those adapters +use `convert_quat()` explicitly. Newton/Warp transforms expose position plus an +`xyzw` quaternion and therefore need no component-order conversion. Use a +non-symmetric rotation when testing an adapter; an identity or 180-degree +single-axis rotation can hide an incorrect order. + +Deformables use the same public hierarchy for both topologies: +`DeformableObjectCfg` is specialized by `VolumeDeformableObjectCfg` and +`SurfaceDeformableObjectCfg`; `SoftObjectCfg` and `ClothObjectCfg` remain +compatibility subclasses. `objects/deformable/` owns the common +`DeformableObject`/`DeformableObjectData` contract and the DexSim volume and +surface implementations. Consumers should use `data.nodal_pos_w`, +`data.nodal_vel_w`, `data.nodal_state_w`, `get_surface_vertices()`, and +`get_surface_triangles()`. Legacy soft/cloth methods delegate to that contract. +At the Spawn boundary, volume and surface configs translate to DexSim's typed +`SoftBodyDesc` and `ClothDesc` particle-set descriptors. Their Default-native +attributes are carried by `DexsimSoftBodyPhysicsDesc` and +`DexsimClothPhysicsDesc`; volume voxel settings use `SoftBodyMeshingDesc`. + +`SimulationManager` stores both topologies once in `_deformable_objects` and +exposes `add/get_deformable_object()` plus filtered legacy soft/cloth APIs. +Only the Default backend is registered today and still requires CUDA. +Backend capability flags and `_DEFORMABLE_BACKEND_IMPLEMENTATIONS` reserve the +Newton integration boundary; Newton volume/surface support must remain disabled +until native object and data adapters are implemented and validated. + `BaseEnv._setup_scene()` temporarily constructs the manager headlessly so the scene can be assembled before a native window is opened. It sets `SimulationManagerCfg.num_envs` from `EnvCfg.num_envs`. -`SimulationManager` enables physics, selects manual physics updates, creates -the configured arenas, installs default plane/background/lighting resources, -and starts configured visualization during initialization. A Viser backend -forces `headless=True`; Viser and the native DexSim window are mutually -exclusive. +`SimulationManager` enables physics, selects manual physics updates, prepares +the configured Arena layout, and owns a thin Spawn scene coordinator. With the +Default backend, preparing the Arena layout lets `add_*` materialize native +entities immediately, so articulation metadata and render nodes are available +before finalization. A source-backed articulation added to an eager Default +result is loaded first and then receives its exact-name typed properties on the +live native articulation. Newton defers physical materialization until +`prepare()`: EmbodiChain first reads exact URDF metadata through a disposable +render-only skeleton, applies the source-name overlays, and then builds the +immutable Newton model once. A Viser backend forces `headless=True`; Viser and +the native DexSim window are mutually exclusive. -`SimulationManager.update()` initializes GPU physics lazily if needed and -then advances the world for the requested number of physics steps. Each -environment control step normally calls it with -`sim_steps_per_control`. +`SpawnScene` always requests DexSim replication with +`collision_policy="isolated"`. Consequently, when `num_envs > 1`, all +per-environment dynamic, kinematic, and static rigid shapes and every +articulation link shape collide only with entities in the same Arena. Global +`per_env=False` physics resources still collide with every Arena. EmbodiChain +owns this policy choice; DexSim's `ReplicatePlan` and backend adapters own the +effective Default filter data and Newton collision groups. Do not duplicate +the backend-specific group calculation in object facades or task configs. + +The default ground plane authors its repeated texture coordinates in the Spawn +render descriptor before materialization, so native and offscreen render paths +receive identical UV data on their first GPU upload. + +`SimulationManager.prepare()` is the backend-neutral readiness boundary for +Default CPU, Direct GPU, and Newton. It is idempotent. Topology is committed +only when dirty. Newton source resolution and exact-name configuration precede +the first commit; Default source configuration follows native materialization. +A failed resolver or configurator remains pending and retryable. Runtime +preparation is recorded by committed topology revision: Default CUDA calls +`World.init_gpu_physics()`, while Default CPU and Newton need no additional +manager call after Spawn commit. After facade binding/reset, the active physics +backend publishes current state to render resources once per committed topology +revision. This is a no-op for Default and invokes Newton's render bridge without +advancing simulation time. Facade binding, render publication, and sensor +attachment remain retryable; already completed declarations are not +reconfigured or rebound. `init_gpu_physics()` and +`finalize_newton_physics()` remain compatibility aliases, but new code should +call `prepare()`. + +Standalone callers must call `prepare()` after their last `add_*()` and before +reading link/joint metadata, object state, or advancing physics. `BaseEnv` +provides this boundary automatically between `_setup_scene()` and +metadata-dependent setup. `SimulationManager.update()` still calls the +readiness path defensively before advancing the requested physics steps. ## Module Boundaries | Area | Owner | Routed topic | |------|-------|--------------| | World, arenas, asset registries, physics update, cleanup | `sim_manager.py` | `simulation-system` | -| Shared object, render, physics, drive, and URDF configs | `cfg.py` | `configclass-pattern` for config mechanics | -| Rigid, deformable, articulation, robot, light, constraint, gizmo | `objects/` | `robot-system` for robots | +| Backend activation and configured/resolved solver state | `physics/` | `simulation-system` | +| Spawn declaration, source resolution, commit/rebuild, and facade binding | `spawn/scene.py`, `spawn/source.py`, `spawn/descriptors.py` | `simulation-system` | +| Backend-neutral batched state/property access | `objects/backends/spawn.py` | `simulation-system` | +| Shared object, render, physics, drive, and URDF configs | `cfg/` domain modules; `cfg/__init__.py` preserves the public import surface | `configclass-pattern` for config mechanics | +| Rigid, articulation, robot, light, constraint, gizmo | `objects/` | `robot-system` for robots | +| Common deformable contract and DexSim volume/surface adapters | `objects/deformable/` | `sim-visualization` for export | | Camera, stereo camera, contact sensor | `sensors/` | `sensor-system` | | Robot-specific configuration | `robots/` | `robot-system` | | Inverse kinematics | `solvers/` | `ik-solvers` | @@ -98,25 +216,281 @@ origin, axis, and optional limits. Consumers must not reach into ## Configuration Flow -`SimulationManagerCfg` owns window size, headless mode, rendering, GPU/CPU -selection, arena count and spacing, physics timestep, physics and GPU-memory -settings, recording, profiling, and browser visualization. +`SimulationManagerCfg.physics_cfg` is the backend selector as well as the +backend config. `PhysicsBackendCfg` owns common timing, device, and gravity; +`DefaultPhysicsCfg` adds default-backend scene settings, while +`NewtonPhysicsCfg` adds the Newton solver, substeps, gradient/CUDA-graph +behavior, and a grouped `NewtonCollisionPipelineCfg`. +Do not add a second backend string that can disagree with the config type. +Leaving `NewtonPhysicsCfg.solver_cfg=None` preserves DexSim's +`AutoSolverCfg` default. A DexSim build exporting `AutoSolverCfg` is required; +EmbodiChain does not substitute a concrete solver. DexSim resolves that +placeholder from the complete Spawn scene during finalization: rigid-only +scenes select XPBD, scenes with an articulation select MuJoCo Warp, and +supported particle families select their matching particle/deformable solver. +A mapping with `solver_type: auto` or +`class_type: AutoSolverCfg` is the explicit equivalent. Gradient mode must +still select `semi_implicit` explicitly because AutoSolver does not choose a +differentiable solver. Before finalization, EmbodiChain treats `auto` as +unresolved; after finalization, `NewtonPhysicsBackend.solver_type` reads the +concrete type from DexSim's World-owned backend. +MuJoCo-Warp mappings may set `enable_multiccd: true`; EmbodiChain forwards it +to DexSim's `MJWarpSolverCfg`, which passes it to Newton `SolverMuJoCo`. +Enabling it changes contact generation (up to four contacts per geometry pair) +without changing the collision geometry authored by EmbodiChain. DexSim must +export an `MJWarpSolverCfg` version that declares the field. +The `open_drawer.py` tutorial combines this option with 20 Newton substeps per +10 ms control step, while keeping its authored robot gains, collision geometry, +pull trajectory, success criteria, and push trajectory identical to Default. +Atomic-action tutorials configure their shared Newton simulation in +`scripts/tutorials/atomic_action/tutorial_utils.py`: they retain 20 substeps +while following Newton's brick-stacking contact profile (`solver=newton`, +`integrator=implicitfast`, 15 solver iterations, 100 line-search iterations, +an elliptic friction cone, `impratio=50`, and the Newton collision pipeline +with contact reduction and an `nxn` broad phase). The shared factory leaves +the Default backend configuration unchanged. +The package dependency must identify the exact DexSim dev build containing +this API; a base `==0.4.3` requirement also accepts older local-version wheels +that do not export `AutoSolverCfg` and is therefore insufficient. +Newton's `suppress_warp_kernel_logs=True` suppresses Warp's one-time runtime +banner plus module compile/load chatter during manager startup, build, facade +initialization, and physics updates, then restores the process-wide setting. +It does not suppress DexSim native startup output or genuine Warp/Newton +warnings and errors. + +EmbodiChain-authored Newton collision shapes use a default margin and gap of +`0.001 m` each only when no portable or Newton-native envelope is authored. +`CollisionPropertiesCfg.contact_offset/rest_offset` are portable: Default uses +them directly, while the Spawn compiler maps `rest_offset → margin` and +`contact_offset - rest_offset → gap`. Both values must be present to derive a +Newton gap; an active Newton configuration rejects an ambiguous standalone +`contact_offset` unless a native margin or gap completes the intent. Explicit +`NewtonCollisionPropertiesCfg.margin/gap` values take precedence over this +translation. `EnvCfg` embeds `SimulationManagerCfg` and supplies the control-to-physics step ratio. CLI and task config loaders may override runtime fields before constructing the environment. Trace those overrides through the caller rather than changing a default in the manager blindly. -Object-specific configuration belongs in `lab/sim/cfg.py` or the -corresponding robot/sensor module. Scene composition belongs in +Object-specific configuration belongs in the matching `lab/sim/cfg/` domain +module or the corresponding robot/sensor module. Scene composition belongs in `EmbodiedEnv` or a task config, not in `SimulationManagerCfg`. +Deformable configs use an explicit `deformable_type: volume|surface` +discriminator. Common source mesh and pose fields stay on +`DeformableObjectCfg`; tetrahedral voxelization/soft-body attributes stay on +the volume subclass, and cloth attributes stay on the surface subclass. Do not +add backend conditionals to one monolithic deformable config. Add a backend +implementation at the manager dispatch boundary when its runtime exists. + +New rigid-body configs use `RigidBodyPhysicsCfg`. Portable intent is organized +by physical concept: + +- `mass_props`: `MassPropertiesCfg` (`mass`, `density`, inertia, COM, and the + source-inertia recomputation policy); +- `rigid_props`: `DefaultRigidBodyPropertiesCfg`; Newton currently exposes no + additional body-level property group beyond common mass properties; +- `collision_props`: common collision enablement and the portable + `contact_offset/rest_offset` envelope, optionally specialized by + `DefaultCollisionPropertiesCfg` or `NewtonCollisionPropertiesCfg`; +- `material_props`: common friction/restitution or a backend material subclass. + +Each concept has exactly one slot. Backend-native fields are represented by the +slot's concrete subclass or its local `backend: default|newton` discriminator; +`default_props` and `newton_props` were removed. Every grouped field defaults to +`None`, meaning “do not author this field”; source USD/URDF values and backend +defaults therefore survive partial overlays. Dynamic and kinematic mass +priority is explicit inertia with positive mass, then mass, then density; +static descriptors omit mass properties. + +Mesh collision construction is geometry-owned. `MeshCfg.collision` contains a +`MeshCollisionCfg` with an explicit `convex_hull`, `convex_decomposition`, +`triangle_mesh`, or `sdf` approximation. Strategy-specific fields are validated +when the config is constructed; numerical values never infer the strategy in +the canonical schema. Newton SDF and hydroelastic mesh settings share this +single owner. `RigidBodyPhysicsCfg` and articulation link overlays do not carry +mesh cooking. An imported articulation retains its source mesh approximation +until a named source-shape overlay API is introduced. + +`MassPropertiesCfg.recompute_inertia=True` discards source-authored inertia so +the backend derives it from collision geometry and the effective mass or +density. The default `None` inherits an outer per-body overlay and otherwise +preserves source inertia. Explicit inertia and recomputation are mutually +exclusive. The policy lives with mass properties so global articulation, +per-link articulation, and rigid USD overlays share the same behavior; +`LinkPhysicsOverrideCfg` only selects links and carries their partial `attrs`. + +Polymorphic collision and material slots use a local +`backend: common|default|newton` discriminator; a unique native field may infer +the subtype. `rigid_props` currently accepts only `backend: default`. +`MeshCfg.from_dict()` temporarily normalizes the deprecated flat +`max_convex_hull_num`, `acd_method`, and `sdf_resolution` inputs to +`MeshCfg.collision` with a deprecation warning. `RigidObjectCfg.from_dict()` +also migrates the former `attrs.mesh_collision_props` input when it has the +owning mesh shape. Serialization emits only the new nested geometry form. + +`RigidBodyPhysicsCfg` is the only user-facing rigid-body physics schema. +Flat `attrs` keys such as `mass`, `dynamic_friction`, and `enable_collision` +are rejected at the parsing boundary; place them in `mass_props`, +`material_props`, or `collision_props` instead. `LinkPhysicsOverrideCfg.attrs` +uses the same partial schema, so global and per-link overlays share one model. +COM quaternions in every EmbodiChain config and public runtime API are `xyzw`. +The Spawn/Default adapter alone converts them to DexSim's native `wxyz` order. + +Robot configs normally keep these portable values on one ordinary `RobotCfg`. +For a genuine backend-specific asset or actuator difference, subclass +`RobotPresetCfg` and declare complete `default`, `newton`, or +`newton_` alternatives. +`SimulationManager.add_robot()` derives the +selection from its existing `physics_cfg`, deep-copies the selected complete +robot config, and never merges alternatives. While AutoSolver is unresolved, +only the generic `newton` and `default` alternatives are eligible; do not guess +a solver-specific preset before DexSim has inspected the complete scene. This +is the only robot preset selection boundary; do not add a second backend +selector to robot configs. + +File-backed rigid objects and articulations share one source-independent +physics policy: `asset_physics_mode="preserve"` keeps properties resolved from +the asset, while `asset_physics_mode="overlay"` applies only non-`None` +EmbodiChain fields after DexSim has translated the real materialized source. +This policy applies equally to USD rigid objects and USD/URDF articulations. +Generic `RigidObjectCfg` and `ArticulationCfg` default to `preserve`; `RobotCfg` +defaults to `overlay` to retain its established configured-drive behavior. +If an articulation in preserve mode contains explicit `attrs`, `link_attrs`, +`joint_drive_props`, or `qpos_limits`, configuration emits a warning +naming the ignored overlay fields instead of silently discarding them. +Import concerns that the source format does not author, such as URDF root +fixation and body scale, remain controlled by their dedicated fields. An +articulation defaults to `root_props.fixed_base=True` and +`root_props.self_collision_enabled=False`, so both URDF and USD assets are +fixed to the world with self-collision disabled unless configured otherwise. +Setting either field explicitly to `None` preserves the corresponding USD +property and selects the established URDF import default. + +`ArticulationRootPropertiesCfg` is the single root-property definition. Spawn +consumes its portable fixed-base and self-collision intent through common +articulation descriptor fields. Its `sleep_threshold`, `min_position_iters`, +and `min_velocity_iters` fields are Default-only: EmbodiChain applies them to +the materialized native articulation before Direct GPU initialization and the +first reset, while Newton ignores them. PhysX Direct GPU runtime setup captures +the articulation solver iteration counts; applying them only during facade +binding leaves the active GPU solver at its source/default values and can make +mimic constraints much softer than CPU. The preparation is idempotent per +Spawn topology revision. The two iteration counts must be configured together +because the Default native API exposes one atomic setter. This remains distinct +from `DefaultRigidBodyPropertiesCfg`, whose same-named values configure +individual rigid bodies or articulation links. `root_props` is the only +root-property interface; `fix_base`, `disable_self_collision`, and the former +flat root solver fields are removed. `JointDrivePropertiesCfg` keeps the +original `drive_type` (`force`, Default-only `acceleration`, or `none`) and adds +the portable actuator `target_mode` (`none`, `position`, `velocity`, +`position_velocity`, or `effort`), stiffness/damping gains, effort/velocity +limits, passive friction, and armature. `ArticulationCfg.joint_drive_props` is +the single joint-property entry point. Every field is optional; `None` means source-owned, +which permits sparse overlays without resetting unrelated source values. If +`target_mode` is unset, +`drive_type="force"` or `"acceleration"` defaults it to `position_velocity`, +while `drive_type="none"` defaults it to `none`. +`NewtonJointDrivePropertiesCfg` is only a serialized configuration +compatibility subtype; new robot definitions use the common class. Common +effort/velocity/armature values stay on `JointDesc` instead of being duplicated +in both backend blocks. `link_attrs` accepts the same grouped rigid-body schema +for partial per-link overrides. + +Spawn resolves drive intent per source-resolved joint before lowering it. +Default selects its force/acceleration enum and masks inactive gains; Newton +authors `JointTargetMode` values 0 through 4. `none` and `effort` always clear +both target gains, and `velocity` clears the position gain, so Newton solvers +that ignore `joint_target_mode` still receive deterministic passive, +effort-only, and velocity-only behavior. Non-MuJoCo Newton position mode is an +explicit gain-based emulation that assumes a zero velocity target. An active +`drive_type="acceleration"` is rejected for Newton because it has no exact +equivalent. + +For articulations, `SimulationManager._declare_spawn_articulation()` supplies +`configure_articulation_desc()` as the source-configuration callback. Preserve +mode leaves source descriptors untouched, while overlay mode applies +exact-name link/joint fields. Both regex dictionaries and flattened +`(num_dofs, 2)` arrays in `qpos_limits` are compiled into the resolved +joint descriptors before either backend builds. Default obtains those names +from its loaded native articulation and applies the typed properties live. +Newton resolves the same metadata first and consumes the configured descriptor +during its initial immutable-model build, so initial source configuration must +not be implemented as finalize-then-rebuild. Do not duplicate these link/joint +descriptor writes in `Articulation._apply_spawn_config()`; that hook is +reserved for Default-native root setters, finalized Newton runtime adaptation, +and render work requiring finalized resources. + +Spawn articulation state IDs follow the final batch `qpos`/`qvel` layout, which +can differ from Newton's source-articulation traversal order. Initial `qpos`, +mimic child/parent metadata, control groups, and every batch mutation must be +mapped by joint name into that state layout. Public `joint_names` uses this +same state-buffer order; use the Spawn handle's source-name query only when +resolving source topology. Newton solvers without configured mimic compliance +project reset positions onto the authored relation before the first step. The +MuJoCo-Warp compliance path preserves the authored current position, matching +Default's initial hand state. +`SpawnArticulationView` filters Newton root-pose rows that already match the +requested translation and rotation before calling the Spawn batch write. This +keeps ordinary fixed-root resets from invalidating a captured CUDA graph while +still forwarding genuine root-pose changes, which refresh Newton solver +constants and recapture the graph as required. +Initialization code that intentionally changes fixed-root poses should do so +after `prepare()` but before the first `update()`, allowing the first Newton +CUDA graph to capture the final anchors instead of immediately invalidating a +graph captured from transient poses. + +MuJoCo-Warp lowers URDF mimic joints to native joint equality constraints, but +its default equality solver reference is underdamped compared with Default's +PhysX mimic. During +`Articulation._apply_spawn_config()`, +`_configure_newton_mimic_compliance()` in `objects/backends/newton.py` resolves +only that articulation's constraint rows and approximates Default's natural +frequency/damping ratio with MuJoCo's positive, effective-mass-scaled +`(timeconst, dampratio)` `solref`; the time constant observes MuJoCo's +two-solver-timestep safety floor. The native rows remain enabled, preserving +contact force coupling between follower and leader joints. A very weak +follower drive (one percent of its leader's target gains; `ke=1`, `kd=0.1` +for the W1 hand) stabilizes the equality between solver updates; target +`set_qpos()` and `set_qvel()` writes propagate the authored leader relation to +that drive. Never copy measured follower state or disable +the native equality: doing either turns mimic into an independent servo and +loses the Default backend's mechanical coupling. Other Newton solvers, +gradient mode, and Default retain native behavior. Keep private Newton +runtime/solver access inside this backend helper; the generic `Articulation` +owns state-order metadata, reset behavior, and target propagation only. + +DexSim 0.4.3's Newton `RigidBodyBatch.apply_pose()` writes maximal `body_q` +state but does not update the standalone body's reduced FREE-joint state read +by MuJoCo-Warp on the next step. `SpawnRigidBodyView` therefore caches a +`StandaloneRigidStateSync` for its stable batch and projects both Newton state +buffers after pose writes. Invalidate that cache on a Spawn topology revision; +remove the compatibility path once DexSim's public batch operation guarantees +the same synchronization. + +The `grasp_cup_to_caffe.py` comparison demo seeds its XY perturbations after +`prepare()` (default seed `0`). This placement makes the scene independent of +random numbers consumed by backend initialization. Pass a negative `--seed` +to restore non-deterministic perturbations. + +Rigid USD objects follow the same overlay rule: parsed source descriptors are +updated field-by-field, never replaced wholesale by a partial config. The +former flat `RigidBodyAttributesCfg` and `RigidBodyAttributesOverrideCfg` +types have been removed. New and migrated definitions use the grouped schema, +where `None` means “leave the source/backend value unchanged.” + ## Where to Make Changes | Change | Primary location | |--------|------------------| | Global world, renderer, device, arena, or physics lifecycle | `sim_manager.py` | -| Shared object or physics config type | `cfg.py` | +| Spawn source translation or typed link/joint overrides | `spawn/descriptors.py` plus the DexSim Spawn descriptor/adapter boundary | +| Declaration-to-result binding or retry behavior | `spawn/scene.py` and the object's `bind_spawn()` | +| Batched row/DOF selection or backend property parity | `objects/backends/spawn.py` and the DexSim Spawn batch facade | +| Newton object/runtime adaptation | `objects/backends/newton.py` | +| Shared object or physics config type | Matching domain module under `cfg/`, then re-export from `cfg/__init__.py` | +| Deformable nodal/surface contract or topology-specific buffers | `objects/deformable/` | | Add/get/remove behavior for a scene entity | `sim_manager.py` plus its `objects/` implementation | | Task scene composition | `embodied_env.py` or the task config | | Environment timing, reset, or control-step behavior | `base_env.py` and `env-framework` | @@ -127,15 +501,33 @@ corresponding robot/sensor module. Scene composition belongs in - Configure `num_envs`, device, renderer, and physics settings before constructing `SimulationManager`. +- Treat `add_*()` as declaration. Call `prepare()` before consuming native + handles, link/joint metadata, batched state, or physics results. +- Keep `prepare()` convergent and retryable: do not mark a declaration bound + until its full facade construction succeeds. +- Keep backend render-state publication free of physics steps. Newton's initial + state sync must not advance its simulation step or time. - Treat resource UIDs as registry identities; retrieve and mutate resources through the manager instead of maintaining a parallel scene registry. - Keep batched object and sensor state aligned with the manager's arena count. -- Build scene assets before explicitly initializing GPU physics. The manager - will warn and initialize lazily on the first update if this was missed. +- Add the initial physical scene before `prepare()`. Calls to the legacy + `init_gpu_physics()` and `finalize_newton_physics()` aliases are equivalent to + `prepare()` and do not cause a second build. +- Delegate environment and DOF selections to DexSim Spawn batches instead of + full-batch read/modify/write loops in object facades. +- Newton descriptor or topology mutations that cannot update the immutable + runtime model live remain pending until the next `prepare()` rebuild. +- Apply Newton collision and articulation-joint configuration to the + source-translated Spawn descriptors before the first model build; post-bind + object initialization is only for state and supported live batch properties. - Manual update is the default; normal environment stepping must advance physics through `SimulationManager.update()`. - Reset only the requested environment rows and honor `excluded_uids` for resources detached from automatic reset. +- Keep the `default_mass`, `default_inertia`, and `default_com_pose` values in + `RigidBodyData`, `ArticulationData`, and `RigidBodyGroupData` as immutable + initialization snapshots; runtime setters and randomizers must not mutate + them. - `destroy()` queues deferred cleanup. Tests and non-exiting standalone callers that use `exit_process=False` must call `SimulationManager.flush_cleanup_queue()`. @@ -147,7 +539,9 @@ corresponding robot/sensor module. Scene composition belongs in | Symptom | Likely cause | |---------|--------------| | Scene resource cannot be found or the wrong object is returned | UID mismatch or code bypassed the manager registry | -| CUDA physics data is stale on the first step | GPU physics was initialized before all assets were added, or not initialized explicitly | +| Link/joint metadata is empty or state access fails after `add_*()` | The declared facade has not crossed `SimulationManager.prepare()` yet | +| CUDA/Newton physics data is stale after a topology or descriptor mutation | Call `prepare()` so the dirty Spawn result can rebuild and rebind runtime views | +| Warp module compile/load lines appear during Newton initialization | `NewtonPhysicsCfg.suppress_warp_kernel_logs` was explicitly disabled, or compilation happened outside the managed preparation scope | | Native window does not open | `headless=True`, often forced by the Viser backend | | Device and renderer use the wrong GPU | `sim_device` and `gpu_id` disagree; the device index takes precedence for CUDA simulation | | Simulation advances at the wrong control rate | `physics_dt` and `sim_steps_per_control` were configured inconsistently; see `env-framework` | diff --git a/design/newton-backend-design.md b/design/newton-backend-design.md new file mode 100644 index 000000000..4cdb2e9b5 --- /dev/null +++ b/design/newton-backend-design.md @@ -0,0 +1,441 @@ +# 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 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 (Default/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. + +### Grouped Newton and Default physics attributes + +`RigidBodyPhysicsCfg` is the single public schema for rigid-object and +articulation link physics. It separates portable values into `mass_props`, +`rigid_props`, `collision_props`, and `material_props`. Each concept has one +slot, and backend-native values use that slot's concrete subtype; parallel +`default_props`/`newton_props` blocks are not supported. Every field is +optional, so source-authored values survive sparse USD/URDF overlays. The same +partial schema is used by `LinkPhysicsOverrideCfg`, eliminating the former flat +compatibility/override type pair. + +Mesh collision construction is owned by `MeshCfg.collision`, whose explicit +approximation selects convex hull, convex decomposition, triangle mesh, or SDF. +Newton SDF/hydroelastic cooking fields live there rather than in rigid-body +physics. Imported articulation links keep their source mesh approximation until +a named source-shape overlay is available. + +Spawn compiles these groups into its backend-neutral rigid-body and shape +descriptors, then projects Default- or Newton-specific values at the selected +backend boundary. The remaining raw Default path uses a private +`PhysicalAttr` adapter only at that boundary. User-facing COM quaternions stay +in `xyzw` order; adapters convert to DexSim's `wxyz` order when writing native +attributes and convert back on reads. + +### 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 default-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. Procedural SDF and CoACD geometry is compiled +from `MeshCfg.collision` through the Spawn descriptor path. + +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. +- 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 local property-slot discriminators; the Spawn compiler + projects common and Newton-native fields; per-solver warnings (`xpbd` ignores + `ke`/`kd`; `mujoco_warp` ignores `restitution`) 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. +- Single-slot physics properties and `MeshCfg.collision` spawn through the + descriptor path; the 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 Spawn compiler authors a positive configured density on the + rigid-body descriptor to avoid the mass gap for dynamic bodies without an + explicit mass and inertia. 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/scripts/check_api_docs.py b/docs/scripts/check_api_docs.py index 24a0735ee..56a271f69 100644 --- a/docs/scripts/check_api_docs.py +++ b/docs/scripts/check_api_docs.py @@ -221,7 +221,7 @@ def discover_public_modules( ) if relative_parts[-1] == "__init__": relative_parts.pop() - if any(part.startswith("_") for part in relative_parts): + if any(part.startswith(("_", ".")) for part in relative_parts): continue tree = ast.parse( diff --git a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst index b80a2b364..107a90955 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.gym.envs.rst @@ -21,6 +21,7 @@ through :func:`~embodichain.lab.gym.utils.registration.make`. .. autosummary:: demo + differentiable_env expert_program managers types @@ -56,6 +57,21 @@ Environment Classes :members: :exclude-members: __init__, class_type +Differentiable Environment +-------------------------- + +``DifferentiableEmbodiedEnv`` keeps the standard environment lifecycle while +bridging Newton trajectories into PyTorch autograd for analytic policy-gradient +tasks. Dynamics and explicit kinematics subclasses provide the action and +output kernels; the base class owns tape-aware stepping and deferred resets. + +.. currentmodule:: embodichain.lab.gym.envs.differentiable_env + +.. autoclass:: DifferentiableEmbodiedEnv + :members: + :inherited-members: + :show-inheritance: + Controller-ready Actions ------------------------ 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 14b832079..ed19ce2b1 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.cfg.rst @@ -18,24 +18,48 @@ Entity configs form a small inheritance hierarchy rooted at ``ObjectBaseCfg`` (``LightCfg``, ``RigidObjectCfg``, ``SoftObjectCfg``, ``ClothObjectCfg``, ``ArticulationCfg`` and its ``RobotCfg`` subclass), while ``URDFCfg`` and ``RigidConstraintCfg`` describe multi-component assembly and constraints. +``RobotPresetCfg`` provides replace-only complete robot alternatives when a +backend-specific asset or actuator definition is unavoidable. +Public backend selectors use only ``default`` and ``newton``. Nested physical +property groups may additionally use ``common`` for backend-neutral intent; +DexSim names belong to the runtime and Spawn SDK adapter boundary. + +.. rubric:: Type aliases + +.. autosummary:: + + AssetPhysicsMode + MeshCollisionApproximation .. rubric:: Classes .. autosummary:: RenderCfg - PhysicsCfg + PhysicsBackendCfg + DefaultPhysicsCfg + NewtonPhysicsCfg + NewtonCollisionPipelineCfg MarkerCfg WindowRecordCfg WindowCameraPoseCfg GPUMemoryCfg - RigidBodyAttributesCfg - RigidBodyAttributesOverrideCfg + MassPropertiesCfg + DefaultRigidBodyPropertiesCfg + CollisionPropertiesCfg + DefaultCollisionPropertiesCfg + NewtonCollisionPropertiesCfg + RigidBodyMaterialCfg + NewtonRigidBodyMaterialCfg + MeshCollisionCfg + RigidBodyPhysicsCfg + ArticulationRootPropertiesCfg LinkPhysicsOverrideCfg SoftbodyVoxelAttributesCfg SoftbodyPhysicalAttributesCfg ClothPhysicalAttributesCfg JointDrivePropertiesCfg + NewtonJointDrivePropertiesCfg ObjectBaseCfg LightCfg RigidObjectCfg @@ -46,3 +70,4 @@ Entity configs form a small inheritance hierarchy rooted at ``ObjectBaseCfg`` URDFCfg ArticulationCfg RobotCfg + RobotPresetCfg diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst index 7f9577485..125b878fd 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.objects.rst @@ -198,3 +198,174 @@ Rigid Constraint :members: :inherited-members: :show-inheritance: + +Backend Views +------------- + +Backend views normalize tensor layouts and row selection across the default, +Newton, and DexSim Spawn runtimes. The package import path exposes the common +contracts, concrete adapters, and Newton collision-filter helpers. + +.. currentmodule:: embodichain.lab.sim.objects.backends + +.. autosummary:: + + ArticulationViewBase + RigidBodyViewBase + DefaultArticulationView + DefaultRigidBodyView + NewtonArticulationView + NewtonRigidBodyView + apply_collision_filter_for_entities + apply_collision_filter_for_envs + is_newton_scene + SpawnArticulationView + SpawnRigidBodyView + +.. autoclass:: ArticulationViewBase + :members: + +.. autoclass:: RigidBodyViewBase + :members: + +.. autoclass:: DefaultArticulationView + :members: + :show-inheritance: + +.. autoclass:: DefaultRigidBodyView + :members: + :show-inheritance: + +.. autoclass:: NewtonArticulationView + :members: + :show-inheritance: + +.. autoclass:: NewtonRigidBodyView + :members: + :show-inheritance: + +.. autoclass:: SpawnArticulationView + :members: + :show-inheritance: + +.. autoclass:: SpawnRigidBodyView + :members: + :show-inheritance: + +.. autofunction:: apply_collision_filter_for_entities + +.. autofunction:: apply_collision_filter_for_envs + +.. autofunction:: is_newton_scene + +Backend implementation import paths +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: embodichain.lab.sim.objects.backends.base + +.. autosummary:: + + RigidBodyViewBase + ArticulationViewBase + +.. currentmodule:: embodichain.lab.sim.objects.backends.default + +.. autosummary:: + + DefaultRigidBodyView + DefaultArticulationView + +.. currentmodule:: embodichain.lab.sim.objects.backends.newton + +.. autosummary:: + + NewtonRigidBodyView + NewtonArticulationView + apply_collision_filter_for_entities + apply_collision_filter_for_envs + is_newton_scene + +.. currentmodule:: embodichain.lab.sim.objects.backends.spawn + +.. autosummary:: + + SpawnArticulationView + SpawnRigidBodyView + +Unified Deformable Objects +-------------------------- + +The deformable package provides a backend-neutral nodal-state contract and +canonical surface/volume names. ``Cloth*`` and ``Soft*`` remain compatibility +aliases for existing environments and tutorials. + +.. currentmodule:: embodichain.lab.sim.objects.deformable + +.. autosummary:: + + ClothBodyData + ClothObject + DeformableObject + DeformableObjectData + SoftBodyData + SoftObject + SurfaceDeformableData + SurfaceDeformableObject + VolumeDeformableData + VolumeDeformableObject + +.. autoclass:: DeformableObject + :members: + :show-inheritance: + +.. autoclass:: DeformableObjectData + :members: + +.. autoclass:: SurfaceDeformableData + :members: + :show-inheritance: + +.. autoclass:: SurfaceDeformableObject + :members: + :show-inheritance: + +.. autoclass:: VolumeDeformableData + :members: + :show-inheritance: + +.. autoclass:: VolumeDeformableObject + :members: + :show-inheritance: + +Deformable implementation import paths +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. currentmodule:: embodichain.lab.sim.objects.deformable.base + +.. autosummary:: + + DeformableObject + +.. currentmodule:: embodichain.lab.sim.objects.deformable.data + +.. autosummary:: + + DeformableObjectData + +.. currentmodule:: embodichain.lab.sim.objects.deformable.surface + +.. autosummary:: + + ClothBodyData + ClothObject + SurfaceDeformableData + SurfaceDeformableObject + +.. currentmodule:: embodichain.lab.sim.objects.deformable.volume + +.. autosummary:: + + SoftBodyData + SoftObject + VolumeDeformableData + VolumeDeformableObject diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.shapes.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.shapes.rst index ebe5e3170..469524925 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.shapes.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.shapes.rst @@ -9,8 +9,15 @@ Overview Geometry configuration objects used to build the collision and visual shapes of rigid bodies. :class:`ShapeCfg` is the common base; :class:`MeshCfg`, :class:`CubeCfg`, and :class:`SphereCfg` describe triangle-mesh, box, and -sphere primitives respectively, and :class:`LoadOption` controls how mesh -assets are loaded and decomposed. +sphere primitives respectively. :class:`MeshCollisionCfg` explicitly selects +the collision representation and its cooking settings, while +:class:`LoadOption` controls mesh loading. + +.. rubric:: Type aliases + +.. autosummary:: + + MeshCollisionApproximation .. rubric:: Classes @@ -18,6 +25,7 @@ assets are loaded and decomposed. CubeCfg LoadOption + MeshCollisionCfg MeshCfg ShapeCfg SphereCfg @@ -36,6 +44,12 @@ assets are loaded and decomposed. :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate +.. autoclass:: MeshCollisionCfg + :members: + :undoc-members: + :show-inheritance: + :exclude-members: __init__, copy, replace, to_dict, validate + .. autoclass:: CubeCfg :members: :undoc-members: diff --git a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst index 088dcb39e..611539f88 100644 --- a/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst +++ b/docs/source/api_reference/embodichain/embodichain.lab.sim.sim_manager.rst @@ -19,6 +19,7 @@ instance registry instead of passing it around explicitly. SimulationManager SimulationManagerCfg + get_physics_scene .. currentmodule:: embodichain.lab.sim.sim_manager @@ -37,3 +38,8 @@ instance registry instead of passing it around explicitly. :undoc-members: :show-inheritance: :exclude-members: __init__, copy, replace, to_dict, validate + +Active Physics Scene +-------------------- + +.. autofunction:: get_physics_scene diff --git a/docs/source/api_reference/public_api.rst b/docs/source/api_reference/public_api.rst index 2fb836b97..840c9c24c 100644 --- a/docs/source/api_reference/public_api.rst +++ b/docs/source/api_reference/public_api.rst @@ -818,6 +818,40 @@ embodichain.lab.sim.atomic_actions.transports EndpointCommandRouter EndpointCommandTransport +embodichain.lab.sim.diff +------------------------ + +Public differentiable-stepping bridge from manager-owned Newton trajectories +and Warp tapes into PyTorch autograd. + +.. currentmodule:: embodichain.lab.sim.diff + +.. autosummary:: + + NewtonStepFunc + differentiable_step + tape_context + +embodichain.lab.sim.diff.bridge +------------------------------- + +.. currentmodule:: embodichain.lab.sim.diff.bridge + +.. autosummary:: + + NewtonStepFunc + differentiable_step + tape_context + +embodichain.lab.sim.diff.runtime +-------------------------------- + +.. currentmodule:: embodichain.lab.sim.diff.runtime + +.. autosummary:: + + NewtonDifferentiableRuntime + embodichain.lab.sim.objects.articulation ---------------------------------------- @@ -839,6 +873,9 @@ embodichain.lab.sim.objects.cloth_object ClothBodyData ClothObject ClothObjectCfg + SurfaceDeformableData + SurfaceDeformableObject + SurfaceDeformableObjectCfg embodichain.lab.sim.objects.constraint -------------------------------------- @@ -901,6 +938,51 @@ embodichain.lab.sim.objects.soft_object SoftBodyData SoftObject SoftObjectCfg + VolumeDeformableData + VolumeDeformableObject + VolumeDeformableObjectCfg + +embodichain.lab.sim.physics +--------------------------- + +Manager-level physics backend selection and lifecycle contracts for the +Default and Newton implementations integrated through DexSim. + +.. currentmodule:: embodichain.lab.sim.physics + +.. autosummary:: + + PhysicsBackend + DefaultPhysicsBackend + NewtonPhysicsBackend + make_physics_backend + +embodichain.lab.sim.physics.base +-------------------------------- + +.. currentmodule:: embodichain.lab.sim.physics.base + +.. autosummary:: + + PhysicsBackend + +embodichain.lab.sim.physics.default +----------------------------------- + +.. currentmodule:: embodichain.lab.sim.physics.default + +.. autosummary:: + + DefaultPhysicsBackend + +embodichain.lab.sim.physics.newton +---------------------------------- + +.. currentmodule:: embodichain.lab.sim.physics.newton + +.. autosummary:: + + NewtonPhysicsBackend embodichain.lab.sim.planners.base_planner ----------------------------------------- @@ -1300,6 +1382,68 @@ embodichain.lab.sim.solvers.srs_solver SRSSolver SRSSolverCfg +embodichain.lab.sim.spawn +------------------------- + +Translation boundary from EmbodiChain object configs and singleton USD assets +into DexSim Spawn descriptors. + +.. currentmodule:: embodichain.lab.sim.spawn + +.. autosummary:: + + articulation_desc_from_cfg + articulation_desc_from_usd + cloth_desc_from_cfg + rigid_desc_from_cfg + rigid_desc_from_usd + soft_desc_from_cfg + surface_deformable_desc_from_cfg + volume_deformable_desc_from_cfg + +embodichain.lab.sim.spawn.descriptors +------------------------------------- + +.. currentmodule:: embodichain.lab.sim.spawn.descriptors + +.. autosummary:: + + articulation_desc_from_cfg + cloth_desc_from_cfg + configure_articulation_desc + rigid_desc_from_cfg + soft_desc_from_cfg + surface_deformable_desc_from_cfg + volume_deformable_desc_from_cfg + +embodichain.lab.sim.spawn.scene +------------------------------- + +.. currentmodule:: embodichain.lab.sim.spawn.scene + +.. autosummary:: + + SpawnScene + +embodichain.lab.sim.spawn.source +-------------------------------- + +.. currentmodule:: embodichain.lab.sim.spawn.source + +.. autosummary:: + + resolve_articulation_source + +embodichain.lab.sim.spawn.usd +----------------------------- + +.. currentmodule:: embodichain.lab.sim.spawn.usd + +.. autosummary:: + + articulation_desc_from_usd + rigid_desc_from_usd + embodichain.lab.sim.utility.render_utils ---------------------------------------- @@ -2002,6 +2146,18 @@ embodichain_tasks.manipulation.tableware.stack_cups StackCupsEnv +embodichain_tasks.special.franka_reach_apg +------------------------------------------- + +Differentiable Franka FR3 reach environment that demonstrates the explicit +kinematics route used by analytic policy-gradient experiments. + +.. currentmodule:: embodichain_tasks.special.franka_reach_apg + +.. autosummary:: + + FrankaReachApgEnv + embodichain_tasks.special.simple_task ------------------------------------- diff --git a/docs/source/features/workspace_analyzer/workspace_analyzer.md b/docs/source/features/workspace_analyzer/workspace_analyzer.md index 2e49ff0e0..e5416350e 100644 --- a/docs/source/features/workspace_analyzer/workspace_analyzer.md +++ b/docs/source/features/workspace_analyzer/workspace_analyzer.md @@ -28,7 +28,7 @@ from embodichain.lab.sim.workspace 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({ @@ -170,7 +170,7 @@ from embodichain.lab.sim.workspace import ( from embodichain.lab.sim.workspace.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/guides/add_robot.rst b/docs/source/guides/add_robot.rst index 3601c37d4..c075e816f 100644 --- a/docs/source/guides/add_robot.rst +++ b/docs/source/guides/add_robot.rst @@ -13,7 +13,7 @@ Every robot config subclasses :class:`~embodichain.lab.sim.cfg.RobotCfg` and overrides two hooks: - ``_build_defaults(self, init_dict=None)`` — populate ``urdf_cfg``, - ``control_parts``, ``solver_cfg``, ``drive_pros`` and ``attrs`` from variant + ``control_parts``, ``solver_cfg``, ``joint_drive_props`` and ``attrs`` from variant fields read out of ``init_dict``. - ``build_pk_serial_chain(self, device=...)`` — return a ``{control_part: pk.SerialChain}`` mapping, reading the PK URDF from a single @@ -34,7 +34,7 @@ Checklist 1. **Prepare the URDF** — place the URDF (+ meshes) in the assets directory. 2. **Override** ``_build_defaults(self, init_dict=None)`` — set variant fields from ``init_dict``, then populate ``urdf_cfg`` / ``control_parts`` / ``solver_cfg`` / - ``drive_pros`` / ``attrs``. + ``joint_drive_props`` / ``attrs``. 3. **Define control parts** — group joints into logical sets (e.g. ``arm``, ``gripper``). 4. **Configure the IK solver** — ``OPWSolverCfg`` (6-DOF), ``SRSSolverCfg`` (7-DOF), or a generic ``SolverCfg``. @@ -68,9 +68,9 @@ Key parameters +---------------------+----------------------------------+----------------------------------+ | ``solver_cfg`` | Dict[str, SolverCfg] | IK solver configurations | +---------------------+----------------------------------+----------------------------------+ -| ``drive_pros`` | JointDrivePropertiesCfg | Joint stiffness, damping, force | +| ``joint_drive_props`` | JointDrivePropertiesCfg | Joint drive, limits, friction | +---------------------+----------------------------------+----------------------------------+ -| ``attrs`` | RigidBodyAttributesCfg | Rigid-body physics attributes | +| ``attrs`` | RigidBodyPhysicsCfg | Grouped rigid-body physics | +---------------------+----------------------------------+----------------------------------+ | variant fields | enum / str / bool | Optional subclass fields | | | | (e.g. ``version``) | diff --git a/docs/source/guides/configuration.md b/docs/source/guides/configuration.md index d15cd6ae4..49f8e29ee 100644 --- a/docs/source/guides/configuration.md +++ b/docs/source/guides/configuration.md @@ -34,12 +34,12 @@ EmbodiChain configs form a nested hierarchy: EmbodiedEnvCfg ├── sim_cfg: SimulationManagerCfg │ ├── render_cfg: RenderCfg -│ ├── physics_config: PhysicsCfg -│ ├── gpu_memory_config: GPUMemoryCfg +│ ├── physics_cfg: DefaultPhysicsCfg | NewtonPhysicsCfg +│ │ └── gpu_memory: GPUMemoryCfg # Default backend only │ └── visualization: VisualizationCfg ├── robot: RobotCfg │ ├── urdf_cfg: URDFCfg -│ ├── drive_pros: JointDrivePropertiesCfg +│ ├── joint_drive_props: JointDrivePropertiesCfg │ └── solver_cfg: Dict[str, SolverCfg] ├── sensor: List[SensorCfg] ├── events: EventCfg @@ -194,8 +194,8 @@ When a training config references a gym config (via `trainer.gym_config`), the n "height": 540, "width": 960 } - ], - "env": { + ], + "env": { "control_parts": ["arm"], "actions": { "delta_qpos": { diff --git a/docs/source/guides/preview_asset.md b/docs/source/guides/preview_asset.md index a83fa49bc..a0fd2ef6f 100644 --- a/docs/source/guides/preview_asset.md +++ b/docs/source/guides/preview_asset.md @@ -146,7 +146,7 @@ asset.set_local_pose(pose) | `--init_pos X Y Z` | `0 0 0.5` | Initial position of the first asset. | | `--init_rot RX RY RZ` | `0 0 0` | Initial rotation in degrees. | | `--body_type` | `kinematic` | Rigid body type: `dynamic`, `kinematic`, or `static`. | -| `--use_usd_properties` | disabled | Use physical properties stored in the USD file. | +| `--asset-physics-mode {preserve,overlay}` | `overlay` | Preserve source-authored physics or overlay explicitly configured values. | | `--fix_base` / `--no-fix_base` | fixed | Fix or unfix articulation bases. | | `--sim_device` | `cpu` | Simulation device. | | `--renderer` | `hybrid` | Renderer: `hybrid`, `fast-rt`, or `rt`. | diff --git a/docs/source/overview/gym/action_functors.md b/docs/source/overview/gym/action_functors.md index 375c0a3c9..bd971f87f 100644 --- a/docs/source/overview/gym/action_functors.md +++ b/docs/source/overview/gym/action_functors.md @@ -59,7 +59,7 @@ This page lists all available action terms that can be used with the Action Mana * - Action Term - Description * - {class}`~actions.EefPoseTerm` - - End-effector pose (6D or 7D) -> IK -> qpos. The policy outputs target end-effector poses which are converted to joint positions via inverse kinematics. Returns ``ik_success`` in the output so reward/observation can penalize or condition on IK failures. Supports both 6D (euler angles) and 7D (quaternion) pose representations. + - End-effector pose (6D or 7D) -> IK -> qpos. The policy outputs target end-effector poses which are converted to joint positions via inverse kinematics. Returns ``ik_success`` in the output so reward/observation can penalize or condition on IK failures. Supports both 6D (euler angles) and 7D (``x, y, z, qx, qy, qz, qw``) pose representations. ```json {"func": "EefPoseTerm", "params": {"scale": 0.1, "pose_dim": 7}} @@ -129,7 +129,7 @@ actions = { func="EefPoseTerm", params={ "scale": 0.1, - "pose_dim": 7, # 7D (position + quaternion) + "pose_dim": 7, # 7D (x, y, z, qx, qy, qz, qw) }, ), } diff --git a/docs/source/overview/gym/observation_functors.md b/docs/source/overview/gym/observation_functors.md index 3b6ead061..ba6f04a7c 100644 --- a/docs/source/overview/gym/observation_functors.md +++ b/docs/source/overview/gym/observation_functors.md @@ -25,7 +25,7 @@ This page lists all available observation functors that can be used with the Obs * - Functor Name - Description * - {func}`~observations.get_object_pose` - - Get the arena poses of objects. Returns 4x4 transformation matrices of shape (num_envs, 4, 4) by default, or (num_envs, 7) as [x, y, z, qw, qx, qy, qz] when ``to_matrix=False``. Returns zero tensor if object doesn't exist. + - Get the arena poses of objects. Returns 4x4 transformation matrices of shape (num_envs, 4, 4) by default, or (num_envs, 7) as [x, y, z, qx, qy, qz, qw] when ``to_matrix=False``. Returns zero tensor if object doesn't exist. ```json {"func": "get_object_pose", "mode": "add", @@ -33,7 +33,7 @@ This page lists all available observation functors that can be used with the Obs "params": {"entity_cfg": {"uid": "bottle"}, "to_matrix": true}} ``` * - {func}`~observations.get_rigid_object_pose` - - Get the arena poses of rigid objects. Returns 4x4 transformation matrices of shape (num_envs, 4, 4) by default, or (num_envs, 7) when ``to_matrix=False``. If the object doesn't exist, returns a zero tensor. (Deprecated: use ``get_object_pose`` instead.) + - Get the arena poses of rigid objects. Returns 4x4 transformation matrices of shape (num_envs, 4, 4) by default, or (num_envs, 7) as [x, y, z, qx, qy, qz, qw] when ``to_matrix=False``. If the object doesn't exist, returns a zero tensor. (Deprecated: use ``get_object_pose`` instead.) ```json {"func": "get_rigid_object_pose", "mode": "add", diff --git a/docs/source/overview/sim/atomic_actions/expert_programs.md b/docs/source/overview/sim/atomic_actions/expert_programs.md index 32dbceab8..d997b997b 100644 --- a/docs/source/overview/sim/atomic_actions/expert_programs.md +++ b/docs/source/overview/sim/atomic_actions/expert_programs.md @@ -116,7 +116,7 @@ targets: kind: cyclic_pose values: - position: [-0.40, 0.48, 0.10] - quaternion_wxyz: [1.0, 0.0, 0.0, 0.0] + quaternion_xyzw: [0.0, 0.0, 0.0, 1.0] program: kind: repeat count: 3 diff --git a/docs/source/overview/sim/planners/motion_generator.md b/docs/source/overview/sim/planners/motion_generator.md index 08f71b6fb..92b263251 100644 --- a/docs/source/overview/sim/planners/motion_generator.md +++ b/docs/source/overview/sim/planners/motion_generator.md @@ -66,7 +66,7 @@ sim_cfg = SimulationManagerCfg( width=1920, height=1080, physics_dt=1.0 / 100.0, - sim_device="cpu", + device="cpu", ) sim = SimulationManager(sim_cfg) @@ -92,7 +92,7 @@ robot_cfg = RobotCfg( dt=0.1, ) }, - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness={"Joint[1-6]": 1e4}, damping={"Joint[1-6]": 1e3}, ), diff --git a/docs/source/overview/sim/semantic_skills.md b/docs/source/overview/sim/semantic_skills.md index fdd611c26..39902490c 100644 --- a/docs/source/overview/sim/semantic_skills.md +++ b/docs/source/overview/sim/semantic_skills.md @@ -68,7 +68,7 @@ three curated call values: | {class}`HandOver` | Transfer a held object to another robot resource. | Uses a robot-profile-selected provider for the middle and default final pose; an explicit `final_target` overrides the latter. | {class}`SemanticPose` expresses an absolute object-space pose with a position -and normalized WXYZ quaternion. Scene objects and affordances use typed +and normalized XYZW quaternion. Scene objects and affordances use typed {class}`SceneObjectRef` and {class}`SceneAffordanceRef` values, so aliases are resolved at the registry boundary instead of being propagated into execution. @@ -203,7 +203,7 @@ calls = ( object=workpiece, at=SemanticPose( position=(-0.40, 0.48, 0.025), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ), ), ) diff --git a/docs/source/overview/sim/sim_articulation.md b/docs/source/overview/sim/sim_articulation.md index eec111ded..b47b4ec62 100644 --- a/docs/source/overview/sim/sim_articulation.md +++ b/docs/source/overview/sim/sim_articulation.md @@ -8,19 +8,19 @@ The {class}`~objects.Articulation` class represents the fundamental physics enti ## Configuration Articulations are configured using the {class}`~cfg.ArticulationCfg` dataclass. + | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | | `fpath` | `str` | `None` | Path to the asset file (URDF/USD). | | `init_pos` | `tuple` | `(0,0,0)` | Initial root position `(x, y, z)`. | | `init_rot` | `tuple` | `(0,0,0)` | Initial root rotation `(r, p, y)` in degrees. | -| `fix_base` | `bool` | `True` | Whether to fix the base of the articulation. | -| `use_usd_properties` | `bool` | `False` | If True, use physical properties from USD file; if False, override with config values. Only effective for usd files. | +| `root_props` | `ArticulationRootPropertiesCfg` | all fields `None` | Fixed-base/self-collision are portable; root sleep and paired solver iterations are Default-only and ignored by Newton. `None` preserves source/backend values. | +| `asset_physics_mode` | `"preserve" \| "overlay"` | `"preserve"` | Preserve source link/joint physics, or apply explicitly configured overlays after source resolution. | | `init_qpos` | `List[float]` | `None` | Initial joint positions. | -| `qpos_limits` | `Tensor` / `Dict[str, List[float]]` | `None` | Override joint position limits. Replaces asset limits and may either tighten or expand the range. | +| `qpos_limits` | `Tensor` / `Dict[str, List[float]]` | `None` | Override limits by flattened source-resolved DOF order or joint-name/regex rules before backend build. | | `body_scale` | `List[float]` | `[1.0, 1.0, 1.0]` | Scaling factors for the articulation links. | -| `disable_self_collisions` | `bool` | `True` | Whether to disable self-collisions. | -| `drive_pros` | `JointDrivePropertiesCfg` | `drive_type="none"` | Default drive properties. | -| `attrs` | `RigidBodyAttributesCfg` | `...` | Default rigid body attributes applied to all links. | +| `joint_drive_props` | `JointDrivePropertiesCfg` | `None` | Optional sparse joint drive, limit, friction, and armature overlay. | +| `attrs` | `RigidBodyPhysicsCfg` | empty groups | Grouped rigid-body physics applied to all links. | | `link_attrs` | `dict[str, LinkPhysicsOverrideCfg]` | `None` | Optional per-link overrides keyed by group name; each group matches link names via regex. | @@ -32,20 +32,23 @@ override specific links (matched by regex, same rules as joint drive dict keys): ```python from embodichain.lab.sim.cfg import ( ArticulationCfg, + CollisionPropertiesCfg, LinkPhysicsOverrideCfg, - RigidBodyAttributesCfg, - RigidBodyAttributesOverrideCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, ) art_cfg = ArticulationCfg( fpath="path/to/robot.urdf", - attrs=RigidBodyAttributesCfg(static_friction=0.5), + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(static_friction=0.5), + ), link_attrs={ "eef": LinkPhysicsOverrideCfg( link_names_expr=[".*(hand|finger|ee).*"], - attrs=RigidBodyAttributesOverrideCfg( - static_friction=0.95, - contact_offset=0.001, + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(static_friction=0.95), + collision_props=CollisionPropertiesCfg(contact_offset=0.001), ), ), }, @@ -57,7 +60,7 @@ for the same partial-override behavior. ### Drive Configuration -The `drive_pros` parameter controls the joint physics behavior. It is defined using the `JointDrivePropertiesCfg` class. Generic articulations default to `drive_type="none"`, so passive assets such as cabinets and drawers do not receive internal drive forces unless explicitly configured. +The `joint_drive_props` parameter controls the joint physics behavior. It is defined using the `JointDrivePropertiesCfg` class. Generic articulations default to `drive_type="none"`, so passive assets such as cabinets and drawers do not receive internal drive forces unless explicitly configured. | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | @@ -117,18 +120,19 @@ articulation layer. ```python import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.objects import Articulation, ArticulationCfg +from embodichain.lab.sim.cfg import ArticulationCfg, ArticulationRootPropertiesCfg +from embodichain.lab.sim.objects import Articulation # 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 art_cfg = ArticulationCfg( fpath="assets/robots/franka/franka.urdf", init_pos=(0, 0, 0.5), - fix_base=True + root_props=ArticulationRootPropertiesCfg(fixed_base=True), ) # 3. Spawn Articulation @@ -151,7 +155,7 @@ from embodichain.data import get_data_path usd_art_cfg = ArticulationCfg( fpath=get_data_path("path/to/robot.usd"), init_pos=(0, 0, 0.5), - use_usd_properties=True # Keep USD drive/physics properties + asset_physics_mode="preserve", ) usd_robot = sim.add_articulation(cfg=usd_art_cfg) @@ -159,8 +163,8 @@ usd_robot = sim.add_articulation(cfg=usd_art_cfg) usd_art_cfg_override = ArticulationCfg( fpath=get_data_path("path/to/robot.usd"), init_pos=(0, 0, 0.5), - use_usd_properties=False, # Use config instead - drive_pros=JointDrivePropertiesCfg(stiffness=5000, damping=500) + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(stiffness=5000, damping=500), ) robot = sim.add_articulation(cfg=usd_art_cfg_override) ``` @@ -179,8 +183,8 @@ State data is accessed via getter methods that return batched tensors (`N` envir | Method | Shape / Return Type | Description | | :--- | :--- | :--- | -| `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Root link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | -| `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qw, qx, qy, qz]` or a 4x4 matrix. | +| `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Root link pose `[x, y, z, qx, qy, qz, qw]` or a 4x4 matrix. | +| `get_link_pose(link_name, to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Specific link pose `[x, y, z, qx, qy, qz, qw]` or a 4x4 matrix. | | `get_qpos(target=False)` | `(N, dof)` | Current joint positions (or joint targets if `target=True`). | | `get_qvel(target=False)` | `(N, dof)` | Current joint velocities (or velocity targets if `target=True`). | | `get_joint_drive()` | `Tuple[Tensor, ...]` | Returns `(stiffness, damping, max_effort, max_velocity, friction, armature)`, each shaped `(N, dof)`. | @@ -244,8 +248,8 @@ sim.update() ### Pose Control ```python # Teleport the articulation root to a new pose -# shape: (N, 7) formatted as [x, y, z, qw, qx, qy, qz] -new_root_pose = torch.tensor([[0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0]], device=device).repeat(sim.num_envs, 1) +# shape: (N, 7) formatted as [x, y, z, qx, qy, qz, qw] +new_root_pose = torch.tensor([[0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0]], device=device).repeat(sim.num_envs, 1) articulation.set_local_pose(new_root_pose) ``` diff --git a/docs/source/overview/sim/sim_assets.md b/docs/source/overview/sim/sim_assets.md index 30ab3e0b4..bd21ef75f 100644 --- a/docs/source/overview/sim/sim_assets.md +++ b/docs/source/overview/sim/sim_assets.md @@ -98,31 +98,31 @@ Configured via {class}`~cfg.RigidObjectCfg`. | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | | `shape` | `ShapeCfg` | `ShapeCfg()` | Shape configuration (e.g., Mesh, Box). | -| `attrs` | `RigidBodyAttributesCfg` | `RigidBodyAttributesCfg()` | Physical attributes. | +| `attrs` | `RigidBodyPhysicsCfg` | `RigidBodyPhysicsCfg()` | Grouped physical attributes. | | `body_type` | `Literal` | `"dynamic"` | "dynamic", "kinematic", or "static". | -| `max_convex_hull_num` | `int` | `1` | Max convex hulls for decomposition (CoACD). | -| `sdf_resolution` | `int` | `0` | Resolution for signed distance field. In most cases, a resolution of around 250 produces good results; resolutions exceeding 1000 are rarely necessary.| +| `shape.collision` | `MeshCollisionCfg \| None` | `None` | Explicit mesh collision geometry: convex hull, convex decomposition, triangle mesh, or SDF. `None` uses one convex hull. | | `body_scale` | `tuple` | `(1.0, 1.0, 1.0)` | Scale of the rigid body. | -### Rigid Body Attributes +### Rigid Body Physics -The {class}`~cfg.RigidBodyAttributesCfg` class defines physical properties for rigid bodies. +{class}`~cfg.RigidBodyPhysicsCfg` keeps physical settings in optional groups. +An unset field leaves the source asset or backend default intact, which makes +the same configuration usable as either a complete procedural definition or a +sparse USD/URDF overlay. -| Parameter | Type | Default | Description | -| :--- | :--- | :--- | :--- | -| `mass` | `float` | `1.0` | Mass in kg. Set to 0 to use density. | -| `density` | `float` | `1000.0` | Density in kg/m^3. | -| `angular_damping` | `float` | `0.7` | Angular damping coefficient. | -| `linear_damping` | `float` | `0.7` | Linear damping coefficient. | -| `max_depenetration_velocity` | `float` | `10.0` | Maximum depenetration velocity. | -| `sleep_threshold` | `float` | `0.001` | Threshold below which the body can go to sleep. | -| `enable_ccd` | `bool` | `False` | Enable continuous collision detection. | -| `contact_offset` | `float` | `0.002` | Contact offset for collision detection. | -| `rest_offset` | `float` | `0.001` | Rest offset for collision detection. | -| `enable_collision` | `bool` | `True` | Enable collision for the rigid body. | -| `restitution` | `float` | `0.0` | Restitution (bounciness) coefficient. | -| `dynamic_friction` | `float` | `0.5` | Dynamic friction coefficient. | -| `static_friction` | `float` | `0.5` | Static friction coefficient. | +| Group | Type | Contents | +| :--- | :--- | :--- | +| `mass_props` | `MassPropertiesCfg` | Mass, density, inertia, and COM pose. | +| `rigid_props` | `DefaultRigidBodyPropertiesCfg` | Rigid-body behavior such as damping, CCD, and solver iterations. | +| `collision_props` | `CollisionPropertiesCfg` | Collision enablement, contact/rest offsets, and concrete-backend contact properties. | +| `material_props` | `RigidBodyMaterialCfg` | Restitution, friction, and concrete-backend material properties. | + +COM quaternions in configuration use `xyzw`. The Spawn adapter converts to the +native backend order only when it writes an engine descriptor. + +Mesh cooking is owned by `MeshCfg.collision`, not by rigid-body physics. Its +`approximation` field selects the representation explicitly; strategy-specific +fields such as `max_hulls` and `sdf_resolution` are validated against it. For a runnable rigid-object example, see the {doc}`Create Scene ` tutorial. diff --git a/docs/source/overview/sim/sim_cloth.md b/docs/source/overview/sim/sim_cloth.md index dfc19caad..41f7983db 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) @@ -173,7 +173,7 @@ You can set the global pose of a cloth object (which transforms all its vertices ```python # Reset or Move the Cloth Object -target_pose = torch.tensor([[0, 0, 1.0, 1, 0, 0, 0]], device=device) # (x, y, z, qw, qx, qy, qz) +target_pose = torch.tensor([[0, 0, 1.0, 0, 0, 0, 1]], device=device) # (x, y, z, qx, qy, qz, qw) cloth_object.set_local_pose(target_pose) # Important: Step simulation to apply changes diff --git a/docs/source/overview/sim/sim_manager.md b/docs/source/overview/sim/sim_manager.md index cd280c26e..c9d82d491 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,16 +42,29 @@ 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. | +| `physics_cfg` | `DefaultPhysicsCfg` \| `NewtonPhysicsCfg` | `DefaultPhysicsCfg()` | Physics backend configuration (class selects default vs Newton). | | `profiler` | `ProfilerCfg` \| `None` | `None` | Optional hierarchical wall-time profiler for simulation updates. | -| `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. | | `visualization` | `VisualizationCfg` | `VisualizationCfg()` | Browser visualization, opt-in Gizmo commands, and Viser server settings. | ### Physics Configuration -The {class}`~cfg.PhysicsCfg` class controls the global physics simulation parameters. +Use {class}`~cfg.DefaultPhysicsCfg` for the Default backend or {class}`~cfg.NewtonPhysicsCfg` for the Newton backend. Both are integrated through the DexSim runtime. GPU memory settings are on {class}`~cfg.DefaultPhysicsCfg` as ``gpu_memory``. + +`default` and `newton` are the only public physics-backend identifiers. +Backend-neutral nested property groups may additionally use `common`. DexSim +is the runtime and Spawn SDK integration layer, not another selectable physics +backend; SDK-native `Dexsim*Desc` names remain confined to that adapter boundary. + +All physics backends inherit these base parameters from {class}`~cfg.PhysicsBackendCfg`: + +| 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. | + +#### Default Backend + +The {class}`~cfg.DefaultPhysicsCfg` class controls the global default-backend physics simulation parameters. | Parameter | Type | Default | Description | | :--- | :--- | :--- | :--- | @@ -60,7 +76,98 @@ The {class}`~cfg.PhysicsCfg` class controls the global physics simulation parame PCM and TGS remain enabled, enhanced determinism remains disabled, and friction is evaluated on every solver iteration. These solver implementation details use -fixed defaults and are not exposed by `PhysicsCfg`. +fixed defaults and are not exposed by `DefaultPhysicsCfg`. + +#### Newton Backend and Automatic Solver Selection + +Use {class}`~cfg.NewtonPhysicsCfg` to enable the Newton backend. Its +`solver_cfg` defaults to `None` intentionally: EmbodiChain leaves the +`solver_cfg` argument unset when it creates DexSim's `NewtonCfg`, preserving +DexSim's `AutoSolverCfg` default. + +```python +from embodichain.lab.sim import SimulationManagerCfg +from embodichain.lab.sim.cfg import NewtonPhysicsCfg + +sim_config = SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + device="cuda:0", + physics_dt=0.01, + num_substeps=10, + ) +) +``` + +AutoSolver is resolved when DexSim finalizes the complete Spawn scene during +{meth}`SimulationManager.prepare`. Add all initial robots and objects before +calling `prepare()` so the selection sees the complete scene. An explicit +`{"solver_type": "auto"}` or `{"class_type": "AutoSolverCfg"}` mapping has +the same effect as leaving `solver_cfg` unset. + +:::{important} +This integration requires a DexSim build that exports `AutoSolverCfg`. +EmbodiChain does not fall back to a hard-coded concrete solver when that API is +unavailable. +::: + +DexSim applies the following scene-content rules. Independent rigid objects and +articulation links are classified separately. + +| Finalized scene contents | Selected configuration | Solver type | Active collision path | +| :--- | :--- | :--- | :--- | +| Empty scene or independent rigid bodies only | `XPBDSolverCfg` | `xpbd` | Newton collision pipeline | +| Articulations, with or without independent rigid bodies | `MJWarpSolverCfg` | `mujoco_warp` | MuJoCo Warp collision pipeline | +| Cloth or soft bodies, optionally with rigid bodies | `VBDSolverCfg` | `vbd` | Newton collision pipeline; VBD may handle deformable self-contact | +| Cloth or soft bodies with articulations, optionally with rigid bodies | `MJVBDSolverCfg` | `mjvbd` | Newton particle-shape soft contacts; MuJoCo rigid collision is disabled | +| Fluid particles, optionally with rigid SDF boundaries | `SPHSolverCfg` | `sph` | SPH one-way SDF boundary handling; rigid contacts are not consumed | +| MPM particles, optionally with rigid colliders | `ImplicitMPMSolverCfg` | `implicit_mpm` | Implicit-MPM collider projection; the rigid collision pipeline is not stepped | + +The current MJVBD path does not generate rigid-rigid or rigid-ground contacts. +MuJoCo Warp still advances rigid bodies and articulations, while Newton's soft +contact kernels handle deformable particle-shape contacts. + +:::{note} +The table documents DexSim's resolver. EmbodiChain currently exposes Newton +runtime adapters for rigid bodies and articulations. Newton soft-body and cloth +adapters remain disabled, and fluid/MPM assets do not yet have public +EmbodiChain APIs; those rows describe upstream selection behavior rather than +an EmbodiChain support guarantee. +::: + +AutoSolver rejects scene combinations for which one solver cannot represent +all coupled systems: + +- more than one particle family among deformable, fluid, and MPM; +- fluid particles combined with articulations; +- MPM particles combined with articulations. + +Selection is based on scene contents, not the configured device. DexSim reports +device incompatibility after resolution; cloth, soft-body, fluid, and MPM +solvers currently require CUDA. The selected type is also written to the +DexSim log, for example `Newton AutoSolver selected 'mujoco_warp'.` + +Pass a concrete solver configuration when an algorithm or solver-specific +parameter must be fixed: + +```python +sim_config = SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + device="cpu", + solver_cfg={ + "solver_type": "xpbd", + "iterations": 8, + }, + ) +) +``` + +EmbodiChain mapping configs recognize `auto`, `mujoco_warp` (or `mjwarp`), +`xpbd`, `semi_implicit`, `featherstone`, and `vbd`. A DexSim +`NewtonSolverCfg` object may also be assigned directly when another explicit +solver class is required. AutoSolver never selects `DFSPHSolverCfg`, +`FeatherstoneSolverCfg`, or `SemiImplicitSolverCfg`. In particular, +`requires_grad=True` requires an explicit `semi_implicit` configuration; +automatic selection is rejected for differentiable simulation. ### Render Configuration @@ -239,14 +346,14 @@ EmbodiChain supports importing USD files (`.usd`, `.usda`, `.usdc`) for both rig # Import rigid object with USD properties rigid_cfg = RigidObjectCfg( shape=MeshCfg(fpath=get_data_path("path/to/object.usd")), - use_usd_properties=True # Use properties from USD file + asset_physics_mode="preserve", ) obj = sim.add_rigid_object(cfg=rigid_cfg) # Import articulation with USD properties robot_cfg = ArticulationCfg( fpath=get_data_path("path/to/robot.usd"), - use_usd_properties=True # Use joint drive properties from USD + asset_physics_mode="preserve", ) robot = sim.add_articulation(cfg=robot_cfg) ``` @@ -291,7 +398,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 diff --git a/docs/source/overview/sim/sim_rigid_object.md b/docs/source/overview/sim/sim_rigid_object.md index b80f267ba..6106c66f8 100644 --- a/docs/source/overview/sim/sim_rigid_object.md +++ b/docs/source/overview/sim/sim_rigid_object.md @@ -13,30 +13,27 @@ Configured via the {class}`~cfg.RigidObjectCfg` class. | :--- | :--- | :--- | :--- | | `shape` | {class}`~shapes.ShapeCfg` | `ShapeCfg()` | Geometry configuration for visual and collision shapes. Use `MeshCfg` for mesh files or primitive cfgs (e.g., `CubeCfg`). | | `body_type` | `Literal["dynamic","kinematic","static"]` | `"dynamic"` | Actor type for the rigid body. See `{class}`~cfg.RigidObjectCfg.to_dexsim_body_type` for conversion. | -| `attrs` | {class}`~cfg.RigidBodyAttributesCfg` | defaults in code | Physical attributes (mass, damping, friction, restitution, collision offsets, CCD, etc.). | +| `attrs` | {class}`~cfg.RigidBodyPhysicsCfg` | empty groups | Grouped physical attributes (mass, damping, friction, restitution, collision offsets, CCD, etc.). | | `init_pos` | `Sequence[float]` | `(0,0,0)` | Initial root position (x, y, z). | | `init_rot` | `Sequence[float]` | `(0,0,0)` (Euler degrees) | Initial root orientation (Euler angles in degrees) or provide `init_local_pose`. | -| `use_usd_properties` | `bool` | `False` | If True, use physical properties from USD file; if False, override with config values. Only effective for usd files. | +| `asset_physics_mode` | {class}`~cfg.AssetPhysicsMode` | `"preserve"` | Preserve source-authored physics or overlay explicitly configured values. | | `uid` | `str` | `None` | Optional unique identifier for the object; manager will assign one if omitted. | -### Rigid Body Attributes ({class}`~cfg.RigidBodyAttributesCfg`) +### Rigid Body Physics ({class}`~cfg.RigidBodyPhysicsCfg`) -The full attribute set lives in `{class}`~cfg.RigidBodyAttributesCfg`. Common fields shown in code include: +Physical properties are grouped by intent. Every field is optional: `None` +means that a source asset or the active backend keeps ownership of that value. -| Parameter | Type | Default (from code) | Description | -| :--- | :--- | :---: | :--- | -| `mass` | `float` | `1.0` | Mass of the rigid body in kilograms (set to 0 to use density). | -| `density` | `float` | `1000.0` | Density used when mass is negative/zero. | -| `linear_damping` | `float` | `0.7` | Linear damping coefficient. | -| `angular_damping` | `float` | `0.7` | Angular damping coefficient. | -| `dynamic_friction` | `float` | `0.5` | Dynamic friction coefficient. | -| `static_friction` | `float` | `0.5` | Static friction coefficient. | -| `restitution` | `float` | `0.0` | Restitution (bounciness). | -| `contact_offset` | `float` | `0.002` | Contact offset for collision detection. | -| `rest_offset` | `float` | `0.001` | Rest offset for collision detection. | -| `enable_ccd` | `bool` | `False` | Enable continuous collision detection. | +| Group | Example fields | +| :--- | :--- | +| `mass_props` | `mass`, `density`, `inertia`, `com_position`, `com_quaternion` | +| `rigid_props` | `linear_damping`, `angular_damping`, `enable_ccd` | +| `collision_props` | `collision_enabled`, `contact_offset`, `rest_offset` | +| `material_props` | `dynamic_friction`, `static_friction`, `restitution` | -Use the `.attr()` helper to convert to `dexsim.PhysicalAttr` when interfacing with the engine. +COM quaternions are always authored in `xyzw` order. Native engine attributes +are an internal adapter detail. Backend-specific values use the concrete type in +the corresponding property slot rather than a second backend block. ## Setup & Initialization @@ -45,15 +42,26 @@ import torch from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg from embodichain.lab.sim.shapes import CubeCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.cfg import ( + MassPropertiesCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, +) # 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) -physics_attrs = RigidBodyAttributesCfg(mass=1.0, dynamic_friction=0.5, static_friction=0.5, restitution=0.1) +physics_attrs = RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg( + dynamic_friction=0.5, + static_friction=0.5, + restitution=0.1, + ), +) cfg = RigidObjectCfg( uid="cube", @@ -85,7 +93,7 @@ from embodichain.data import get_data_path usd_cfg = RigidObjectCfg( shape=MeshCfg(fpath=get_data_path("path/to/object.usd")), body_type="dynamic", - use_usd_properties=True # Keep USD properties + asset_physics_mode="preserve", # Keep USD properties ) obj = sim.add_rigid_object(cfg=usd_cfg) @@ -93,8 +101,8 @@ obj = sim.add_rigid_object(cfg=usd_cfg) usd_cfg_override = RigidObjectCfg( shape=MeshCfg(fpath=get_data_path("path/to/object.usd")), body_type="dynamic", - use_usd_properties=False, # Use config instead - attrs=RigidBodyAttributesCfg(mass=2.0) + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=2.0)), ) obj2 = sim.add_rigid_object(cfg=usd_cfg_override) ``` @@ -107,18 +115,18 @@ Rigid objects are observed and controlled via single poses and linear/angular ve | Method / Property | Return / Args | Description | | :--- | :--- | :--- | -| `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Get object local pose as (x, y, z, qw, qx, qy, qz) or 4x4 matrix per environment. | +| `get_local_pose(to_matrix=False)` | `(N, 7)` or `(N, 4, 4)` | Get object local pose as (x, y, z, qx, qy, qz, qw) or 4x4 matrix per environment. | | `set_local_pose(pose, env_ids=None)` | `pose: (N, 7)` or `(N, 4, 4)` | Teleport object to given pose (requires calling `sim.update()` to apply). | -| `body_data.pose` | `(N, 7)` | Access object pose directly (for dynamic/kinematic bodies). | +| `body_data.pose` | `(N, 7)` | Access object pose as `[x, y, z, qx, qy, qz, qw]` (for dynamic/kinematic bodies). | | `body_data.lin_vel` | `(N, 3)` | Access linear velocity of object root (for dynamic bodies). | | `body_data.ang_vel` | `(N, 3)` | Access angular velocity of object root (for dynamic bodies). | | `body_data.vel` | `(N, 6)` | Concatenated linear and angular velocities. | | `body_data.lin_acc` | `(N, 3)` | Access linear acceleration of object root (for dynamic bodies). | | `body_data.ang_acc` | `(N, 3)` | Access angular acceleration of object root (for dynamic bodies). | | `body_data.acc` | `(N, 6)` | Concatenated linear and angular accelerations. | -| `body_data.com_pose` | `(N, 7)` | Get center of mass pose of rigid bodies. | -| `body_data.default_com_pose` | `(N, 7)` | Default center of mass pose. | -| `body_state` | `(N, 13)` | Get full body state: [x, y, z, qw, qx, qy, qz, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z]. | +| `body_data.com_pose` | `(N, 7)` | Get center of mass pose as `[x, y, z, qx, qy, qz, qw]`. | +| `body_data.default_com_pose` | `(N, 7)` | Default center of mass pose as `[x, y, z, qx, qy, qz, qw]`. | +| `body_state` | `(N, 13)` | Get full body state: [x, y, z, qx, qy, qz, qw, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z]. | ### Dynamics Control @@ -132,7 +140,7 @@ Rigid objects are observed and controlled via single poses and linear/angular ve | Method / Property | Return / Args | Description | | :--- | :--- | :--- | -| `set_attrs(attrs, env_ids=None)` | `attrs: RigidBodyAttributesCfg` | Set physical attributes (mass, friction, damping, etc.). | +| `set_attrs(attrs, env_ids=None)` | `attrs: RigidBodyPhysicsCfg` | Set grouped physical attributes (mass, friction, damping, etc.). | | `set_mass(mass, env_ids=None)` | `mass: (N,)` | Set mass for rigid object. | | `get_mass(env_ids=None)` | `(N,)` | Get mass for rigid object. | | `set_friction(friction, env_ids=None)` | `friction: (N,)` | Set dynamic and static friction. | @@ -185,7 +193,7 @@ When a rigid object is loaded, its material assignment is captured without repla ### Observation Shapes -- Pose: `(N, 7)` per-object pose (position + quaternion). +- Pose: `(N, 7)` per-object pose `[x, y, z, qx, qy, qz, qw]`. - Velocities: `(N, 3)` for linear and angular velocities respectively. N denotes the number of parallel environments when using vectorized simulation (`SimulationManagerCfg.num_envs`). @@ -195,8 +203,8 @@ N denotes the number of parallel environments when using vectorized simulation ( - When moving objects programmatically via `set_local_pose`, call `sim.update()` (or step the sim) to ensure transforms and collision state are synchronized. - 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. +- For complex meshes, configure `MeshCfg.collision` with `approximation="convex_decomposition"` and a bounded `max_hulls`, or provide a simplified collision mesh. +- 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 e6d61f82f..7287ebb8d 100644 --- a/docs/source/overview/sim/sim_rigid_object_group.md +++ b/docs/source/overview/sim/sim_rigid_object_group.md @@ -19,7 +19,7 @@ Configured via the {class}`~cfg.RigidObjectGroupCfg` class. | `ext` | `str` | `".obj"` | File extension filter when loading assets from `folder_path`. | | `init_pos` / `init_rot` | `Sequence` (optional) | group-level transform | Optional transform to apply as a base offset to all members. | -Refer to {class}`~cfg.RigidObjectCfg` and {class}`~cfg.RigidBodyAttributesCfg` for per-member configuration options (mass, friction, restitution, collision options, shapes, etc.). +Refer to {class}`~cfg.RigidObjectCfg` and {class}`~cfg.RigidBodyPhysicsCfg` for per-member configuration options (mass, friction, restitution, collision options, shapes, etc.). ### Folder-based initialization @@ -40,19 +40,25 @@ from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.objects import ( RigidObjectGroup, RigidObjectGroupCfg, RigidObjectCfg ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg +from embodichain.lab.sim.cfg import ( + MassPropertiesCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, +) # 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 -physics_attrs = RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, +physics_attrs = RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg( + dynamic_friction=0.5, + static_friction=0.5, + restitution=0.1, + ), ) # 3. Create group config with multiple members @@ -83,9 +89,9 @@ A group provides batch operations on multiple rigid objects. Key APIs include: | :--- | :--- | :--- | | `num_objects` | `int` | Number of objects in each group instance. | | `body_data` | `RigidBodyGroupData` | Data manager providing `pose`, `lin_vel`, `ang_vel` properties. | -| `body_state` | `(N, M, 13)` | Full body state of all members: [x, y, z, qw, qx, qy, qz, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z]. | -| `get_local_pose(to_matrix=False)` | `(N, M, 7)` or `(N, M, 4, 4)` | Poses of all members across N envs; M = number of members. | -| `set_local_pose(pose, env_ids=None, obj_ids=None)` | `pose: (N, M, 7)` or `(N, M, 4, 4)` | Set poses for specific environments and/or objects; requires `sim.update()` to apply. | +| `body_state` | `(N, M, 13)` | Full body state of all members: [x, y, z, qx, qy, qz, qw, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z]. | +| `get_local_pose(to_matrix=False)` | `(N, M, 7)` or `(N, M, 4, 4)` | Poses of all members as `[x, y, z, qx, qy, qz, qw]` or matrices; M = number of members. | +| `set_local_pose(pose, env_ids=None, obj_ids=None)` | `pose: (N, M, 7)` or `(N, M, 4, 4)` | Set poses in `[x, y, z, qx, qy, qz, qw]` or matrix form; requires `sim.update()` to apply. | | `get_user_ids()` | `(N, M)` | Get user IDs tensor for all members in the group. | | `clear_dynamics(env_ids=None)` | - | Reset velocities and clear all forces/torques for the group. | | `set_visual_material(mat, env_ids=None)` | `mat: VisualMaterial` | Change visual appearance for all members. | @@ -106,10 +112,10 @@ Use these shapes when collecting vectorized observations for multi-environment t - Groups are convenient for batch operations: resetting, setting visibility, and applying transforms to multiple objects together. - Use `obj_ids` parameter in `set_local_pose()` to control specific objects within the group rather than all members. -- Prefer providing simplified collision meshes or enabling convex decomposition (`max_convex_hull_num` > 1) for complex visual meshes to improve physics stability. +- Prefer simplified collision meshes or an explicit `MeshCfg.collision` convex-decomposition strategy with a bounded `max_hulls` for complex visual meshes. - `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 2e8e36261..5574858f7 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 @@ -79,7 +79,7 @@ print(f"EE Pose: {ee_pose}") Compute the required joint positions to reach a target pose. ```python # Compute IK -# pose: Target pose (N, 7) or (N, 4, 4) +# pose: Target pose (N, 7) as [x, y, z, qx, qy, qz, qw], or (N, 4, 4) target_pose = ee_pose.clone() # Example target target_pose[:, 2] += 0.1 # Move up 10cm diff --git a/docs/source/overview/sim/sim_sensor.md b/docs/source/overview/sim/sim_sensor.md index 1a8388f5f..f3ea1b724 100644 --- a/docs/source/overview/sim/sim_sensor.md +++ b/docs/source/overview/sim/sim_sensor.md @@ -33,7 +33,7 @@ The `ExtrinsicsCfg` class defines the position and orientation of the camera. | :--- | :--- | :--- | :--- | | `parent` | `str` | `None` | Name of the link to attach to (e.g., `"ee_link"`). If `None`, camera is fixed in world. | | `pos` | `list` | `[0.0, 0.0, 0.0]` | Position offset `[x, y, z]`. | -| `quat` | `list` | `[1.0, 0.0, 0.0, 0.0]` | Orientation quaternion `[w, x, y, z]`. | +| `quat` | `list` | `[0.0, 0.0, 0.0, 1.0]` | Orientation quaternion `[x, y, z, w]`. | | `eye` | `tuple` | `None` | (Optional) Camera eye position for look-at mode. | | `target` | `tuple` | `None` | (Optional) Target position for look-at mode. | | `up` | `tuple` | `None` | (Optional) Up vector for look-at mode. | @@ -55,7 +55,7 @@ camera_cfg = CameraCfg( extrinsics=CameraCfg.ExtrinsicsCfg( parent="ee_link", # Attach to robot end-effector pos=[0.09, 0.05, 0.04], # Relative position - quat=[0, 1, 0, 0], # Relative rotation [w, x, y, z] + quat=[1, 0, 0, 0], # Relative rotation [x, y, z, w] ), enable_color=True, enable_depth=True, @@ -236,4 +236,4 @@ env_positions = contact_report["position"][env_id, :num_valid] ### Additional Methods - **`filter_by_user_ids(item_user_ids, env_ids=None)`**: Filter contact report to include only contacts involving specific user IDs. Optionally filter by specific environment IDs. -- **`set_contact_point_visibility(visible, rgba, point_size, env_ids=None)`**: Enable/disable visualization of contact points with customizable color and size. Optionally visualize only specific environments. \ No newline at end of file +- **`set_contact_point_visibility(visible, rgba, point_size, env_ids=None)`**: Enable/disable visualization of contact points with customizable color and size. Optionally visualize only specific environments. diff --git a/docs/source/overview/sim/sim_soft_object.md b/docs/source/overview/sim/sim_soft_object.md index 5936321d4..5d9fa04d8 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 @@ -118,7 +118,7 @@ You can set the global pose of a soft object (which transforms all its vertices) ```python # Reset or Move the Soft Object -target_pose = torch.tensor([[0, 0, 1.0, 1, 0, 0, 0]], device=device) # (x, y, z, qw, qx, qy, qz) +target_pose = torch.tensor([[0, 0, 1.0, 0, 0, 0, 1]], device=device) # (x, y, z, qx, qy, qz, qw) soft_object.set_local_pose(target_pose) # Important: Step simulation to apply changes diff --git a/docs/source/overview/sim/viser_visualization.md b/docs/source/overview/sim/viser_visualization.md index efefaea50..2b1411270 100644 --- a/docs/source/overview/sim/viser_visualization.md +++ b/docs/source/overview/sim/viser_visualization.md @@ -177,7 +177,7 @@ sampled independently from rigid-body poses: - **Cloth** uses the physical cloth vertices and a welded mapping of the source render triangles. Its browser topology matches the simulated surface. -- **Soft bodies** expose live PhysX collision vertices through DexSim, but +- **Soft bodies** expose live DexSim collision vertices, but DexSim does not expose the collision triangle connectivity. EmbodiChain therefore visualizes a stable convex-hull surface over those vertices. The preview follows deformation but omits concave render-mesh details. diff --git a/docs/source/resources/robot/cobotmagic.md b/docs/source/resources/robot/cobotmagic.md index de23dd2a6..d17608f6d 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) @@ -56,7 +56,7 @@ robot = sim.add_robot(cfg=CobotMagicCfg().from_dict({})) - **urdf_cfg**: URDF configuration, supports multi-component assembly (e.g., dual arms) - **control_parts**: Control groups for independent control of each arm and gripper - **solver_cfg**: Inverse kinematics solver configuration, customizable end-effector and base -- **drive_pros**: Joint drive properties (stiffness, damping, max effort, etc.) +- **joint_drive_props**: Joint drive properties (stiffness, damping, max effort, etc.) - **attrs**: Rigid body physical attributes (mass, friction, damping, etc.) ### 2. Custom Usage Example diff --git a/docs/source/resources/task/index.rst b/docs/source/resources/task/index.rst index e3241a8eb..1f6a54e9f 100644 --- a/docs/source/resources/task/index.rst +++ b/docs/source/resources/task/index.rst @@ -72,9 +72,6 @@ Environment catalog * - Tableware - ``StackCups-v1`` - ``embodichain_tasks/configs/tasks/manipulation/tableware/stack_cups/env.json`` - * - Tableware - - ``Rearrangement-v3`` - - ``embodichain_tasks/configs/tasks/manipulation/tableware/rearrangement/env.json`` The value of ``id`` inside a conventional gym config must match a registered environment ID. A supported ``expert_program_runtime`` config registers its own diff --git a/docs/source/tutorial/articulation.rst b/docs/source/tutorial/articulation.rst index c5477c064..50645a39b 100644 --- a/docs/source/tutorial/articulation.rst +++ b/docs/source/tutorial/articulation.rst @@ -44,7 +44,7 @@ Loading the URDF Resolve the bundled drawer asset, then pass its path to :class:`cfg.ArticulationCfg`. The example intentionally does not set -``drive_pros``. Therefore the configuration uses the Articulation default, +``joint_drive_props``. Therefore the configuration uses the Articulation default, ``drive_type="none"``. ``SimulationManager.add_articulation`` loads one drawer into each configured environment and returns a batched :class:`objects.Articulation` handle. @@ -62,7 +62,7 @@ effective physics limit used by both the backend and the force-control loop. Verifying the constructed drive type ------------------------------------ -Checking ``articulation.cfg.drive_pros.drive_type`` confirms the requested +Checking ``articulation.cfg.joint_drive_props.drive_type`` confirms the requested configuration, but it does not prove what the physics backend received. The example therefore calls :meth:`objects.Articulation.get_joint_drive_type`, which reads the drive type from every constructed DexSim entity. It raises an @@ -158,7 +158,7 @@ articulation needs an actuator, opt in with articulation_cfg = ArticulationCfg( fpath="path/to/articulation.urdf", - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( drive_type="force", stiffness=1.0e4, damping=1.0e3, @@ -170,8 +170,8 @@ For a controllable robot, prefer :class:`cfg.RobotCfg` and .. attention:: - For USD assets, ``use_usd_properties=True`` preserves the drive types stored - in the USD file instead of applying the Articulation configuration default. + For file-backed assets, ``asset_physics_mode="preserve"`` keeps source + physics, while ``"overlay"`` applies explicitly configured values. Next Steps ~~~~~~~~~~ diff --git a/docs/source/tutorial/create_cloth.rst b/docs/source/tutorial/create_cloth.rst index 181a1eaf4..d5aa28439 100644 --- a/docs/source/tutorial/create_cloth.rst +++ b/docs/source/tutorial/create_cloth.rst @@ -68,7 +68,7 @@ The grid mesh generated earlier is saved to disk and then passed to :meth:`Simul Adding a rigid body for interaction ------------------------------------- -A small cubic rigid body (``padding_box``) is placed beneath the cloth so the cloth drapes over it. It is added with :meth:`SimulationManager.add_rigid_object` using :class:`cfg.RigidObjectCfg` and :class:`cfg.RigidBodyAttributesCfg`: +A small cubic rigid body (``padding_box``) is placed beneath the cloth so the cloth drapes over it. It is added with :meth:`SimulationManager.add_rigid_object` using :class:`cfg.RigidObjectCfg` and :class:`cfg.RigidBodyPhysicsCfg`: - :class:`cfg.CubeCfg` — defines the box dimensions - ``body_type="dynamic"`` — the box responds to physics; change to ``"static"`` for a fixed obstacle diff --git a/docs/source/tutorial/rigid_constraint.rst b/docs/source/tutorial/rigid_constraint.rst index 500558a6c..15afc78d2 100644 --- a/docs/source/tutorial/rigid_constraint.rst +++ b/docs/source/tutorial/rigid_constraint.rst @@ -47,7 +47,7 @@ Adding two cubes Two dynamic cubes are added with :meth:`SimulationManager.add_rigid_object`. Each uses a :class:`CubeCfg` shape (a primitive cube, so no mesh asset file is -needed) and a :class:`RigidBodyAttributesCfg` for mass and friction. ``cube_a`` +needed) and a :class:`RigidBodyPhysicsCfg` for mass and friction. ``cube_a`` is placed slightly higher than ``cube_b`` so that, once detached, the lower cube lands first and the relative pose visibly changes. diff --git a/docs/source/tutorial/robot.rst b/docs/source/tutorial/robot.rst index cd3f277ac..5875f9b3a 100644 --- a/docs/source/tutorial/robot.rst +++ b/docs/source/tutorial/robot.rst @@ -63,7 +63,7 @@ Drive properties control how the robot's joints behave during simulation, includ .. literalinclude:: ../../../scripts/tutorials/sim/create_robot.py :language: python - :start-at: drive_pros=JointDrivePropertiesCfg( + :start-at: joint_drive_props=JointDrivePropertiesCfg( :end-at: ) You can set different stiffness values for different joint groups using regex patterns. More details on drive properties can be found in :class:`cfg.JointDrivePropertiesCfg`. diff --git a/docs/source/tutorial/semantic_skills.rst b/docs/source/tutorial/semantic_skills.rst index e4d93d0d7..f3779671c 100644 --- a/docs/source/tutorial/semantic_skills.rst +++ b/docs/source/tutorial/semantic_skills.rst @@ -122,7 +122,7 @@ no robot control-part names: object=workpiece, at=SemanticPose( position=(-0.40, 0.48, 0.025), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ), ), ) 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 000000000..06856d256 --- /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/plans/2026-07-13-newton-runtime-contracts-stage-1.md b/docs/superpowers/plans/2026-07-13-newton-runtime-contracts-stage-1.md new file mode 100644 index 000000000..049f6cc96 --- /dev/null +++ b/docs/superpowers/plans/2026-07-13-newton-runtime-contracts-stage-1.md @@ -0,0 +1,2617 @@ +# Newton Runtime Contracts Stage 1 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:** Replace the current private, stale-ID Newton integration with a +public DexSim contract that supports exact prepare/step semantics, +generation-aware rigid/articulation bindings, transactional runtime rebuild, +full multi-arena frames, and independent simultaneous simulation managers. + +**Architecture:** DexSim remains the sole owner of Newton model/state/control +truth and publishes stable entity references, immutable generation-tagged +bindings, prepare results, and rebuild events. EmbodiChain owns environment +semantics through an explicit `BackendSceneContext`, rebinds views after a +generation change, initializes only newly added objects, and keeps its existing +public object and manager calls as compatibility wrappers. + +**Tech Stack:** Python 3.11+, DexSim, NVIDIA Newton, NVIDIA Warp, PyTorch, +Gymnasium, pytest, Sphinx Markdown, Black 26.3.1. + +## Global Constraints + +- Source of truth: `docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md`. +- EmbodiChain branch: `feature/newton-physics-backend` at or after `1d4b9eb`. +- DexSim branch: `feature/embodichain-newton-contracts` from `dev@5281bce`. +- DexSim package version is exactly `0.4.4`; `NEWTON_INTEGRATION_API_VERSION` is exactly `2`. +- `prepare()` creates a ready runtime and never advances simulation time; a requested update performs exactly the requested number of physics steps after prepare. +- Model generation starts at `0`, first successful finalize commits `1`, and only successful model replacement increments it. +- Runtime topology mutation is supported for rigid bodies and articulations; soft-body and cloth mutation fails explicitly. +- Runtime IDs are generation-scoped. Stable public identity is `(world_token, entity_handle, entity_kind)`. +- Public poses use `float32`, `xyzw`, and explicit world/arena-local frames; arena conversion uses the full SE(3) transform. +- Existing `SimulationManager`, `RigidObject`, and `Articulation` public call surfaces remain source compatible. +- Core Newton paths must not use `dexsim.default_world()`, global `get_physics_scene()`, or the default `SimulationManager` instance. +- Default-backend behavior remains unchanged. +- Stage 2 differentiable execution is not implemented by this plan; its plan is written only after this stage's merge gate passes. + +--- + +## Repository Baselines and Reference Files + +Run before Task 1: + +```bash +git -C /root/sources/dexsim branch --show-current +git -C /root/sources/dexsim rev-parse HEAD +git -C /root/sources/EmbodiChain branch --show-current +git -C /root/sources/EmbodiChain rev-parse HEAD +``` + +Expected branch names are `feature/embodichain-newton-contracts` and +`feature/newton-physics-backend`. Record the actual starting SHAs in the +execution log; do not reset either repository to the SHAs above. + +Use these implementations as focused references, without copying their +singleton/Omniverse assumptions: + +- DexSim runtime: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- DexSim rebuild: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rebuild.py` +- DexSim articulation spans: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/articulation.py` +- IsaacLab clone mapping: `/root/sources/IsaacLab/source/isaaclab_newton/isaaclab_newton/cloner/newton_replicate.py` +- IsaacLab rebinding: `/root/sources/IsaacLab/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation_data.py` +- IsaacLab FK invalidation: `/root/sources/IsaacLab/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py` + +## File Map + +### DexSim files created + +- `python/dexsim/engine/newton_physics/contracts.py` — API version, stable references, prepare/rebuild results, subscriptions, model leases, and public errors. +- `python/dexsim/engine/newton_physics/bindings.py` — immutable rigid/articulation bindings and explicit q/qd spans. +- `python/dexsim/engine/newton_physics/runtime_snapshot.py` — generation-independent rigid/articulation snapshot and restore data. +- `python/test/engine/newton_physics/newton_contract_test_utils.py` — shared world, rigid, articulation, and state-array test helpers. +- `python/test/engine/newton_physics/test_newton_public_contract.py` — public API/version/prepare/attachment tests. +- `python/test/engine/newton_physics/test_newton_bindings.py` — generation, cross-world, static-body, and joint-span tests. +- `python/test/engine/newton_physics/test_newton_transactional_rebuild.py` — rigid/articulation preservation and rollback tests. +- `python/test/engine/newton_physics/test_newton_multi_world_runtime.py` — same-device isolation and cleanup tests. + +### DexSim files modified + +- `version.txt` — package patch version `0.4.4`. +- `python/dexsim/engine/newton_physics/__init__.py` — public exports. +- `python/dexsim/engine/newton_physics/newton_manager.py` — generation, prepare, public attachment/binding delegation, candidate commit, close, events. +- `python/dexsim/engine/newton_physics/rebuild.py` — candidate build/restore/validate/atomic commit. +- `python/dexsim/engine/newton_physics/registry.py` — world-token ownership and public owner lookup. +- `python/dexsim/engine/newton_physics/rigid_body/add_body.py` — legacy patch delegates to public attachment. +- `python/dexsim/engine/newton_physics/rigid_body/registration.py` — canonical descriptor replay into a chosen build target. +- `python/dexsim/engine/newton_physics/articulation/articulation.py` — stable ref and explicit q/qd span export. +- `python/dexsim/engine/newton_physics/articulation/skeleton_bridge.py` — canonical articulation replay and removal delta. +- `python/dexsim/engine/newton_physics/world.py` — prepare-then-step behavior and closed-world checks. +- `python/dexsim/engine/newton_physics/integration.py` — deterministic per-world teardown. +- `python/dexsim/engine/newton_physics/capture_coordinator.py` — weak, device-scoped capture coordination. + +### EmbodiChain files created + +- `embodichain/lab/sim/physics/context.py` — explicit backend/world/scene/arena ownership and full transform table. +- `tests/sim/newton_contract_test_utils.py` — deterministic Newton manager and asset-config fixtures used by the Stage 1 tests. +- `tests/sim/test_newton_scene_context.py` — ownership and full-frame conversion tests. +- `tests/sim/test_newton_rebuild_bindings.py` — pending initialization and generation refresh tests. +- `tests/sim/test_newton_multi_manager.py` — two-manager isolation and cleanup tests. + +### EmbodiChain files modified + +- `pyproject.toml` — exact `dexsim_engine==0.4.4` dependency. +- `embodichain/lab/sim/cfg.py` — strict Newton configuration validation. +- `embodichain/lab/sim/common.py` — remove base-constructor virtual reset. +- `embodichain/lab/sim/physics/__init__.py` — export the context/capability types. +- `embodichain/lab/sim/physics/base.py` — structured capability and prepare-result contracts. +- `embodichain/lab/sim/physics/default.py` — default-backend compatibility implementation. +- `embodichain/lab/sim/physics/newton.py` — API handshake, event subscription, generation, pending initialization, close. +- `embodichain/lab/sim/sim_manager.py` — context construction, exact update lifecycle, remove invalidation, idempotent close/reset. +- `embodichain/lab/sim/utility/sim_utils.py` — public `attach_rigid_body` use; remove private registry/meta writes. +- `embodichain/lab/sim/objects/rigid_object.py` — explicit context and deferred initialization. +- `embodichain/lab/sim/objects/articulation.py` — explicit context, separate position/velocity widths, deferred initialization. +- `embodichain/lab/sim/objects/robot.py` — pass owner context and defer reset until fully constructed. +- `embodichain/lab/sim/objects/backends/newton.py` — generation-aware bindings, full frames, FK, explicit unsupported errors. +- `embodichain/lab/sim/objects/backends/default.py` — accept explicit context without behavior changes. +- `tests/sim/test_backend_parity.py` — structured capability matrix. +- `tests/sim/test_newton_finalize_lifecycle.py` — exact prepare and initialization semantics. +- `tests/sim/objects/test_rigid_object.py` — multi-arena/rebuild compatibility cases. +- `tests/sim/objects/test_articulation.py` — q/qd span, FK, link/root frame, rebuild cases. +- `docs/source/overview/sim/sim_manager.md` — public lifecycle, capabilities, topology mutation, and cleanup contract. +- `design/newton-backend-design.md` — replace obsolete Target 4 claims with verified Stage 1 status. + +--- + +### Task 1: Publish the DexSim integration contract and generation lifecycle + +**Files:** + +- Create: `/root/sources/dexsim/python/dexsim/engine/newton_physics/contracts.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/newton_contract_test_utils.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_public_contract.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/__init__.py` +- Modify: `/root/sources/dexsim/version.txt` + +**Interfaces:** + +- Produces: `NEWTON_INTEGRATION_API_VERSION = 2`. +- Produces: `NewtonEntityRef`, `NewtonPrepareResult`, `NewtonPrepareFailure`, `NewtonModelRebuiltEvent`, `NewtonRuntimeStatus`, `NewtonSubscription`, `NewtonModelLease`. +- Produces: `NewtonIntegrationError` subclasses for stale generation, cross-world use, unsupported operation, closed runtime, rebuild failure, and active model lease. +- Produces: `NewtonManager.world_token`, `NewtonManager.model_generation`, `NewtonManager.runtime_status`, and `NewtonManager.prepare()`. +- Produces: `NewtonManager.acquire_model_lease() -> NewtonModelLease`; Stage 2 sessions consume this primitive without changing its rebuild-safety semantics. + +- [ ] **Step 1: Add shared deterministic test helpers** + +Create `newton_contract_test_utils.py` and import its symbols explicitly from +each new DexSim test file: + +```python +DT = 0.01 + + +def make_world(device: str = "cpu"): + config = dexsim.WorldConfig() + config.open_windows = False + config.use_default_physics = False + config.backend = dexsim.types.Backend.OPENGL + config.renderer = dexsim.types.Renderer.HYBRID + world = dexsim.World(config) + world.set_physics_backend("Newton", cfg=NewtonCfg(device=device)) + manager = get_newton_manager(world) + manager._dexsim_renderer = config.renderer + manager._visualizer_disabled = True + return world, world.get_env(), manager + + +@pytest.fixture +def newton_world(): + world, env, manager = make_world() + try: + yield world, env + finally: + world.quit() + + +@pytest.fixture +def two_newton_worlds(): + first = make_world() + second = make_world() + try: + yield first, second + finally: + first[0].quit() + second[0].quit() + + +@pytest.fixture +def two_cuda_worlds(): + if not wp.is_cuda_available(): + pytest.skip("CUDA is required for same-device capture coordination.") + first = make_world("cuda:0") + second = make_world("cuda:0") + try: + yield first, second + finally: + first[0].quit() + second[0].quit() + + +def dynamic_box(arena, name: str, z: float = 1.0): + obj = arena.create_cube(0.1, 0.1, 0.1) + obj.set_name(name) + obj.set_location(0.0, 0.0, z) + attr = PhysicalAttr() + attr.mass = 1.0 + obj.add_rigidbody(ActorType.DYNAMIC, RigidBodyShape.BOX, attr) + return obj + + +def static_plane(arena, name: str): + obj = arena.create_plane(0.0, 10.0) + obj.set_name(name) + obj.add_rigidbody(ActorType.STATIC, RigidBodyShape.PLANE, PhysicalAttr()) + return obj + + +def test_articulation(arena, name: str): + path = get_resources_data_path("Robot", "UR5GPI", "UR5_pgi.urdf") + articulation = arena.load_urdf(path) + articulation.set_name(name) + return articulation + + +def assign_body_state(state, body_id: int, pose, velocity, acceleration) -> None: + body_q = state.body_q.numpy() + body_qd = state.body_qd.numpy() + body_qdd = state.body_qdd.numpy() + body_q[body_id] = np.asarray(pose, dtype=np.float32) + body_qd[body_id] = np.asarray(velocity, dtype=np.float32) + body_qdd[body_id] = np.asarray(acceleration, dtype=np.float32) + state.body_q.assign(body_q) + state.body_qd.assign(body_qd) + state.body_qdd.assign(body_qdd) +``` + +Each test file imports `dynamic_box as _dynamic_box`, +`static_plane as _static_plane`, `test_articulation as +_test_urdf_articulation`, and `assign_body_state as _assign_body_state`, so all +helper names in the following snippets are defined. + +- [ ] **Step 2: Write contract and prepare tests that fail on the current API** + +Add tests with these assertions: + +```python +def test_public_contract_version_and_initial_generation(newton_world): + world, _ = newton_world + mgr = get_newton_manager(world) + assert Version(dexsim.__version__).base_version == "0.4.4" + assert NEWTON_INTEGRATION_API_VERSION == 2 + assert mgr.model_generation == 0 + assert mgr.runtime_status.model_finalized is False + + +def test_prepare_finalizes_without_advancing_time(newton_world): + world, env = newton_world + _dynamic_box(env, "box") + mgr = get_newton_manager(world) + before = mgr._sim_time + result = mgr.prepare() + assert result.generation == 1 + assert result.did_build is True + assert result.did_rebuild is False + assert len(result.added_entities) == 1 + assert result.removed_entities == () + assert mgr._sim_time == before + assert mgr.model_generation == 1 + assert mgr.runtime_status.solver_ready is True + + +def test_second_prepare_is_idempotent(newton_world): + world, env = newton_world + _dynamic_box(env, "box") + mgr = get_newton_manager(world) + first = mgr.prepare() + second = mgr.prepare() + assert first.generation == second.generation == 1 + assert second.did_build is False + assert second.did_rebuild is False +``` + +- [ ] **Step 3: Run the focused test and confirm the missing-contract failure** + +Run: + +```bash +cd /root/sources/dexsim +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py +``` + +Expected: collection fails because the new public symbols and `prepare()` do +not exist. + +- [ ] **Step 4: Add the public value types and errors** + +Implement the contract with immutable, typed values: + +```python +from __future__ import annotations + + +NEWTON_INTEGRATION_API_VERSION = 2 + + +class NewtonIntegrationError(RuntimeError): + """Base error for the public Newton integration contract.""" + + +class NewtonStaleBindingError(NewtonIntegrationError): + pass + + +class NewtonCrossWorldError(NewtonIntegrationError): + pass + + +class NewtonUnsupportedOperationError(NewtonIntegrationError): + pass + + +class NewtonClosedError(NewtonIntegrationError): + pass + + +class NewtonRebuildError(NewtonIntegrationError): + def __init__(self, failure: NewtonPrepareFailure) -> None: + self.failure = failure + super().__init__(failure.message) + + +class NewtonActiveLeaseError(NewtonIntegrationError): + pass + + +@dataclass(frozen=True, slots=True) +class NewtonEntityRef: + world_token: int + entity_handle: int + entity_kind: Literal["rigid", "articulation"] + + +@dataclass(frozen=True, slots=True) +class NewtonPrepareResult: + generation: int + did_build: bool + did_rebuild: bool + added_entities: tuple[NewtonEntityRef, ...] = () + removed_entities: tuple[NewtonEntityRef, ...] = () + + +@dataclass(frozen=True, slots=True) +class NewtonModelRebuiltEvent: + old_generation: int + new_generation: int + added_entities: tuple[NewtonEntityRef, ...] + removed_entities: tuple[NewtonEntityRef, ...] + + +@dataclass(frozen=True, slots=True) +class NewtonRuntimeStatus: + model_finalized: bool + solver_ready: bool + running: bool + stale: bool + closed: bool + + +@dataclass(frozen=True, slots=True) +class NewtonPrepareFailure: + generation: int + operation: Literal["build", "rebuild"] + message: str + + +class NewtonSubscription: + def __init__(self, unsubscribe: Callable[[], None]) -> None: + self._unsubscribe = unsubscribe + self._closed = False + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._unsubscribe() + + +class NewtonModelLease: + def __init__( + self, + *, + world_token: int, + generation: int, + model: object, + release: Callable[[], None], + ) -> None: + self.world_token = world_token + self.generation = generation + self.model = model + self._release = release + self._closed = False + + def __enter__(self) -> NewtonModelLease: + return self + + def __exit__(self, *exc_info: object) -> None: + self.close() + + def close(self) -> None: + if self._closed: + return + self._closed = True + self._release() +``` + +Keep `NewtonModelState.BUILDER` and `NewtonModelState.READY` working for +existing callers. Add `_model_generation`, `_has_stepped`, `_closed`, and a +monotonic per-manager `world_token`. `prepare()` delegates build/rebuild to the +later transactional helper, returns an idempotent result, and never calls +`simulate()` or increments `_sim_time`. +Add +`subscribe_model_rebuilt(callback: Callable[[NewtonModelRebuiltEvent], None]) -> NewtonSubscription`; +callbacks are stored per manager and invoked only after an atomic successful +commit. A failed build raises `NewtonRebuildError` carrying a +`NewtonPrepareFailure` and emits no rebuilt event. +`acquire_model_lease()` requires a finalized, non-stale model, increments a +per-manager active-lease counter, and returns a lease that strongly owns that +exact model and generation. Lease release is idempotent and decrements the +counter exactly once. This is a lifecycle primitive only; no differentiable +session or tape behavior is added in Stage 1. + +- [ ] **Step 5: Export the contract and bump the package patch version** + +Export the new symbols from `newton_physics/__init__.py` and change only: + +```text +DEXSIM_VERSION_MAJOR 0 +DEXSIM_VERSION_MINOR 4 +DEXSIM_VERSION_PATCH 4 +``` + +- [ ] **Step 6: Refresh the editable DexSim package metadata** + +The development install points at `build_Release/lib/python_package`; its +Python modules are source symlinks, but its `dexsim/version.txt` is a generated +copy. Refresh it through the repository-supported setup script after changing +the root version: + +```bash +cd /root/sources/dexsim +PYTHON_BIN="$(command -v python)" ./setup_dev_python.sh -j12 +python - <<'PY' +from packaging.version import Version +import dexsim +assert Version(dexsim.__version__).base_version == "0.4.4" +print(dexsim.__version__, dexsim.__file__) +PY +``` + +Expected: the editable install reports base version `0.4.4` and imports from +the local development package. Generated `build_Release` contents are never +staged or committed. + +- [ ] **Step 7: Run the focused tests** + +Run: + +```bash +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py +``` + +Expected: all tests in the file pass; no simulation-time change occurs during +prepare. + +- [ ] **Step 8: Commit the public contract** + +```bash +git -C /root/sources/dexsim add version.txt python/dexsim/engine/newton_physics/contracts.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/__init__.py python/test/engine/newton_physics/newton_contract_test_utils.py python/test/engine/newton_physics/test_newton_public_contract.py +git -C /root/sources/dexsim commit -m "feat(newton): publish runtime integration contract" +``` + +--- + +### Task 2: Replace private rigid registration with stable public attachment + +**Files:** + +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/contracts.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rigid_body/add_body.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rigid_body/registration.py` +- Modify: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_public_contract.py` + +**Interfaces:** + +- Produces: `NewtonManager.attach_rigid_body(...) -> NewtonEntityRef`. +- Produces: immutable `NewtonRigidDescriptor` replay records owned by DexSim. +- Produces: `NewtonManager.entity_ref(entity)`, `descriptor_for(ref)`, and read-only `descriptors`. +- Consumes: `NewtonEntityRef` and generation lifecycle from Task 1. + +- [ ] **Step 1: Add failing public attachment, clone-world, and descriptor-isolation tests** + +```python +def test_attach_rigid_body_returns_stable_ref(newton_world): + world, env = newton_world + cube = env.create_cube(0.1, 0.2, 0.3) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + cube, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.BOX, + physical_attr=PhysicalAttr(), + ) + assert ref.world_token == mgr.world_token + assert ref.entity_handle == cube.get_native_handle() + assert ref.entity_kind == "rigid" + assert mgr.entity_ref(cube) == ref + + +def test_child_arena_attachment_uses_child_world(newton_world): + world, env = newton_world + arena = env.add_arena("arena_a") + cube = arena.create_cube(0.1, 0.1, 0.1) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + cube, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.BOX, + physical_attr=PhysicalAttr(), + ) + mgr.prepare() + binding = mgr.bind_rigid_entities((ref,)) + assert binding.world_ids_host.tolist() == [0] + assert mgr._model.body_world.numpy()[binding.body_ids_host[0]] == 0 + + +def test_clone_descriptors_do_not_alias(newton_world): + world, env = newton_world + arena_a = env.add_arena("arena_a") + arena_b = env.add_arena("arena_b") + prototype = arena_a.create_sphere(0.1) + prototype.set_name("prototype") + prototype.add_rigidbody( + ActorType.DYNAMIC, RigidBodyShape.SPHERE, PhysicalAttr() + ) + clone = arena_a.clone_actor_to( + "prototype", arena_b, "clone", ObjectCloneOptions() + ) + mgr = get_newton_manager(world) + source = mgr.descriptor_for(mgr.entity_ref(prototype)) + target = mgr.descriptor_for(mgr.entity_ref(clone)) + assert source is not target + assert source.world_id == 0 + assert target.world_id == 1 + + +def test_mesh_geometry_descriptor_is_owned_by_dexsim(newton_world): + world, env = newton_world + entity = env.create_cube(0.1, 0.1, 0.1) + vertices = np.array( + [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32 + ) + triangles = np.array( + [[0, 1, 2], [0, 1, 3], [0, 2, 3], [1, 2, 3]], dtype=np.int32 + ) + geometry = GeometryDesc.mesh(vertices=vertices, triangles=triangles) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + entity, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.MESH, + physical_attr=PhysicalAttr(), + geometry_desc=geometry, + ) + vertices[0] = 99.0 + owned = mgr.descriptor_for(ref).geometry_desc + assert np.allclose(owned.vertices[0], [0.0, 0.0, 0.0]) + + +def test_desc_native_box_prepares_without_shape_parameter_fallback(newton_world): + world, env = newton_world + entity = env.create_cube(0.1, 0.2, 0.3) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + entity, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.BOX, + body_desc=RigidBodyPhysicsDesc.dynamic(mass=1.0), + shape_desc=NewtonCollisionDesc(ke=1000.0, kd=50.0, margin=0.01), + geometry_desc=GeometryDesc.cube((0.1, 0.2, 0.3)), + ) + result = mgr.prepare() + assert result.generation == 1 + assert mgr.bind_rigid_entities((ref,)).body_ids_host[0] >= 0 + + +def test_desc_native_sphere_prepares_from_owned_geometry(newton_world): + world, env = newton_world + entity = env.create_sphere(0.2) + mgr = get_newton_manager(world) + ref = mgr.attach_rigid_body( + entity, + actor_type=ActorType.DYNAMIC, + shape_type=RigidBodyShape.SPHERE, + body_desc=RigidBodyPhysicsDesc.dynamic(mass=1.0), + shape_desc=NewtonCollisionDesc(ke=1000.0, kd=50.0), + geometry_desc=GeometryDesc.sphere(0.2), + ) + mgr.prepare() + descriptor = mgr.descriptor_for(ref) + assert descriptor.geometry_desc.radius == pytest.approx(0.2) + assert mgr.bind_rigid_entities((ref,)).body_ids_host[0] >= 0 +``` + +- [ ] **Step 2: Run the tests and confirm public attachment is absent** + +Run the attachment tests above with `pytest -q`. Expected: failures report +missing `attach_rigid_body`, `entity_ref`, or `descriptor_for`. + +- [ ] **Step 3: Add canonical descriptor ownership and the public method** + +Add an immutable descriptor that deep-copies mutable desc-native inputs: + +```python +@dataclass(frozen=True, slots=True) +class NewtonRigidDescriptor: + entity_ref: NewtonEntityRef + arena_handle: int + world_id: int + actor_type: ActorType + shape_type: RigidBodyShape + node_scale: tuple[float, float, float] + body_scale: tuple[float, float, float] + physical_attr: PhysicalAttr | None + body_desc: object | None + shape_desc: object | None + geometry_desc: object | None +``` + +The `object` annotations above stand for the existing concrete DexSim spawn +descriptor types, not borrowed arbitrary objects. At attachment time, normalize +legacy `PhysicalAttr` into owned body/collision values and recursively copy all +desc-native data. Copy NumPy arrays with canonical `float32`/`int32` dtypes and +mark them read-only. Resolve file-backed mesh data, convex/ACD hulls, and SDF +mesh/config inputs into descriptor-owned replay data so a rebuild neither reads +mutable entity metadata nor depends on a later file change. `descriptor_for()` +returns this immutable snapshot; it never returns `dexsim_meta`. + +Implement the public method with this exact surface: + +```python +def attach_rigid_body( + self, + entity, + *, + actor_type, + shape_type, + physical_attr=None, + body_desc=None, + shape_desc=None, + geometry_desc=None, +) -> NewtonEntityRef: + self._assert_open() + ref = self._ref_for_entity(entity, "rigid") + descriptor = make_rigid_descriptor( + manager=self, + entity=entity, + entity_ref=ref, + actor_type=actor_type, + shape_type=shape_type, + physical_attr=physical_attr, + body_desc=body_desc, + shape_desc=shape_desc, + geometry_desc=geometry_desc, + ) + self._rigid_descriptors[ref] = descriptor + replay_rigid_descriptor(self, descriptor, entity) + self._record_added(ref) + self.mark_runtime_model_stale() + return ref +``` + +The legacy `MeshObject.add_rigidbody` patches call this method and return its +reference. `register_mesh_object_to_newton_patch` remains temporarily importable +for DexSim compatibility but delegates through descriptor replay and is no +longer called by EmbodiChain. +Legacy `add_sdf_rigidbody` and `add_acd_rigidbody` project their SDF config or +resolved convex hulls into the same `geometry_desc` contract before delegation; +the existing SDF/ACD regression tests must keep passing. + +- [ ] **Step 4: Run attachment and existing scene mutation tests** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py python/test/engine/newton_physics/test_newton_scene_mutations.py +``` + +Expected: both files pass, including global world `-1` and child world IDs. + +- [ ] **Step 5: Commit public rigid attachment** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/contracts.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/rigid_body/add_body.py python/dexsim/engine/newton_physics/rigid_body/registration.py python/test/engine/newton_physics/test_newton_public_contract.py +git -C /root/sources/dexsim commit -m "feat(newton): add stable rigid attachment API" +``` + +--- + +### Task 3: Add immutable generation-aware rigid and articulation bindings + +**Files:** + +- Create: `/root/sources/dexsim/python/dexsim/engine/newton_physics/bindings.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_bindings.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/articulation.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/__init__.py` + +**Interfaces:** + +- Produces: `NewtonJointSpan`, `NewtonArticulationSpan`, `RigidEntityBinding`, `ArticulationBinding`. +- Produces: `joint_span_from_builder(builder, joint_id) -> NewtonJointSpan`. +- Produces: `NewtonManager.bind_rigid_entities(refs, device=None)`. +- Produces: `NewtonManager.bind_articulations(refs, device=None)`. +- Consumes: stable references from Task 2. + +- [ ] **Step 1: Add failing binding contract tests** + +```python +def test_rigid_binding_is_int32_and_generation_scoped(newton_world): + world, env = newton_world + dynamic = _dynamic_box(env, "dynamic") + static = _static_plane(env, "static") + mgr = get_newton_manager(world) + result = mgr.prepare() + binding = mgr.bind_rigid_entities( + (mgr.entity_ref(dynamic), mgr.entity_ref(static)) + ) + assert binding.generation == result.generation + assert binding.body_ids_host.dtype == np.int32 + assert binding.shape_ids_host.dtype == np.int32 + assert binding.body_ids_host[1] == -1 + binding.assert_current(mgr) + + +def test_binding_rejects_other_world(two_newton_worlds): + (world_a, env_a, mgr_a), (_, _, mgr_b) = two_newton_worlds + box = _dynamic_box(env_a, "box") + mgr_a.prepare() + with pytest.raises(NewtonCrossWorldError): + mgr_b.bind_rigid_entities((mgr_a.entity_ref(box),)) + + +def test_articulation_binding_exposes_distinct_q_and_qd_spans(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + binding = mgr.bind_articulations((mgr.entity_ref(art),)) + assert binding.qpos_width > 0 + assert binding.qvel_width > 0 + assert all( + span.q_width >= 0 for spans in binding.joint_spans for span in spans + ) + assert all( + span.qd_width >= 0 for spans in binding.joint_spans for span in spans + ) + assert binding.joint_spans_wp.dtype == wp.int32 +``` + +Add builder-only unit cases for revolute, spherical, and free-joint widths: + +```python +@pytest.mark.parametrize( + "q_width, qd_width", + [(1, 1), (4, 3), (7, 6)], +) +def test_joint_span_uses_distinct_position_and_velocity_widths( + q_width, qd_width +): + builder = SimpleNamespace( + joint_q_start=[0, q_width], + joint_qd_start=[0, qd_width], + joint_q=[0.0] * q_width, + joint_qd=[0.0] * qd_width, + joint_target_pos=[0.0] * qd_width, + joint_target_vel=[0.0] * qd_width, + ) + span = joint_span_from_builder(builder, joint_id=0) + assert span.q_start == span.qd_start == 0 + assert span.q_width == q_width + assert span.qd_width == qd_width + assert span.target_q_start == span.qd_start + assert span.target_q_width == qd_width + assert span.target_qd_start == span.qd_start + assert span.target_qd_width == qd_width +``` + +- [ ] **Step 2: Run the binding file and confirm collection/API failures** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_bindings.py +``` + +Expected: missing binding classes/methods. + +- [ ] **Step 3: Implement immutable bindings with O(1) validation** + +```python +@dataclass(frozen=True, slots=True) +class NewtonJointSpan: + joint_id: int + q_start: int + q_width: int + qd_start: int + qd_width: int + target_q_start: int + target_q_width: int + target_qd_start: int + target_qd_width: int + + +@dataclass(frozen=True, slots=True) +class NewtonArticulationSpan: + ref: NewtonEntityRef + q_start: int + q_width: int + qd_start: int + qd_width: int + target_q_start: int + target_q_width: int + target_qd_start: int + target_qd_width: int + control_start: int + control_width: int + + +@dataclass(frozen=True, slots=True) +class RigidEntityBinding: + world_token: int + generation: int + refs: tuple[NewtonEntityRef, ...] + body_ids_host: np.ndarray + shape_ids_host: np.ndarray + world_ids_host: np.ndarray + body_ids_wp: wp.array + shape_ids_wp: wp.array + world_ids_wp: wp.array + + def assert_current(self, manager: NewtonManager) -> None: + assert_binding_owner_and_generation(self, manager) + + +@dataclass(frozen=True, slots=True) +class ArticulationBinding: + world_token: int + generation: int + refs: tuple[NewtonEntityRef, ...] + articulation_ids_host: np.ndarray + root_body_ids_host: np.ndarray + link_body_ids_host: np.ndarray + world_ids_host: np.ndarray + articulation_spans: tuple[NewtonArticulationSpan, ...] + joint_spans: tuple[tuple[NewtonJointSpan, ...], ...] + joint_spans_wp: wp.array + qpos_width: int + qvel_width: int + target_qpos_width: int + target_qvel_width: int +``` + +Resolve all host arrays once, upload `int32` device arrays once, make NumPy +arrays read-only, and validate a binding by comparing only `world_token` and +`generation`. Do not resolve entity IDs in steady-state fetch/apply calls. +`joint_span_from_builder` computes each end from the next start entry, or from +the corresponding flat array length for the final joint. Current position uses +`joint_q_start/q_width`; current velocity uses `joint_qd_start/qd_width`. +Newton `joint_target_pos`, `joint_target_vel`, and `joint_f` are per-DOF, so +their starts and widths use the qd span even for spherical/free joints. Do not +reuse the coordinate-width q span for position targets. + +- [ ] **Step 4: Run binding and simulation-index regressions** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_bindings.py python/test/engine/newton_physics/test_newton_sim_index.py +``` + +Expected: all tests pass; old `get_sim_index()` remains a generation-local +compatibility view, not a stable handle. + +- [ ] **Step 5: Commit bindings** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/bindings.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/articulation/articulation.py python/dexsim/engine/newton_physics/__init__.py python/test/engine/newton_physics/test_newton_bindings.py +git -C /root/sources/dexsim commit -m "feat(newton): add generation-aware entity bindings" +``` + +--- + +### Task 4: Make rigid rebuild transactional and preserve both state buffers + +**Files:** + +- Create: `/root/sources/dexsim/python/dexsim/engine/newton_physics/runtime_snapshot.py` +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_transactional_rebuild.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rebuild.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/world.py` + +**Interfaces:** + +- Produces: `NewtonRuntimeSnapshot`, `RigidRuntimeSnapshot`, and candidate-runtime commit. +- Produces: successful rebuild event after commit only. +- Produces: rebuild rejection while a model-generation lease is active. +- Consumes: descriptors and bindings from Tasks 2–3. + +- [ ] **Step 1: Add failing state-preservation, rollback, and exact-step tests** + +```python +def test_rigid_rebuild_preserves_both_states_and_wrench(newton_world): + world, env = newton_world + box = _dynamic_box(env, "box") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(box) + body_id = mgr.bind_rigid_entities((ref,)).body_ids_host[0] + pose0 = np.array([0.1, 0.2, 1.3, 0.0, 0.0, 0.0, 1.0], np.float32) + pose1 = np.array([0.2, 0.3, 1.4, 0.0, 0.0, 0.0, 1.0], np.float32) + _assign_body_state( + mgr._state_0, + body_id, + pose0, + [1, 2, 3, 4, 5, 6], + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + ) + _assign_body_state( + mgr._state_1, + body_id, + pose1, + [6, 5, 4, 3, 2, 1], + [0.6, 0.5, 0.4, 0.3, 0.2, 0.1], + ) + external_forces = mgr._external_forces.numpy() + external_forces[body_id] = [1, 2, 3, 4, 5, 6] + mgr._external_forces.assign(external_forces) + _dynamic_box(env, "new_box") + result = mgr.prepare() + new_id = mgr.bind_rigid_entities((ref,)).body_ids_host[0] + assert result.did_rebuild is True + assert np.allclose(mgr._state_0.body_q.numpy()[new_id], pose0) + assert np.allclose(mgr._state_1.body_q.numpy()[new_id], pose1) + assert np.allclose( + mgr._state_0.body_qdd.numpy()[new_id], + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6], + ) + assert np.allclose( + mgr._state_1.body_qdd.numpy()[new_id], + [0.6, 0.5, 0.4, 0.3, 0.2, 0.1], + ) + assert np.allclose(mgr._external_forces.numpy()[new_id], [1, 2, 3, 4, 5, 6]) + + +def test_failed_candidate_keeps_old_runtime_and_generation(newton_world, monkeypatch): + world, env = newton_world + box = _dynamic_box(env, "box") + mgr = get_newton_manager(world) + mgr.prepare() + old_model = mgr._model + old_generation = mgr.model_generation + _dynamic_box(env, "new_box") + monkeypatch.setattr(rebuild, "_validate_candidate", lambda candidate: (_ for _ in ()).throw(ValueError("injected"))) + with pytest.raises(NewtonRebuildError, match="injected"): + mgr.prepare() + assert mgr._model is old_model + assert mgr.model_generation == old_generation + assert mgr.lifecycle_state is NewtonModelState.STALE + + +def test_world_update_prepares_then_executes_one_step(newton_world): + world, env = newton_world + _dynamic_box(env, "box", z=1.0) + mgr = get_newton_manager(world) + before = mgr._sim_time + world.update(0.01) + assert mgr.model_generation == 1 + assert mgr._sim_time == pytest.approx(before + 0.01) + + +def test_active_model_lease_blocks_rebuild_until_released(newton_world): + world, env = newton_world + _dynamic_box(env, "box") + mgr = get_newton_manager(world) + mgr.prepare() + lease = mgr.acquire_model_lease() + assert lease.world_token == mgr.world_token + assert lease.generation == mgr.model_generation + assert lease.model is mgr._model + _dynamic_box(env, "new_box") + with pytest.raises(NewtonActiveLeaseError, match="generation 1"): + mgr.prepare() + assert mgr.model_generation == 1 + lease.close() + lease.close() + assert mgr.prepare().generation == 2 +``` + +- [ ] **Step 2: Run the new file and confirm current rebuild destroys/aliases state** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_transactional_rebuild.py +``` + +Expected: failures show missing snapshot types, non-transactional clearing, or +the current first-update skip; the active-lease case currently rebuilds or has +no lease surface. + +- [ ] **Step 3: Add complete rigid snapshots keyed by stable reference** + +```python +@dataclass(frozen=True, slots=True) +class RigidRuntimeSnapshot: + ref: NewtonEntityRef + state_0_pose: np.ndarray + state_0_velocity: np.ndarray + state_0_acceleration: np.ndarray + state_1_pose: np.ndarray + state_1_velocity: np.ndarray + state_1_acceleration: np.ndarray + external_wrench: np.ndarray + + +@dataclass(frozen=True, slots=True) +class NewtonRuntimeSnapshot: + generation: int + rigid: dict[NewtonEntityRef, RigidRuntimeSnapshot] + articulations: dict[NewtonEntityRef, ArticulationRuntimeSnapshot] +``` + +Capture only references present in the old finalized model. Restore surviving +references after candidate finalization using the candidate binding. Copy +arrays into both candidate states; do not make both buffers identical when the +old buffers differed. + +- [ ] **Step 4: Build and validate a candidate before mutating the live manager** + +Use an unregistered candidate manager that shares only the stable world token: + +```python +candidate = NewtonManager( + cfg=copy.deepcopy(manager.cfg), + world_token=manager.world_token, + register_live=False, +) +candidate.set_dexsim_world(world) +candidate.replay_descriptors(manager.descriptors) +candidate.start_simulation() +restore_runtime_snapshot(candidate, snapshot) +_validate_candidate(candidate) +manager._commit_candidate(candidate, delta) +``` + +Before constructing a candidate, `prepare()` checks the active-lease count. A +topology-changing prepare raises `NewtonActiveLeaseError` with the world token +and leased generation while the count is non-zero; idempotent prepare of an +unchanged READY model remains allowed. The live model, generation, topology +delta, and STALE status remain unchanged after rejection. Stage 2 must acquire +this lease before recording a tape and close it only after backward or explicit +session detach. + +`_commit_candidate` swaps builder/model/states/control/contacts/pipeline/solver, +entity mappings, descriptors, articulation runtime bindings, caches, and graph +as one non-raising assignment block. It increments generation once, sets READY, +detaches the transferred resources from the candidate, then emits one rebuilt +event. Retire the old runtime only after the swap; cleanup failure is logged +without rolling back a runtime already published to subscribers. Candidate +failure before the swap closes candidate resources, +leaves the live runtime fields unchanged, keeps STALE, and raises +`NewtonRebuildError` chained from the original exception. + +- [ ] **Step 5: Make World.update call prepare and still step** + +Replace skip-on-build behavior with: + +```python +prepare_result = mgr.prepare() +mgr.step(dt) +if mgr.should_sync_to_dexsim(step_override=sync_to_dexsim): + _push_newton_state_to_dexsim(mgr) +``` + +The result is available to integrations but never consumes the requested +physics step. + +- [ ] **Step 6: Run lifecycle and mutation regressions** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_scene_lifecycle.py python/test/engine/newton_physics/test_newton_scene_mutations.py python/test/engine/newton_physics/test_newton_body_dynamics.py +``` + +Expected: all pass; update time increments on the first call and after rebuild. + +- [ ] **Step 7: Commit transactional rigid rebuild** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/runtime_snapshot.py python/dexsim/engine/newton_physics/rebuild.py python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/world.py python/test/engine/newton_physics/test_newton_transactional_rebuild.py +git -C /root/sources/dexsim commit -m "refactor(newton): rebuild rigid runtime transactionally" +``` + +--- + +### Task 5: Preserve articulation runtime state and rebind explicit spans + +**Files:** + +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/runtime_snapshot.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/rebuild.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/articulation.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/articulation/skeleton_bridge.py` +- Modify: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_transactional_rebuild.py` +- Modify: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_bindings.py` + +**Interfaces:** + +- Produces: complete `ArticulationRuntimeSnapshot` and canonical articulation replay. +- Produces: `NewtonManager.forward_kinematics(articulation_mask=None)` as the public FK synchronization entry point; it updates both ping-pong states. +- Consumes: `ArticulationBinding` from Task 3 and candidate transaction from Task 4. + +- [ ] **Step 1: Add failing articulation preservation and removal tests** + +Define these local test helpers above the tests: + +```python +def _assign_slice(owner, name: str, start: int, values: np.ndarray) -> None: + source = getattr(owner, name) + data = source.numpy() + data[start : start + len(values)] = values + source.assign(data) + + +def _read_slice(owner, name: str, start: int, width: int) -> np.ndarray: + return getattr(owner, name).numpy()[start : start + width].copy() + + +def _assign_existing_owners( + manager, owner_names: tuple[str, ...], name: str, start: int, values +) -> None: + assigned = False + for owner_name in owner_names: + owner = getattr(manager, owner_name, None) + if owner is None or getattr(owner, name, None) is None: + continue + _assign_slice(owner, name, start, values) + assigned = True + assert assigned, f"{name} is unavailable on {owner_names}" + + +def _read_first_owner( + manager, owner_names: tuple[str, ...], name: str, start: int, width: int +) -> np.ndarray: + for owner_name in owner_names: + owner = getattr(manager, owner_name, None) + if owner is not None and getattr(owner, name, None) is not None: + return _read_slice(owner, name, start, width) + raise AssertionError(f"{name} is unavailable on {owner_names}") + + +def _write_articulation_arrays( + manager, + binding, + state_0_q, + state_0_qd, + state_1_q, + state_1_qd, + model_target_q, + model_target_qd, + control_target_q, + control_target_qd, + generalized_force, + active_control, +) -> None: + span = binding.articulation_spans[0] + _assign_slice(manager._state_0, "joint_q", span.q_start, state_0_q) + _assign_slice(manager._state_0, "joint_qd", span.qd_start, state_0_qd) + _assign_slice(manager._state_1, "joint_q", span.q_start, state_1_q) + _assign_slice(manager._state_1, "joint_qd", span.qd_start, state_1_qd) + _assign_existing_owners( + manager, + ("_model",), + "joint_target_pos", + span.target_q_start, + model_target_q, + ) + _assign_existing_owners( + manager, + ("_model",), + "joint_target_vel", + span.target_qd_start, + model_target_qd, + ) + _assign_existing_owners( + manager, + ("_control",), + "joint_target_pos", + span.target_q_start, + control_target_q, + ) + _assign_existing_owners( + manager, + ("_control",), + "joint_target_vel", + span.target_qd_start, + control_target_qd, + ) + _assign_existing_owners( + manager, + ("_model",), + "joint_f", + span.control_start, + generalized_force, + ) + _assign_existing_owners( + manager, + ("_control",), + "joint_f", + span.control_start, + active_control, + ) + + +def _read_q(manager, binding, state_name="_state_0"): + span = binding.articulation_spans[0] + return _read_slice( + getattr(manager, state_name), "joint_q", span.q_start, span.q_width + ) + + +def _read_qd(manager, binding, state_name="_state_0"): + span = binding.articulation_spans[0] + return _read_slice( + getattr(manager, state_name), "joint_qd", span.qd_start, span.qd_width + ) + + +def _read_target_q(manager, binding, owner_name): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + (owner_name,), + "joint_target_pos", + span.target_q_start, + span.target_q_width, + ) + + +def _read_target_qd(manager, binding, owner_name): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + (owner_name,), + "joint_target_vel", + span.target_qd_start, + span.target_qd_width, + ) + + +def _read_generalized_force(manager, binding): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + ("_model",), + "joint_f", + span.control_start, + span.control_width, + ) + + +def _read_active_control(manager, binding): + span = binding.articulation_spans[0] + return _read_first_owner( + manager, + ("_control",), + "joint_f", + span.control_start, + span.control_width, + ) +``` + +```python +def test_articulation_rebuild_preserves_current_target_and_control(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(art) + binding = mgr.bind_articulations((ref,)) + q0 = np.linspace(0.01, 0.01 * binding.qpos_width, binding.qpos_width, dtype=np.float32) + qd0 = np.linspace(0.02, 0.02 * binding.qvel_width, binding.qvel_width, dtype=np.float32) + q1 = q0 + 0.4 + qd1 = qd0 + 0.5 + model_target_q = np.linspace( + -0.03, + -0.03 * binding.target_qpos_width, + binding.target_qpos_width, + dtype=np.float32, + ) + model_target_qd = np.linspace( + -0.04, + -0.04 * binding.target_qvel_width, + binding.target_qvel_width, + dtype=np.float32, + ) + control_target_q = model_target_q - 0.7 + control_target_qd = model_target_qd - 0.8 + generalized_force = np.full(binding.qvel_width, 0.3, dtype=np.float32) + active_control = np.full(binding.qvel_width, -0.6, dtype=np.float32) + _write_articulation_arrays( + mgr, + binding, + q0, + qd0, + q1, + qd1, + model_target_q, + model_target_qd, + control_target_q, + control_target_qd, + generalized_force, + active_control, + ) + _dynamic_box(env, "topology_change") + mgr.prepare() + rebound = mgr.bind_articulations((ref,)) + assert np.allclose(_read_q(mgr, rebound), q0) + assert np.allclose(_read_qd(mgr, rebound), qd0) + assert np.allclose(_read_q(mgr, rebound, "_state_1"), q1) + assert np.allclose(_read_qd(mgr, rebound, "_state_1"), qd1) + assert np.allclose( + _read_target_q(mgr, rebound, "_model"), model_target_q + ) + assert np.allclose( + _read_target_qd(mgr, rebound, "_model"), model_target_qd + ) + assert np.allclose( + _read_target_q(mgr, rebound, "_control"), control_target_q + ) + assert np.allclose( + _read_target_qd(mgr, rebound, "_control"), control_target_qd + ) + assert np.allclose( + _read_generalized_force(mgr, rebound), generalized_force + ) + assert np.allclose(_read_active_control(mgr, rebound), active_control) + + +def test_removed_articulation_ref_cannot_rebind(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(art) + env.remove_skeleton("arm") + mgr.prepare() + with pytest.raises(NewtonStaleBindingError, match="removed"): + mgr.bind_articulations((ref,)) + + +def test_articulation_rebuild_preserves_drive_limits_and_feedforward(newton_world): + world, env = newton_world + art = _test_urdf_articulation(env, "arm") + mgr = get_newton_manager(world) + mgr.prepare() + ref = mgr.entity_ref(art) + binding = mgr.bind_articulations((ref,)) + span = binding.articulation_spans[0] + width = span.control_width + expected = { + "joint_target_ke": np.full(width, 11.0, np.float32), + "joint_target_kd": np.full(width, 1.2, np.float32), + "joint_friction": np.full(width, 0.13, np.float32), + "joint_armature": np.full(width, 0.07, np.float32), + "joint_target_mode": np.full(width, 1, np.int32), + "joint_effort_limit": np.full(width, 9.0, np.float32), + "joint_velocity_limit": np.full(width, 4.0, np.float32), + "joint_limit_lower": np.full(width, -0.9, np.float32), + "joint_limit_upper": np.full(width, 0.9, np.float32), + } + for name, values in expected.items(): + _assign_slice(mgr._model, name, span.control_start, values) + feedforward = np.full(width, 0.23, np.float32) + _assign_slice(mgr._control, "joint_act", span.control_start, feedforward) + _dynamic_box(env, "topology_change") + mgr.prepare() + rebound = mgr.bind_articulations((ref,)).articulation_spans[0] + for name, values in expected.items(): + actual = _read_slice( + mgr._model, name, rebound.control_start, rebound.control_width + ) + assert np.allclose(actual, values) + assert np.allclose( + _read_slice( + mgr._control, + "joint_act", + rebound.control_start, + rebound.control_width, + ), + feedforward, + ) +``` + +- [ ] **Step 2: Run the articulation tests and confirm runtime data is lost** + +Expected: current rebuild either omits the articulation or loses current/target +state and control. + +- [ ] **Step 3: Implement complete articulation snapshots** + +```python +@dataclass(frozen=True, slots=True) +class ArticulationRuntimeSnapshot: + ref: NewtonEntityRef + state_0_joint_q: np.ndarray + state_0_joint_qd: np.ndarray + state_1_joint_q: np.ndarray + state_1_joint_qd: np.ndarray + model_target_joint_q: np.ndarray | None + model_target_joint_qd: np.ndarray | None + control_target_joint_q: np.ndarray | None + control_target_joint_qd: np.ndarray | None + model_joint_f: np.ndarray | None + control_joint_f: np.ndarray | None + control_joint_act: np.ndarray | None + drive_stiffness: np.ndarray + drive_damping: np.ndarray + drive_friction: np.ndarray + drive_armature: np.ndarray + drive_target_mode: np.ndarray + drive_effort_limit: np.ndarray + drive_velocity_limit: np.ndarray + joint_limit_lower: np.ndarray + joint_limit_upper: np.ndarray + root_state_0: np.ndarray + root_state_1: np.ndarray +``` + +Read/write each field through explicit spans from `ArticulationBinding`. +Model defaults and active `Control` targets/forces are captured separately; +never collapse them just because the normal setter currently writes both. +An owner field is `None` only when that Newton model/control array is genuinely +absent, and restore preserves that absence. +Capture active feed-forward control from `Control.joint_act`. Capture drive and +limit arrays from `joint_target_ke`, `joint_target_kd`, `joint_friction`, +`joint_armature`, `joint_target_mode`, `joint_effort_limit`, +`joint_velocity_limit`, `joint_limit_lower`, and `joint_limit_upper` on their +live owner (`Control` when exposed, otherwise `Model`) and restore them before +candidate validation. Each `root_state_*` is exactly 13 `float32` values: world pose in +`xyz+xyzw` followed by linear and angular velocity. Contacts are not captured; +the candidate collision pipeline regenerates them. +Canonical articulation replay reconstructs links/joints/drives into the +candidate and then refreshes existing `NewtonArticulation` wrapper metadata at +commit. Never interpret an active-joint ordinal as a flattened q/qd index. + +- [ ] **Step 4: Invalidate FK after current q writes** + +When current q is restored or written, build a boolean articulation mask with +shape `(model.articulation_count,)`. Evaluate FK independently against both +ping-pong states so their distinct q/qd histories remain distinct: + +```python +for state in (manager._state_0, manager._state_1): + eval_fk( + manager._model, + state.joint_q, + state.joint_qd, + state, + articulation_mask, + ) +``` + +Run visual synchronization only after FK is current. + +- [ ] **Step 5: Run articulation, binding, and rebuild tests** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_bindings.py python/test/engine/newton_physics/test_newton_physics_scene.py -k articulation +``` + +Expected: all selected tests pass, including spherical/free-joint span cases. + +- [ ] **Step 6: Commit articulation preservation** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/runtime_snapshot.py python/dexsim/engine/newton_physics/rebuild.py python/dexsim/engine/newton_physics/articulation/articulation.py python/dexsim/engine/newton_physics/articulation/skeleton_bridge.py python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_bindings.py +git -C /root/sources/dexsim commit -m "feat(newton): preserve articulation state across rebuild" +``` + +--- + +### Task 6: Isolate same-device worlds and make DexSim cleanup deterministic + +**Files:** + +- Create: `/root/sources/dexsim/python/test/engine/newton_physics/test_newton_multi_world_runtime.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/newton_manager.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/registry.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/integration.py` +- Modify: `/root/sources/dexsim/python/dexsim/engine/newton_physics/capture_coordinator.py` + +**Interfaces:** + +- Produces: idempotent `NewtonManager.close()` and public `manager_for_entity(entity)`. +- Produces: weak device-level CUDA capture coordination without shared physics state. +- Consumes: per-world tokens and generations from Tasks 1–5. + +- [ ] **Step 1: Add failing two-world isolation and close tests** + +```python +def test_two_worlds_build_step_rebuild_and_close_independently(two_newton_worlds): + (world_a, env_a, mgr_a), (world_b, env_b, mgr_b) = two_newton_worlds + box_a = _dynamic_box(env_a, "a") + box_b = _dynamic_box(env_b, "b") + world_a.update(0.01) + world_b.update(0.02) + assert mgr_a.world_token != mgr_b.world_token + assert mgr_a.model_generation == mgr_b.model_generation == 1 + assert mgr_a._model is not mgr_b._model + _dynamic_box(env_a, "a2") + world_a.update(0.01) + assert mgr_a.model_generation == 2 + assert mgr_b.model_generation == 1 + mgr_a.close() + mgr_a.close() + with pytest.raises(NewtonClosedError): + mgr_a.bind_rigid_entities((mgr_a.entity_ref(box_a),)) + assert manager_for_entity(box_b) is mgr_b + world_b.update(0.02) + + +def test_capture_coordinator_holds_only_weak_manager_refs(two_cuda_worlds): + (_, _, mgr_a), (_, _, mgr_b) = two_cuda_worlds + coordinator = capture_coordinator_for_device(mgr_a.device) + assert coordinator.manager_count == 2 + mgr_a.close() + assert mgr_a not in tuple(coordinator.managers) + assert mgr_b in tuple(coordinator.managers) + assert coordinator.manager_count == 1 +``` + +- [ ] **Step 2: Run the file; confirm leakage/cross-world assumptions fail** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_multi_world_runtime.py +``` + +Expected: failures expose absent close/owner lookup or strong global state. + +- [ ] **Step 3: Implement deterministic per-world teardown** + +`NewtonManager.close()` sets closed once, invalidates graphs/caches, clears +callbacks and subscriptions, releases model/state/control/contact/solver and +renderer resources, unregisters arena/entity ownership, and removes only this +manager from weak coordination. Every public method calls `_assert_open()`. + +Expose owner lookup without leaking private registries: + +```python +def manager_for_entity(entity) -> NewtonManager | None: + arena = entity.get_arena() + return manager_for_arena(arena) +``` + +The CUDA coordinator stores `weakref.WeakSet[NewtonManager]`, serializes only +capture operations for a device, and has a finite diagnostic timeout. It owns +no builder/model/state/control/solver arrays. Expose a read-only `managers` +tuple and derived `manager_count` for diagnostics/tests. + +- [ ] **Step 4: Run the DexSim Stage 1 suite** + +```bash +pytest -q python/test/engine/newton_physics/test_newton_public_contract.py python/test/engine/newton_physics/test_newton_bindings.py python/test/engine/newton_physics/test_newton_transactional_rebuild.py python/test/engine/newton_physics/test_newton_multi_world_runtime.py python/test/engine/newton_physics/test_newton_scene_lifecycle.py python/test/engine/newton_physics/test_newton_scene_mutations.py python/test/engine/newton_physics/test_newton_sim_index.py python/test/engine/newton_physics/test_newton_physics_scene.py +``` + +Expected: zero failures and no surviving per-world registrations after fixture +teardown. + +- [ ] **Step 5: Commit world isolation and cleanup** + +```bash +git -C /root/sources/dexsim add python/dexsim/engine/newton_physics/newton_manager.py python/dexsim/engine/newton_physics/registry.py python/dexsim/engine/newton_physics/integration.py python/dexsim/engine/newton_physics/capture_coordinator.py python/test/engine/newton_physics/test_newton_multi_world_runtime.py +git -C /root/sources/dexsim commit -m "fix(newton): isolate and close per-world runtimes" +``` + +--- + +### Task 7: Add EmbodiChain API handshake, structured capabilities, and scene context + +**Files:** + +- Create: `embodichain/lab/sim/physics/context.py` +- Create: `tests/sim/newton_contract_test_utils.py` +- Create: `tests/sim/test_newton_scene_context.py` +- Modify: `pyproject.toml` +- Modify: `embodichain/lab/sim/cfg.py` +- Modify: `embodichain/lab/sim/physics/base.py` +- Modify: `embodichain/lab/sim/physics/default.py` +- Modify: `embodichain/lab/sim/physics/newton.py` +- Modify: `embodichain/lab/sim/physics/__init__.py` +- Modify: `tests/sim/test_backend_parity.py` + +**Interfaces:** + +- Produces: `PhysicsCapabilities`, `PhysicsPrepareResult`, `BackendSceneContext`. +- Produces: exact DexSim package/API validation during Newton activation. +- Consumes: DexSim public contract from Tasks 1–6. + +- [ ] **Step 1: Add shared EmbodiChain contract test fixtures** + +Create `tests/sim/newton_contract_test_utils.py`; each new Stage 1 test file +imports the fixtures and factories it uses: + +```python +ART_PATH = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" + + +def box_cfg(uid: str, z: float = 1.0) -> RigidObjectCfg: + return RigidObjectCfg.from_dict( + { + "uid": uid, + "shape": {"shape_type": "Cube", "size": [0.1, 0.1, 0.1]}, + "attrs": {"mass": 1.0}, + "body_type": "dynamic", + "init_pos": (0.0, 0.0, z), + } + ) + + +def arm_cfg(uid: str) -> ArticulationCfg: + return ArticulationCfg.from_dict( + { + "uid": uid, + "fpath": get_data_path(ART_PATH), + "drive_pros": {"drive_type": "force"}, + } + ) + + +@pytest.fixture +def newton_sim(): + sim = SimulationManager( + SimulationManagerCfg( + headless=True, + device="cpu", + num_envs=2, + physics_cfg=NewtonPhysicsCfg( + device="cpu", num_substeps=2, use_cuda_graph=False + ), + ) + ) + try: + yield sim + finally: + close = getattr(sim, "close", None) + if close is None: + sim.destroy(exit_process=False) + else: + close() + + +@pytest.fixture +def fake_manager(): + identity = np.eye(4, dtype=np.float32) + rotated = np.eye(4, dtype=np.float32) + rotated[:3, :3] = Rotation.from_euler("z", 90, degrees=True).as_matrix() + rotated[:3, 3] = [2.0, 3.0, 0.0] + + class Root: + def __init__(self, pose): + self.pose = pose + def get_world_pose(self): + return self.pose.copy() + + class Arena: + def __init__(self, pose): + self.root = Root(pose) + def get_root_node(self): + return self.root + + class Manager: + def __init__(self): + self._arenas = [Arena(identity), Arena(rotated)] + self.device = torch.device("cpu") + self.is_closed = False + + return Manager() +``` + +- [ ] **Step 2: Add failing version, capability, and context tests** + +```python +def test_newton_backend_requires_exact_contract(monkeypatch): + monkeypatch.setattr(dexsim, "__version__", "0.4.3") + backend = NewtonPhysicsBackend(SimpleNamespace()) + with pytest.raises(RuntimeError, match="0.4.4"): + backend._validate_integration_contract() + + +def test_capabilities_are_structured(): + caps = NewtonPhysicsBackend(SimpleNamespace()).capabilities + assert caps.runtime_topology_mutation == frozenset({"rigid", "articulation"}) + assert caps.multi_world is True + assert caps.articulation_acceleration is False + assert "soft_body" not in caps.asset_kinds + assert "cloth" not in caps.asset_kinds + + +def test_scene_context_uses_full_arena_transform(fake_manager): + context = BackendSceneContext(fake_manager) + transforms = context.arena_transforms + assert transforms.shape == (2, 4, 4) + local = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]) + world = context.local_to_world(local, torch.tensor([1])) + roundtrip = context.world_to_local(world, torch.tensor([1])) + assert torch.allclose(roundtrip, local, atol=1e-6) + assert not torch.allclose(world[:, :3], local[:, :3] + transforms[1, :3, 3]) +``` + +The fake second arena must contain both non-zero translation and a 90-degree Z +rotation so the last assertion detects translation-only conversion. + +- [ ] **Step 3: Run the pure-Python tests and confirm missing types/validation** + +```bash +pytest -q tests/sim/test_backend_parity.py tests/sim/test_newton_scene_context.py -m "not requires_sim" +``` + +Expected: missing capability/context contracts. + +- [ ] **Step 4: Add structured backend contracts** + +```python +@dataclass(frozen=True, slots=True) +class PhysicsCapabilities: + asset_kinds: frozenset[str] + runtime_topology_mutation: frozenset[str] + solver_gradients: frozenset[str] + cuda_graph: bool + partial_reset: bool + forward_kinematics: bool + heterogeneous_joint_spans: bool + runtime_collision_filter: bool + contact_sensor: bool + articulation_acceleration: bool + multi_world: bool + + +@dataclass(frozen=True, slots=True) +class PhysicsPrepareResult: + generation: int | None + did_build: bool + did_rebuild: bool + added_entities: tuple[object, ...] = () + removed_entities: tuple[object, ...] = () +``` + +Keep every existing `supports_*` property as a wrapper over `capabilities`. +The default backend returns `generation=None` and preserves existing behavior. +Add abstract/default implementations for `model_generation`, +`queue_initialization(obj)`, `prepare() -> PhysicsPrepareResult`, and +idempotent `close()` so later tasks do not branch on backend names. + +- [ ] **Step 5: Add strict config and dependency validation** + +Change the dependency to: + +```toml +"dexsim_engine==0.4.4", +``` + +Validate positive `physics_dt`, positive `num_substeps`, normalized device, +recognized solver parameters, gradient/solver compatibility, broad phase, and +CUDA graph combinations before world construction. On backend activation, +require both `dexsim.__version__ == "0.4.4"` and +`NEWTON_INTEGRATION_API_VERSION == 2`. Accept local source builds whose public +version is `0.4.4+` by comparing +`packaging.version.Version(dexsim.__version__).base_version` to `"0.4.4"`; +reject every other base version. + +- [ ] **Step 6: Implement explicit owner context** + +```python +class BackendSceneContext: + def __init__(self, manager: SimulationManager) -> None: + self._manager_ref = weakref.ref(manager) + + @property + def manager(self) -> SimulationManager: + manager = self._manager_ref() + if manager is None or manager.is_closed: + raise RuntimeError("BackendSceneContext owner is closed.") + return manager + + @property + def world(self): + return self.manager.get_world() + + @property + def scene(self): + return self.manager.physics.get_scene() + + @property + def physics(self): + return self.manager.physics + + @property + def generation(self) -> int | None: + return self.manager.physics.model_generation +``` + +Build device-resident `(N, 4, 4)` world transforms and inverses from each +arena root node's full world pose. Provide batched `local_to_world()` and +`world_to_local()` for pose tensors in `xyzw`. Maintain weak maps from arena +and entity native handles to contexts. `register_entities()` and +`for_entities()` provide the source-compatible constructor fallback without +consulting a default world or default SimulationManager; unknown external +entities raise an ownership error that tells the caller to pass `context=`. +Expose `register_entity(entity, ref=None)`, `register_entities(entities)`, and +`entity_ref(entity)`; the last method returns the DexSim stable reference +recorded during Newton attachment. + +Use homogeneous composition for both conversion directions: + +```python +def _xyzw_pose_to_matrix(pose: torch.Tensor) -> torch.Tensor: + matrix = torch.eye(4, dtype=pose.dtype, device=pose.device).repeat( + pose.shape[0], 1, 1 + ) + matrix[:, :3, 3] = pose[:, :3] + matrix[:, :3, :3] = matrix_from_quat( + convert_quat(pose[:, 3:7], to="wxyz") + ) + return matrix + + +def _matrix_to_xyzw_pose(matrix: torch.Tensor) -> torch.Tensor: + quat = convert_quat(quat_from_matrix(matrix[:, :3, :3]), to="xyzw") + return torch.cat((matrix[:, :3, 3], quat), dim=-1) + + +def local_to_world(self, pose: torch.Tensor, env_ids: torch.Tensor) -> torch.Tensor: + local = _xyzw_pose_to_matrix(pose) + world = torch.bmm(self.arena_transforms[env_ids.long()], local) + return _matrix_to_xyzw_pose(world) + + +def world_to_local(self, pose: torch.Tensor, env_ids: torch.Tensor) -> torch.Tensor: + world = _xyzw_pose_to_matrix(pose) + local = torch.bmm(self.inverse_arena_transforms[env_ids.long()], world) + return _matrix_to_xyzw_pose(local) +``` + +- [ ] **Step 7: Run config/context/capability tests** + +```bash +pytest -q tests/sim/test_backend_parity.py tests/sim/test_newton_scene_context.py tests/sim/test_physics_attrs.py +``` + +Expected: all pass without creating a real simulation for pure contract cases. + +- [ ] **Step 8: Commit the EmbodiChain contract foundation** + +```bash +git add pyproject.toml embodichain/lab/sim/cfg.py embodichain/lab/sim/physics/context.py embodichain/lab/sim/physics/base.py embodichain/lab/sim/physics/default.py embodichain/lab/sim/physics/newton.py embodichain/lab/sim/physics/__init__.py tests/sim/newton_contract_test_utils.py tests/sim/test_backend_parity.py tests/sim/test_newton_scene_context.py +git commit -m "refactor(sim): add explicit physics scene contracts" +``` + +--- + +### Task 8: Route EmbodiChain spawning and objects through explicit ownership + +**Files:** + +- Modify: `embodichain/lab/sim/common.py` +- Modify: `embodichain/lab/sim/sim_manager.py` +- Modify: `embodichain/lab/sim/utility/sim_utils.py` +- Modify: `embodichain/lab/sim/objects/rigid_object.py` +- Modify: `embodichain/lab/sim/objects/articulation.py` +- Modify: `embodichain/lab/sim/objects/robot.py` +- Modify: `embodichain/lab/sim/objects/backends/default.py` +- Modify: `tests/sim/test_newton_scene_context.py` +- Modify: `tests/sim/test_newton_finalize_lifecycle.py` + +**Interfaces:** + +- Produces: manager-owned context passed into all physical objects/views. +- Produces: explicit pending-initialization queue; constructors never invoke overridable `reset()`. +- Consumes: `BackendSceneContext` and DexSim `attach_rigid_body`. + +- [ ] **Step 1: Add failing no-global and no-constructor-reset tests** + +```python +def test_rigid_and_articulation_construction_do_not_use_default_world( + monkeypatch, newton_sim +): + monkeypatch.setattr(dexsim, "default_world", lambda: (_ for _ in ()).throw(AssertionError("global"))) + rigid = newton_sim.add_rigid_object(box_cfg("box")) + art = newton_sim.add_articulation(arm_cfg("arm")) + assert rigid.context is newton_sim.scene_context + assert art.context is newton_sim.scene_context + + +def test_physical_object_context_keyword_is_source_compatible(): + assert inspect.signature(RigidObject).parameters["context"].default is None + assert inspect.signature(Articulation).parameters["context"].default is None + + +def test_batch_entity_constructor_never_calls_virtual_reset(): + class Probe(BatchEntity): + def reset(self, env_ids=None): + raise AssertionError("virtual reset from base constructor") + def set_local_pose(self, pose, env_ids=None): + return None + def get_local_pose(self, to_matrix=False): + return torch.zeros(1, 7) + Probe( + cfg=ObjectBaseCfg(uid="probe"), + entities=[object()], + device=torch.device("cpu"), + ) +``` + +- [ ] **Step 2: Run focused tests and confirm global lookup/base reset failures** + +```bash +pytest -q tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +``` + +- [ ] **Step 3: Remove base virtual reset and pass context from the manager** + +`BatchEntity.__init__` stores fields only. Keep the `auto_reset` keyword for +source compatibility, but ignore it and emit a one-time deprecation warning +when `True`; no virtual method is called. + +Create `self.scene_context = BackendSceneContext(self)` immediately after +backend activation and arena construction. Pass it explicitly: + +```python +rigid_obj = RigidObject( + cfg=cfg, + entities=obj_list, + device=self.device, + context=self.scene_context, +) +self.physics.queue_initialization(rigid_obj) +``` + +Use the same pattern for articulations and robots. Default-only light and rigid +group constructors call their own `reset()` explicitly after all subclass +fields are initialized, preserving current behavior. + +Add `context: BackendSceneContext | None = None` as the final keyword to rigid, +articulation, and robot constructors. Resolve `None` with +`BackendSceneContext.for_entities(entities)`. The manager registers spawned +entities against its context before constructing their wrapper, so existing +positional constructor calls retain their signature and no global owner lookup +is needed. + +- [ ] **Step 4: Replace EmbodiChain private Newton attachment** + +In `_attach_newton_rigidbody_desc`, retain EmbodiChain descriptor resolution +and warnings, then call only: + +```python +manager = context.physics.newton_manager +entity_ref = manager.attach_rigid_body( + obj, + actor_type=body_type, + shape_type=shape_type, + body_desc=body, + shape_desc=shape, +) +context.register_entity(obj, entity_ref) +``` + +Delete imports of `register_mesh_object_to_newton_patch`, +`_get_entity_native_handle`, writes to `mgr.dexsim_meta`, and hard-coded world +`-1`. The standard legacy `add_rigidbody` route stores the returned reference +through the same context registration method. + +Thread `context` through `load_mesh_objects_from_cfg`, +`spawn_rigid_object_entities`, `spawn_articulation_entities`, and +`spawn_usd_articulation_entities`. Compatibility defaults resolve from the +provided entities; the SimulationManager core path always passes its context. +Replace `_is_newton_backend_active()`, `_newton_solver_type()`, and +`get_dexsim_arenas()` use in physical spawn/object paths with context +properties. Object `destroy()` methods remove entities from their owning +context arenas, never from a default world. + +- [ ] **Step 5: Run spawn and default-backend regressions** + +```bash +pytest -q tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +pytest -q tests/sim/objects/test_rigid_object.py tests/sim/objects/test_articulation.py -k "constructor or spawn or desc_native" +``` + +Expected: selected tests pass; no core object construction resolves a default +world. + +- [ ] **Step 6: Commit explicit ownership and spawning** + +```bash +git add embodichain/lab/sim/common.py embodichain/lab/sim/sim_manager.py embodichain/lab/sim/utility/sim_utils.py embodichain/lab/sim/objects/rigid_object.py embodichain/lab/sim/objects/articulation.py embodichain/lab/sim/objects/robot.py embodichain/lab/sim/objects/backends/default.py tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +git commit -m "refactor(sim): make physical object ownership explicit" +``` + +--- + +### Task 9: Rebind rigid views by generation and initialize only new objects + +**Files:** + +- Create: `tests/sim/test_newton_rebuild_bindings.py` +- Modify: `embodichain/lab/sim/physics/newton.py` +- Modify: `embodichain/lab/sim/objects/backends/newton.py` +- Modify: `embodichain/lab/sim/objects/rigid_object.py` +- Modify: `tests/sim/objects/test_rigid_object.py` +- Modify: `tests/sim/test_newton_finalize_lifecycle.py` + +**Interfaces:** + +- Produces: `NewtonRigidBodyView._ensure_binding()` with O(1) generation check. +- Produces: event-driven cache invalidation and exact pending initialization. +- Consumes: DexSim `RigidEntityBinding`, prepare result, and rebuild event. + +- [ ] **Step 1: Add failing rebind/state-preservation/initialization tests** + +```python +def test_rigid_view_refreshes_once_after_generation_change(newton_sim): + old = newton_sim.add_rigid_object(box_cfg("old", z=1.0)) + first = newton_sim.physics.prepare() + old_generation = old._data.body_view.binding.generation + old_pose = old.get_local_pose().clone() + new = newton_sim.add_rigid_object(box_cfg("new", z=2.0)) + second = newton_sim.physics.prepare() + assert second.generation == first.generation + 1 + assert old_generation == first.generation + assert old._data.body_view.binding.generation == second.generation + assert new._data.body_view.binding.generation == second.generation + assert torch.allclose(old.get_local_pose(), old_pose, atol=1e-5) + assert torch.allclose(new.get_local_pose()[:, 2], torch.tensor([2.0] * new.num_instances)) + + +def test_existing_object_is_not_reset_on_rebuild(newton_sim, mocker): + old = newton_sim.add_rigid_object(box_cfg("old")) + newton_sim.physics.prepare() + reset = mocker.spy(old, "reset") + newton_sim.add_rigid_object(box_cfg("new")) + newton_sim.physics.prepare() + reset.assert_not_called() +``` + +- [ ] **Step 2: Run the new file and observe stale IDs/global reset** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py +``` + +Expected: stale binding or `_reset_entities_after_finalize` resets existing +objects. + +- [ ] **Step 3: Replace permanent IDs with one binding per view** + +```python +def _ensure_binding(self) -> RigidEntityBinding: + generation = self._context.generation + if self._binding is None or self._binding.generation != generation: + refs = tuple(self._context.entity_ref(entity) for entity in self._entities) + self._binding = self._manager.bind_rigid_entities(refs, device=self._device) + self._invalidate_derived_caches() + self._binding.assert_current(self._manager) + return self._binding + +@property +def binding(self) -> RigidEntityBinding: + return self._ensure_binding() +``` + +Every fetch/apply operation calls this once and passes its batched device IDs +to `NewtonPhysicsScene`. Remove permanent `_body_ids`, sorted-ID, and XY-offset +caches; rebuild derived caches only inside `_invalidate_derived_caches()`. + +- [ ] **Step 4: Subscribe backend lifecycle and initialize only prepare additions** + +`NewtonPhysicsBackend.activate()` subscribes to model rebuilt events and marks +registered views dirty. `queue_initialization(obj)` stores object identity and +its stable refs. After `prepare()` and view rebind, initialize only queued +objects whose refs appear in `result.added_entities`, then remove them from the +queue. First build follows the same path. Existing objects are never reset by a +rebuild. + +- [ ] **Step 5: Convert rigid poses with full context transforms** + +`fetch_pose()` converts DexSim world poses to arena-local with +`context.world_to_local`; `apply_pose()` converts local to world with +`context.local_to_world`. Remove all `[:2, 3]`, XY-only, and +`get_all_arenas()` conversion logic from the Newton view. + +- [ ] **Step 6: Run rigid and lifecycle tests** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_finalize_lifecycle.py tests/sim/objects/test_rigid_object.py -k "Newton or newton or local_pose or reset" +``` + +Expected: selected tests pass for 1, 2, and 8 arenas, including a rotated +arena fixture. + +- [ ] **Step 7: Commit rigid rebinding** + +```bash +git add embodichain/lab/sim/physics/newton.py embodichain/lab/sim/objects/backends/newton.py embodichain/lab/sim/objects/rigid_object.py tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_rigid_object.py tests/sim/test_newton_finalize_lifecycle.py +git commit -m "fix(newton): rebind rigid views after runtime rebuild" +``` + +--- + +### Task 10: Rebind articulation views, separate q/qd widths, and enforce FK/capabilities + +**Files:** + +- Modify: `embodichain/lab/sim/objects/backends/newton.py` +- Modify: `embodichain/lab/sim/objects/articulation.py` +- Modify: `embodichain/lab/sim/objects/robot.py` +- Modify: `tests/sim/test_newton_rebuild_bindings.py` +- Modify: `tests/sim/objects/test_articulation.py` +- Modify: `tests/sim/objects/test_rigid_object.py` +- Modify: `tests/sim/objects/test_robot.py` + +**Interfaces:** + +- Produces: `NewtonArticulationView._ensure_binding()` and real `compute_kinematics()`. +- Produces: `_articulation_buffer_shapes(num_instances, qpos_width, qvel_width, target_qpos_width, target_qvel_width)` used by `ArticulationData`. +- Produces: explicit unsupported-operation errors for q acceleration and other absent capabilities. +- Consumes: DexSim `ArticulationBinding` and full context transforms. + +- [ ] **Step 1: Add failing articulation rebind, frame, width, and FK tests** + +```python +def test_articulation_view_rebinds_and_preserves_targets(newton_sim): + arm = newton_sim.add_articulation(arm_cfg("arm")) + newton_sim.physics.prepare() + q = torch.full((arm.num_instances, arm.qpos_width), 0.1, device=arm.device) + qd = torch.full((arm.num_instances, arm.qvel_width), 0.2, device=arm.device) + target_q = torch.full( + (arm.num_instances, arm.target_qpos_width), + -0.1, + device=arm.device, + ) + target_qd = torch.full( + (arm.num_instances, arm.target_qvel_width), + -0.2, + device=arm.device, + ) + arm.set_qpos(q, target=False) + arm.set_qvel(qd, target=False) + arm.set_qpos(target_q, target=True) + arm.set_qvel(target_qd, target=True) + before_link = arm.get_link_pose(arm.link_names[-1]).clone() + newton_sim.add_rigid_object(box_cfg("topology_change")) + newton_sim.physics.prepare() + assert torch.allclose(arm.get_qpos(), q) + assert torch.allclose(arm.get_qvel(), qd) + assert torch.allclose(arm.get_qpos(target=True), target_q) + assert torch.allclose(arm.get_qvel(target=True), target_qd) + assert torch.allclose(arm.get_link_pose(arm.link_names[-1]), before_link, atol=1e-5) + + +def test_qpos_write_updates_link_pose_without_physics_step(newton_sim): + arm = newton_sim.add_articulation(arm_cfg("arm")) + newton_sim.physics.prepare() + before = arm.get_link_pose(arm.link_names[-1]).clone() + q = arm.get_qpos().clone() + q[:, 0] += 0.2 + arm.set_qpos(q, target=False) + after = arm.get_link_pose(arm.link_names[-1]) + assert not torch.allclose(after, before) + + +def test_newton_qacc_is_explicitly_unsupported(newton_sim): + arm = newton_sim.add_articulation(arm_cfg("arm")) + newton_sim.physics.prepare() + with pytest.raises(NotImplementedError, match="acceleration"): + _ = arm.body_data.qacc +``` + +Add pure allocation coverage for distinct widths: + +```python +def test_articulation_buffer_shapes_keep_q_and_qd_widths_distinct(): + shapes = _articulation_buffer_shapes( + num_instances=3, + qpos_width=7, + qvel_width=6, + target_qpos_width=6, + target_qvel_width=6, + ) + assert shapes["qpos"] == (3, 7) + assert shapes["target_qpos"] == (3, 6) + assert shapes["qvel"] == (3, 6) + assert shapes["target_qvel"] == (3, 6) + assert shapes["qf"] == (3, 6) +``` + +- [ ] **Step 2: Run the selected articulation tests and confirm stale/FK/zero-qacc failures** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_articulation.py -k "Newton or newton or qpos or qacc or link_pose" +``` + +- [ ] **Step 3: Bind articulation IDs, link IDs, and spans as one unit** + +Implement the same O(1) generation guard as the rigid view. Resolve root/link +IDs and joint spans only through `manager.bind_articulations`; delete direct +reads of `dexsim_meta_links`, `get_gpu_index()`, and permanent articulation ID +lists from the Newton view. + +Expose and use separate widths: + +```python +@property +def qpos_width(self) -> int: + return self._ensure_binding().qpos_width + +@property +def qvel_width(self) -> int: + return self._ensure_binding().qvel_width + +@property +def target_qpos_width(self) -> int: + return self._ensure_binding().target_qpos_width + +@property +def target_qvel_width(self) -> int: + return self._ensure_binding().target_qvel_width +``` + +Allocate current `qpos` with `qpos_width`, current `qvel/qf` with +`qvel_width`, and targets with their explicit binding widths. For DexSim 0.4.4 +both target position and target velocity are per-DOF and therefore their +widths equal `qvel_width`, but callers consume the explicit properties rather +than inferring that relationship. Existing `dof` remains a compatibility alias +for all-1-DOF assets and raises a clear error when its old ambiguous meaning +would truncate a non-scalar joint. +Expose matching read-only `Articulation.qpos_width` and +`Articulation.qvel_width` properties plus `target_qpos_width` and +`target_qvel_width` that delegate to the view. + +```python +def _articulation_buffer_shapes( + num_instances: int, + qpos_width: int, + qvel_width: int, + target_qpos_width: int, + target_qvel_width: int, +) -> dict[str, tuple[int, int]]: + return { + "qpos": (num_instances, qpos_width), + "target_qpos": (num_instances, target_qpos_width), + "qvel": (num_instances, qvel_width), + "target_qvel": (num_instances, target_qvel_width), + "qf": (num_instances, qvel_width), + } +``` + +- [ ] **Step 4: Implement frame-correct root/link reads and FK** + +Convert root and link world poses with the complete arena transform. On current +q writes, call the DexSim public FK invalidation/evaluation path for the +affected articulation IDs. Implement `compute_kinematics(env_ids)` by mapping +the selected environment rows through the binding's `articulation_ids_host`, +constructing a boolean mask of length `model.articulation_count`, and calling +the public manager FK method; it must not be a no-op. + +Replace fabricated q acceleration with: + +```python +raise NotImplementedError( + "Newton articulation joint acceleration is not exposed by DexSim 0.4.4." +) +``` + +Use the same explicit pattern for unsupported runtime collision-filter or +sensor operations; do not return plausible zeros or success. + +Add `@pytest.mark.gpu` to the Newton-backed rigid, articulation, and robot test +classes because they configure `device="cuda"` even though their node IDs do +not contain `cuda`. This keeps them out of the CPU job and includes them in the +serial `--run-gpu -m gpu` merge gate. + +- [ ] **Step 5: Run articulation and robot regressions** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_articulation.py tests/sim/objects/test_robot.py -k "Newton or newton" +``` + +Expected: Newton articulation/robot selections pass; skips only correspond to +capabilities explicitly outside Stage 1. + +- [ ] **Step 6: Commit articulation rebinding** + +```bash +git add embodichain/lab/sim/objects/backends/newton.py embodichain/lab/sim/objects/articulation.py embodichain/lab/sim/objects/robot.py tests/sim/test_newton_rebuild_bindings.py tests/sim/objects/test_articulation.py tests/sim/objects/test_rigid_object.py tests/sim/objects/test_robot.py +git commit -m "fix(newton): bind articulation state by model generation" +``` + +--- + +### Task 11: Complete SimulationManager mutation, exact update, multi-instance, and close lifecycle + +**Files:** + +- Create: `tests/sim/test_newton_multi_manager.py` +- Modify: `embodichain/lab/sim/sim_manager.py` +- Modify: `embodichain/lab/sim/physics/base.py` +- Modify: `embodichain/lab/sim/physics/default.py` +- Modify: `embodichain/lab/sim/physics/newton.py` +- Modify: `tests/sim/test_newton_finalize_lifecycle.py` +- Modify: `tests/sim/test_newton_rebuild_bindings.py` + +**Interfaces:** + +- Produces: idempotent `SimulationManager.close()` and close-before-remove `reset()`. +- Produces: exact requested-step update after prepare. +- Consumes: pending initialization and per-world DexSim close. + +- [ ] **Step 1: Add failing update/remove/close/two-manager tests** + +```python +def test_update_runs_exact_requested_steps_after_rebuild(newton_sim): + box = newton_sim.add_rigid_object(box_cfg("box", z=1.0)) + manager = newton_sim.newton_manager + before = manager._sim_time + newton_sim.update(physics_dt=0.01, step=3) + assert manager._sim_time == pytest.approx(before + 0.03) + newton_sim.add_rigid_object(box_cfg("new", z=2.0)) + before = manager._sim_time + newton_sim.update(physics_dt=0.01, step=2) + assert manager._sim_time == pytest.approx(before + 0.02) + + +def test_remove_invalidates_and_survivor_rebinds(newton_sim): + keep = newton_sim.add_rigid_object(box_cfg("keep")) + remove = newton_sim.add_rigid_object(box_cfg("remove")) + newton_sim.physics.prepare() + keep_pose = keep.get_local_pose().clone() + assert newton_sim.remove_asset("remove") is True + result = newton_sim.physics.prepare() + assert result.did_rebuild is True + assert torch.allclose(keep.get_local_pose(), keep_pose, atol=1e-5) + with pytest.raises(Exception, match="removed|closed|stale"): + remove.get_local_pose() + + +def test_close_and_reset_are_idempotent(newton_sim): + instance_id = newton_sim.instance_id + world = newton_sim.get_world() + newton_sim.close() + newton_sim.close() + assert newton_sim.is_closed is True + assert SimulationManager.is_instantiated(instance_id) is False + SimulationManager.reset(instance_id) + assert dexsim.engine.newton_physics.get_newton_manager(world) is None + + +@pytest.mark.gpu +def test_two_same_device_managers_are_isolated(): + def make_sim(): + return SimulationManager( + SimulationManagerCfg( + headless=True, + device="cuda:0", + num_envs=2, + physics_cfg=NewtonPhysicsCfg( + device="cuda:0", use_cuda_graph=False + ), + ) + ) + + first = make_sim() + second = make_sim() + try: + first_box = first.add_rigid_object(box_cfg("first")) + second_box = second.add_rigid_object(box_cfg("second")) + first.update(physics_dt=0.01, step=1) + second.update(physics_dt=0.01, step=1) + assert first.get_world() is not second.get_world() + assert first.get_physics_scene() is not second.get_physics_scene() + assert first.scene_context is not second.scene_context + assert first.newton_manager.world_token != second.newton_manager.world_token + first.add_rigid_object(box_cfg("first_new")) + first.update(physics_dt=0.01, step=1) + assert first.newton_manager.model_generation == 2 + assert second.newton_manager.model_generation == 1 + first.close() + second.update(physics_dt=0.01, step=1) + second_pose = second_box.get_local_pose().clone() + assert torch.isfinite(second_pose).all() + assert first_box.context is not second_box.context + finally: + first.close() + second.close() +``` + +- [ ] **Step 2: Run the lifecycle files and confirm current reset/leak behavior** + +```bash +pytest -q tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py --run-gpu +``` + +- [ ] **Step 3: Return prepare results and preserve exact update count** + +Change `PhysicsBackend.prepare()` and `ensure_initialized()` to return +`PhysicsPrepareResult`. `SimulationManager.update()` calls prepare once, then +executes the existing world update loop exactly `step` times. It never adds a +warmup step and never drops the first requested step. + +- [ ] **Step 4: Make every topology mutation invalidate and every removal close its view** + +After successful rigid/articulation/robot add or remove, call +`physics.invalidate()`. Removal destroys the wrapper, unregisters its context +refs, and leaves surviving refs queued for rebind but not reset. Soft/cloth +add/remove on Newton raises `NotImplementedError` before mutating registries. + +- [ ] **Step 5: Add idempotent close and safe registry allocation** + +```python +def close(self) -> None: + if self._is_closed: + return + self._is_closed = True + first_error = None + try: + self.wait_window_record_saves() + self.physics.close() + except Exception as exc: + first_error = exc + try: + if self._world is not None: + self._world.quit() + except Exception as exc: + if first_error is None: + first_error = exc + finally: + self._instances.pop(self.instance_id, None) + if first_error is not None: + raise first_error +``` + +`reset(instance_id)` calls `instance.close()` before removing it. Preserve +`destroy(exit_process=...)` as a wrapper around close plus its documented +process-exit policy. Allocate instance IDs monotonically rather than from +`len(_instances)`, so closing a non-last manager cannot overwrite a live entry. +DexSim manager close is idempotent, so the backend-close/world-quit sequence is +safe even when `World.quit()` invokes the same integration teardown again. + +- [ ] **Step 6: Run lifecycle, multi-manager, and default-backend tests** + +```bash +pytest -q tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py tests/sim/test_backend_parity.py tests/sim/objects/test_rigid_object.py tests/sim/objects/test_articulation.py --run-gpu +``` + +Expected: zero failures; fixture teardown finds no stale world or manager +registration. + +- [ ] **Step 7: Commit manager lifecycle completion** + +```bash +git add embodichain/lab/sim/sim_manager.py embodichain/lab/sim/physics/base.py embodichain/lab/sim/physics/default.py embodichain/lab/sim/physics/newton.py tests/sim/test_newton_finalize_lifecycle.py tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py +git commit -m "fix(sim): close and rebuild Newton managers safely" +``` + +--- + +### Task 12: Document the contract and run the Stage 1 merge gate + +**Files:** + +- Modify: `docs/source/overview/sim/sim_manager.md` +- Modify: `design/newton-backend-design.md` +- Verify: both repositories' Stage 1 diffs. + +**Interfaces:** + +- Consumes: all Stage 1 interfaces and tests. +- Produces: a verified foundation for the later Stage 2 differentiable plan. + +- [ ] **Step 1: Update public documentation with executable examples** + +Document: + +```python +sim = SimulationManager( + SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg(device="cuda:0"), + num_envs=4, + headless=True, + ) +) +try: + cube = sim.add_rigid_object(cube_cfg) + sim.finalize_newton_physics() + sim.update(step=1) +finally: + sim.close() +``` + +State that prepare/finalize does not advance time, add/remove of rigid bodies +and articulations rebuilds transactionally, old runtime IDs must not be cached, +soft/cloth topology mutation is unsupported, q acceleration is unavailable, +and two managers are isolated. Mark the old Target 4 implementation claims as +superseded by the 2026-07-13 design and record only tests actually passing. + +- [ ] **Step 2: Run DexSim formatting and the full Newton test directory** + +```bash +cd /root/sources/dexsim +black --check --diff python/dexsim/engine/newton_physics python/test/engine/newton_physics +pytest -q python/test/engine/newton_physics +``` + +Expected: Black exits zero and pytest reports zero failures. + +- [ ] **Step 3: Run EmbodiChain focused CPU/headless tests** + +```bash +cd /root/sources/EmbodiChain +pytest -q tests/sim/test_backend_parity.py tests/sim/test_physics_attrs.py tests/sim/test_newton_scene_context.py tests/sim/test_newton_finalize_lifecycle.py +``` + +Expected: zero failures. + +- [ ] **Step 4: Run EmbodiChain serial GPU Newton integration tests** + +```bash +pytest -q tests/sim/test_newton_rebuild_bindings.py tests/sim/test_newton_multi_manager.py tests/sim/objects/test_rigid_object.py tests/sim/objects/test_articulation.py tests/sim/objects/test_robot.py --run-gpu -m gpu +``` + +Expected: zero failures; skips are listed and checked against structured +capabilities. + +- [ ] **Step 5: Run the complete EmbodiChain regression and docs build** + +```bash +pytest -q tests +black --check --diff --color ./ +LC_ALL=C.UTF-8 LANG=C.UTF-8 make -C docs html +``` + +Expected: zero pytest failures, Black leaves all files unchanged, and Sphinx +builds without new warnings/errors. + +- [ ] **Step 6: Inspect both diffs and dependency/API versions** + +```bash +if rg -n "dexsim\.default_world\(\)|get_physics_scene\(\)|SimulationManager\.get_instance\(" embodichain/lab/sim/objects/rigid_object.py embodichain/lab/sim/objects/articulation.py embodichain/lab/sim/objects/backends/newton.py; then + echo "core Newton object/view path still contains global owner lookup" >&2 + exit 1 +fi +if rg -n "register_mesh_object_to_newton_patch|_get_entity_native_handle|dexsim_meta" embodichain/lab/sim/utility/sim_utils.py embodichain/lab/sim/objects/backends/newton.py; then + echo "EmbodiChain still consumes private DexSim Newton integration state" >&2 + exit 1 +fi +git -C /root/sources/dexsim diff --check dev...HEAD +git -C /root/sources/dexsim log --oneline --decorate dev..HEAD +git -C /root/sources/EmbodiChain diff --check main...HEAD +git -C /root/sources/EmbodiChain log --oneline --decorate main..HEAD +python - <<'PY' +import dexsim +from dexsim.engine.newton_physics import NEWTON_INTEGRATION_API_VERSION +assert dexsim.__version__.split("+")[0] == "0.4.4" +assert NEWTON_INTEGRATION_API_VERSION == 2 +print(dexsim.__version__, NEWTON_INTEGRATION_API_VERSION) +PY +``` + +Expected: no whitespace errors, reviewable commit series in each repository, +package base version `0.4.4`, API version `2`. + +- [ ] **Step 7: Commit Stage 1 documentation** + +```bash +git add docs/source/overview/sim/sim_manager.md design/newton-backend-design.md +git commit -m "docs: describe Newton runtime lifecycle contracts" +``` + +- [ ] **Step 8: Stop at the Stage 1 review gate** + +Report exact command outputs, failures/skips, both branch heads, and remaining +known upstream limitations. Do not start Stage 2. After Stage 1 is accepted, +use `superpowers:brainstorming` only if Stage 1 changed the approved design; +otherwise use `superpowers:writing-plans` to create the dependent +differentiable dynamics/kinematics implementation plan. + +--- + +## Spec Coverage Checklist + +| Specification requirement | Implemented by | +|---|---| +| DexSim public API/version/generation/prepare | Tasks 1–3 | +| Public rigid attachment and canonical descriptors | Task 2 | +| Generation-aware rigid/articulation bindings | Tasks 3, 9, 10 | +| Transactional rebuild and rigid state preservation | Task 4 | +| Active model-generation lease blocks rebuild | Tasks 1, 4 | +| Articulation state/control preservation and q/qd spans | Task 5 | +| Full arena/world frame conversion | Tasks 7, 9, 10 | +| No global/default-world core ownership | Tasks 7–8 | +| Pending initialization without constructor virtual reset | Tasks 8–9 | +| Two same-device worlds/managers and deterministic cleanup | Tasks 6, 11 | +| Structured capabilities and explicit unsupported errors | Tasks 7, 10 | +| Exact prepare/step count | Tasks 1, 4, 11 | +| Existing public API and default backend compatibility | Tasks 7–12 | +| Documentation and complete merge gate | Task 12 | 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 000000000..2b93f9bd9 --- /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/docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md b/docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md new file mode 100644 index 000000000..133189299 --- /dev/null +++ b/docs/superpowers/specs/2026-07-13-newton-runtime-contracts-design.md @@ -0,0 +1,749 @@ +# Newton Runtime Contracts and Differentiable Execution Design + +**Status:** Approved in design discussion; awaiting written-spec review + +**Date:** 2026-07-13 + +**EmbodiChain branch:** `feature/newton-physics-backend` + +**DexSim implementation branch:** `feature/embodichain-newton-contracts` + +**Target DexSim package version:** `0.4.4` + +## 1. Purpose and supersession + +This specification replaces the multi-environment and differentiable-runtime +portions of: + +- `docs/superpowers/specs/2026-06-21-newton-backend-pr-design.md` +- `docs/superpowers/plans/2026-06-22-newton-backend-pr.md` +- the corresponding Target 4 and Target 5 status claims in + `design/newton-backend-design.md` + +The previous documents describe mutually incompatible clone-at-finalize and +spawn-time clone designs, treat rebuild-time runtime IDs as permanent, and +conflate a forward-kinematics demonstration with differentiable Newton +dynamics. They remain historical records but are not implementation sources +after this specification is accepted. + +The work is delivered in two sequential stages: + +1. A coordinated DexSim and EmbodiChain refactor covering the Newton public + integration contract, lifecycle, multi-world isolation, runtime topology + mutation, rigid bodies, and articulations. +2. A differentiable execution layer built on the resulting authoritative + state, generation, binding, and lifecycle contracts. It supports both real + solver dynamics and pure kinematics. + +Stage 2 starts only after Stage 1 correctness and lifecycle tests pass. + +## 2. Goals + +### 2.1 Stage 1 goals + +- Make DexSim the sole authority for Newton model, state, control, contacts, + entity metadata, runtime mappings, and model generation. +- Remove EmbodiChain use of private DexSim registry, registration, and + `dexsim_meta` details. +- Preserve the existing public `SimulationManager`, `RigidObject`, and + `Articulation` call surfaces. +- Support initial build and post-finalize add/remove for rigid bodies and + articulations. +- Preserve surviving rigid and articulation state across rebuilds. +- Rebind all runtime IDs after every successful model rebuild. +- Correctly isolate global and child-arena Newton worlds. +- Support two or more simultaneous `SimulationManager`/DexSim `World` + instances in one process, including same-GPU operation and deterministic + teardown. +- Make arena-local and world-frame pose semantics explicit and consistent. +- Support articulation topologies whose position and velocity widths differ, + including spherical and free joints at the binding-contract level. +- Preserve the default physics backend behavior. + +### 2.2 Stage 2 goals + +- Make the default differentiable path execute + `DifferentiableStepper` and the configured Newton solver. +- Match normal simulation time and substep semantics exactly. +- Support pure-kinematics environment steps through `newton.eval_fk` without + misrepresenting them as dynamics. +- Support non-zero action gradients, finite-difference validation, and + continuous multi-step differentiation. +- Preserve the normal environment lifecycle while making only a minimal, + explicit differentiable-output addition to the functor surface. +- Make state-buffer ownership safe across forward, backward, reset, rebuild, + and multiple worlds. + +## 3. Non-goals + +- A general functor-system rewrite. +- Runtime topology mutation for soft bodies or cloth. Such attempts fail + explicitly in this iteration. +- Replacing EmbodiChain's environment framework with IsaacLab's architecture. +- Copying IsaacLab's class-level physics singleton, USD/Fabric coupling, or + backend discovery by class-name convention. +- Broad renderer, sensor, or solver performance optimization unrelated to the + new contracts. +- Heterogeneous articulations inside one `Articulation` batch. Separate + batched assets may have different topologies, while instances within one + batch retain the same topology. + +## 4. Architectural boundary + +The ownership rule is: + +> DexSim owns physical truth; EmbodiChain owns environment semantics. + +```text +SimulationManager + -> PhysicsBackend + -> BackendSceneContext + -> DexSim World / NewtonManager + -> arena transforms + -> model generation + -> entity and view registries +``` + +DexSim owns: + +- `ModelBuilder`, `Model`, both runtime `State` buffers, `Control`, contacts, + collision pipeline, solver, and CUDA graph; +- stable entity references and canonical replay descriptors; +- world assignment, runtime body/shape/articulation/link/joint mappings; +- build, rebuild, snapshot, restore, commit, and generation transitions; +- the public tensor and binding contracts used by integrations. + +EmbodiChain owns: + +- backend selection and environment orchestration; +- public object APIs and backend-independent views; +- arena-local frame semantics and arena transform tables; +- pending initialization of newly added objects; +- environment reset, step count, observations, rewards, hooks, and datasets; +- differentiable environment policy and task-specific action/output kernels. + +Core object and utility code must not resolve its owner through +`dexsim.default_world()` or the default `SimulationManager` instance. + +## 5. Lifecycle and generation + +### 5.1 Lifecycle phases + +The integration exposes the following conceptual phases: + +```text +BUILDING + -> MODEL_FINALIZED + -> VIEWS_BOUND + -> SOLVER_READY + -> RUNNING + -> STALE + -> rebuild + -> VIEWS_BOUND +``` + +DexSim may retain its internal state enum, but public results must distinguish +successful model finalization, successful binding readiness, solver readiness, +staleness, and closure. `READY` alone must not ambiguously mean both “model +exists” and “all external consumers are rebound.” + +### 5.2 Model generation + +Each `NewtonManager` has a public, read-only `model_generation`: + +- it starts at `0` before the first finalized model; +- the first successful finalize commits generation `1`; +- every successful model-replacing rebuild increments it once; +- live writes that do not replace or re-index model arrays do not increment it; +- failed candidate builds do not increment it; +- separate worlds maintain independent generations. + +All runtime bindings carry the generation against which they were resolved. +Using a stale binding raises an explicit generation error or causes the owning +view to rebind before access; it never silently uses old IDs. + +### 5.3 Prepare result and rebuild events + +DexSim exposes a public prepare result equivalent to: + +```python +@dataclass(frozen=True) +class NewtonPrepareResult: + generation: int + did_build: bool + did_rebuild: bool + added_entities: tuple[NewtonEntityRef, ...] + removed_entities: tuple[NewtonEntityRef, ...] +``` + +After an atomic runtime commit, DexSim publishes a `MODEL_REBUILT` event with +the old and new generations and the topology delta. Failed builds publish a +failure result but never a success event. EmbodiChain subscribes when its +backend activates and unsubscribes during close. + +`PhysicsBackend.prepare()` returns an EmbodiChain-level result carrying the +same generation and rebuild facts. It establishes a ready-to-step runtime but +does not itself advance simulation time. + +## 6. DexSim public integration contract + +### 6.1 API version + +DexSim exports: + +```python +NEWTON_INTEGRATION_API_VERSION = 2 +``` + +The package patch version advances to `0.4.4`. EmbodiChain pins the exact +package version and validates the integration API version at backend +activation. This prevents two materially different Newton integrations from +sharing an indistinguishable dependency version. + +### 6.2 Stable entity references + +Registration returns an opaque, stable `NewtonEntityRef` containing enough +identity to reject cross-world use. Runtime integer IDs are not exposed as +stable handles. + +The identity is derived from the owning world and DexSim entity, not from a +model body index. Removal invalidates the reference for new bindings while +allowing rebuild snapshots to identify that the entity should be omitted. + +### 6.3 Public rigid-body attachment + +DexSim provides a supported attachment API equivalent to: + +```python +attach_rigid_body( + entity, + *, + actor_type, + shape_type, + physical_attr=None, + body_desc=None, + shape_desc=None, + geometry_desc=None, +) -> NewtonEntityRef +``` + +This is the only integration entry point EmbodiChain uses for Newton rigid +bodies. It: + +- resolves the owning manager from the entity's actual arena/world; +- derives global world `-1` or the correct child-world index; +- captures mesh, box, sphere, and other supported geometry parameters; +- stores a canonical descriptor that can be replayed during rebuild; +- supports the legacy `PhysicalAttr` projection and Newton-native body/shape + descriptors; +- makes descriptor ownership per entity so clone mutation cannot alias the + prototype; +- marks a finalized runtime stale when topology changes. + +Clone operations recompute target world metadata from the target arena rather +than copying the prototype's world index. + +### 6.4 Public generation-aware bindings + +DexSim provides binding operations equivalent to: + +```python +bind_rigid_entities(refs) -> RigidEntityBinding +bind_articulations(refs) -> ArticulationBinding +``` + +`RigidEntityBinding` includes generation, body IDs where applicable, shape +IDs, and world IDs. Static entities may have shape IDs without body IDs. + +`ArticulationBinding` includes generation, articulation and link body IDs, +world IDs, and explicit per-active-joint spans for: + +- current q position; +- target q position; +- q velocity; +- target q velocity; +- generalized force/control. + +It separately reports `qpos_width` and `dof`/`qvel_width`. An active-joint +index is never interpreted as a flattened DOF index. + +Bindings use `int32` indices and declare device, dtype, shape, ownership, +mutability, and lifetime. Public state data uses `float32`; public quaternions +use `xyzw`. + +## 7. Transactional rebuild and state restoration + +### 7.1 Rebuild sequence + +Runtime topology mutation follows: + +```text +add/remove entity + -> manager STALE + -> snapshot by stable entity reference + -> build candidate builder/model/state/control/solver + -> restore surviving entities into candidate runtime + -> validate candidate mappings and resources + -> atomically commit candidate runtime + -> generation + 1 + -> publish MODEL_REBUILT + -> EmbodiChain rebinds views + -> initialize only newly added entities + -> FK and required DexSim visual synchronization +``` + +The old runtime is not cleared before the candidate is validated. Candidate +construction may temporarily use additional memory; correctness and rollback +take precedence over rebuild-time peak memory. + +If candidate construction or restoration fails: + +- the manager remains `STALE` and stepping is prohibited; +- the previous runtime remains available for diagnostics but is not presented + as current physical truth for the mutated scene; +- generation does not change; +- no rebuilt event is emitted; +- callers may correct the scene/configuration and retry prepare. + +### 7.2 Snapshot coverage + +Snapshots are keyed by stable entity reference rather than runtime body ID. + +Rigid state coverage: + +- pose; +- linear and angular velocity; +- linear and angular acceleration; +- pending external force and torque; +- both ping-pong state buffers where fields exist. + +Articulation coverage: + +- root pose and velocity; +- current and target q position; +- current and target q velocity; +- generalized forces and active controls; +- relevant drive/control state; +- both ping-pong state buffers. + +Contacts are regenerated and are not restored. Removed entities are omitted. +New entities receive descriptor/default state and are reported in the prepare +result for owner-side initialization. + +### 7.3 Differentiable model leases + +A live differentiable session holds a lease on its model generation. A +topology-changing rebuild is rejected while an outstanding tape depends on +that generation. The user or environment must finish backward or explicitly +close/detach the session before rebuilding. This prevents model arrays from +being freed while Warp autograd still references them. + +## 8. EmbodiChain backend and scene context + +### 8.1 BackendSceneContext + +Every simulated object receives its owner explicitly through a context that +contains: + +- `SimulationManager` identity; +- DexSim `World` and physics scene; +- active `PhysicsBackend`; +- arena list and full world transforms; +- current backend/model generation; +- entity/view registration helpers. + +Objects no longer call `dexsim.default_world()`, global +`get_physics_scene()`, or default-instance arena utilities in core paths. + +### 8.2 View rebinding + +Rigid and articulation views store a binding rather than permanent IDs. Before +each batch operation they perform an O(1) generation comparison. A mismatch +causes one binding refresh for the complete batch, invalidating dependent +sorted-ID and arena-transform caches. + +READY steady state does not re-resolve IDs, allocate bindings, or loop over +entities in Python. + +### 8.3 Pending initialization + +EmbodiChain records newly added objects in `pending_initialization`. After a +successful prepare and view rebind, only those objects receive their initial +state/reset. Existing objects retain the state restored by DexSim. + +Base entity constructors do not call overridable `reset()` methods. Object +initialization is an explicit manager/lifecycle phase. + +### 8.4 Frame contract + +The public API distinguishes world and arena-local frames. Existing +`set_local_pose()` and `get_local_pose()` remain arena-local before and after +finalization. + +Conversion uses the complete arena rigid transform, including rotation, not +only XY translation. Root and link pose data returned by DexSim global APIs is +converted in the view. Quaternion convention is consistently `xyzw`. + +Velocity and wrench APIs retain their documented frames; any API whose frame +is currently ambiguous is documented and validated as part of this refactor +rather than inferred from method names. + +### 8.5 Articulation data contract + +EmbodiChain stores separate q-position and velocity/force widths and delegates +active-joint span resolution to `ArticulationBinding`. Current all-1-DOF robot +calls remain source compatible. Spherical and free joints use their actual q +and qd widths. + +Writing current q position triggers required FK invalidation/evaluation before +link pose or visual state is reported. Unsupported data such as articulation +q acceleration is represented through capabilities and raises an explicit +unsupported-operation error rather than returning plausible zeros. + +## 9. Multi-world isolation and cleanup + +Registries, generation counters, entity mappings, solvers, state buffers, +CUDA graphs, and callbacks are keyed by owning world. A reference or binding +from one world cannot be used with another. + +Same-GPU CUDA capture may use a device-level coordinator for capture safety, +but that coordinator does not own simulation state and stores only weak +manager references. Capture timeout and peer diagnostics use existing public +or implemented helpers and cannot wait indefinitely by default. + +EmbodiChain adds an idempotent `SimulationManager.close()` that never exits the +process. It releases backend subscriptions, bindings, DexSim world resources, +CUDA graphs, and instance registry entries. + +Existing cleanup surfaces remain compatible: + +- `destroy()` remains available and preserves its documented exit-process + compatibility behavior; +- `SimulationManager.reset(instance_id)` closes the selected live instance + before removing it, so a new instance cannot inherit its world state; +- repeated close/reset is safe. + +## 10. Capabilities and validation + +Backend capabilities become structured and cover operations in addition to +asset categories. The Newton capability description includes at least: + +- supported asset kinds; +- supported solver and gradient combinations; +- CUDA graph support and invalidation rules; +- partial reset and FK support; +- runtime topology mutation by asset kind; +- heterogeneous q/qd span support; +- runtime collision-filter support; +- contact sensor and acceleration-field support; +- multi-world support. + +Configuration validates positive dt/substeps, device normalization, solver +parameters, gradient requirements, collision pipeline compatibility, and CUDA +graph combinations before finalization. Unconsumed solver parameters are +errors, not silently ignored fields. + +Unsupported operations fail at configuration or API boundaries. In +particular, this iteration rejects runtime topology mutation for soft bodies +and cloth, and reports upstream Newton limitations instead of returning fake +data. + +## 11. Differentiable execution architecture + +### 11.1 Functional core and stateful environment + +The differentiable layer has two levels: + +1. `DifferentiableSession`, a generation-bound functional rollout owner. +2. `DifferentiableEmbodiedEnv`, a stateful Gym/EmbodiChain wrapper. + +The session owns independent state, control, contact, and tape buffers. It +never records directly into buffers that a later environment step will +overwrite before backward. Each forward retains its required buffers in the +autograd context until backward or explicit release. + +The environment maintains a functional session state across steps and mirrors +the resulting state into the normal runtime for non-differentiable consumers, +rendering, and existing object APIs. Mirror writes do not replace tape-owned +buffers. + +Any generation change invalidates the session. Reset detaches the reset +environments from prior episode history. + +### 11.2 Explicit execution modes + +Configuration selects: + +```python +DifferentiableStepCfg( + mode="dynamics" | "kinematics", + bptt_horizon_steps=None, +) +``` + +The existing `truncate_backward_at` input remains accepted as a deprecated +alias for `bptt_horizon_steps`. Its former ambiguous solver-substep meaning is +not retained. Truncation occurs only at environment-step boundaries. + +### 11.3 Dynamics mode + +Dynamics mode must execute `DifferentiableStepper` and the configured solver. +For one environment step: + +```text +write action/control + -> repeat sim_steps_per_control physics steps + -> repeat Newton num_substeps solver steps + -> clear forces + -> apply pending external forces + -> collide + -> DifferentiableStepper.step + -> swap state + -> clear one-shot external inputs +``` + +The total solver step count is: + +```text +sim_steps_per_control * NewtonPhysicsCfg.num_substeps +``` + +The solver dt is `physics_dt / num_substeps`. Control remains applied with the +same cadence as normal simulation. State ownership and final-buffer selection +are independent of odd/even substep count. + +No FK-only fallback is permitted when a task is configured for dynamics. A +zero gradient caused by an unsupported control path fails validation/tests +and must be corrected at the control/solver contract. + +### 11.4 Kinematics mode + +Kinematics mode executes: + +```text +action + -> task-defined q-position update + -> newton.eval_fk + -> body/link state + -> differentiable observations and reward +``` + +It does not run collision or a solver and does not advance physical simulation +time. It does advance the environment episode step. Runtime q position and +DexSim visual state are synchronized after the functional result when enabled +by the environment. + +Kinematics is a first-class, explicitly named mode, not evidence that dynamics +differentiation works. + +### 11.5 Autograd output contract + +The PyTorch/Warp bridge uses explicit outputs equivalent to: + +```python +DifferentiableOutput( + name="reward", + tensor=reward_torch, + source=reward_warp_array, + requires_grad=True, +) +``` + +Every differentiable output has a Warp source whose gradient is seeded by the +custom backward. Observation and reward are handled independently, allowing +both `loss(obs)` and `loss(reward)` to propagate to action. Shape, dtype, +device, contiguity, and finite-value checks occur at the bridge boundary. + +Terminated, truncated, info, and other non-differentiable outputs explicitly +declare that they do not receive a gradient. + +## 12. Environment and minimal functor integration + +`DifferentiableEmbodiedEnv.step()` preserves the normal lifecycle: + +```text +action preprocessing + -> differentiable action mapping + -> dynamics or kinematics execution + -> differentiable observation/reward output + -> ordinary info and termination + -> episode counters + -> hooks and dataset handling + -> reset completed environments +``` + +This iteration does not redesign functors. Tasks provide thin differentiable +action and output adapters, normally implemented with Warp kernels. Existing +ordinary functors continue to run and are detached unless explicitly backed by +a `DifferentiableOutput` source. + +If configuration claims an observation or reward term is differentiable but +no valid Warp source is supplied, construction or the first validated step +raises an error. The system does not silently sever the graph. + +Non-differentiable side effects such as logging, dataset recording, and most +hooks remain outside the tape. + +## 13. Franka reference environments + +The Franka reach task exposes two explicit configurations: + +- **Dynamics:** action maps to a differentiable Newton effort/control path and + must pass through `DifferentiableStepper` and the semi-implicit solver. +- **Kinematics:** action updates joint q position and uses `newton.eval_fk`. + +Both modes use the same documented frame convention. Arena-local targets are +converted consistently against world-frame Newton body state, or body state is +converted to arena-local before reward evaluation. + +The reference task registers through the normal task import path, uses a +deterministic local/fixture asset for required tests, closes its environment in +all test outcomes, and does not depend on a network download for required CI. + +The dynamics acceptance path uses a control mode that is expected to produce a +real action-to-state gradient. The kinematics task remains useful as a faster +smoke test but is reported separately. + +## 14. Error handling + +Errors identify the owning world, entity reference, operation, and expected +versus actual generation where relevant. The design distinguishes: + +- invalid configuration; +- unsupported backend capability; +- stale/removed entity binding; +- closed world/session; +- candidate rebuild failure; +- active differentiable lease blocking rebuild; +- cross-world reference use; +- tensor contract mismatch. + +Runtime collision-filter changes and other setup-only fields either trigger a +documented rebuild or fail explicitly; an API must not report success after +updating metadata that the live model does not consume. + +## 15. Testing and acceptance + +### 15.1 DexSim tests + +- Public rigid attachment for mesh, box, and sphere descriptors. +- Correct global and child-world IDs for one, two, and eight child arenas. +- Clone descriptors are independent and use the target arena's world ID. +- Initial finalize increments generation once. +- Add/remove rebuild increments generation and produces new runtime mappings. +- Rigid pose/velocity/acceleration/external-wrench state survives rebuild. +- Articulation root/current-target q/qd/qf/control state survives rebuild. +- Both state buffers remain valid after restore. +- Candidate build/restore failure leaves a diagnosable non-half-initialized + manager and does not increment generation. +- Revolute, spherical, and free-joint q/qd spans bind correctly. +- Two same-GPU worlds can build, step, rebuild, capture where enabled, and + close independently. +- Repeated close and failed-construction cleanup are safe. + +### 15.2 EmbodiChain tests + +- A view's cached binding changes from the old to the new generation after + rebuild and addresses the correct entity in every environment. +- Adding an entity preserves old-object state and initializes only the new + entity. +- Removing an entity invalidates its binding without changing surviving + object identity or state. +- Rigid, articulation root, and link local poses are correct in every arena, + including non-zero rotations. +- No core object/view lookup resolves through `default_world()`. +- Two `SimulationManager` instances do not share world, scene, arena, + generation, bindings, or cleanup state. +- Existing public object and manager calls remain source compatible. +- Default-backend behavior and tests remain unchanged. +- Capability errors replace silent q-acceleration, collision-filter, or sensor + no-ops covered by the new surface. + +### 15.3 Differentiable tests + +- Dynamics uses the real `DifferentiableStepper` and expected solver-step + count. +- Action gradient is finite, non-zero, and has the expected shape. +- Central finite difference agrees in direction and reasonable tolerance with + autograd for a deterministic small scene. +- At least three consecutive environment steps advance state and backpropagate + safely. +- Odd and even Newton substep counts select the correct final state. +- `loss(obs)` and `loss(reward)` independently propagate to action. +- Kinematics FK pose and gradient match a direct `eval_fk` reference. +- Runtime mirror updates do not corrupt a still-live backward pass. +- BPTT truncation detaches at the requested environment-step boundary. +- Reset prevents cross-episode gradient leakage. +- A generation change invalidates an old session; a live model lease blocks + rebuild until released. +- Dynamics and kinematics handle arena-local targets consistently in multiple + environments. + +### 15.4 Verification order + +The merge gate runs in this order: + +```text +DexSim unit tests + -> EmbodiChain CPU/headless contract tests + -> serial GPU Newton integration tests + -> multi-world lifecycle tests + -> differentiable finite-difference tests + -> complete EmbodiChain regression suite + -> formatting and project pre-commit checks +``` + +GPU and external-simulation tests use the repository's registered markers and +deterministic teardown. Required tests do not silently skip because an asset +download failed. + +## 16. Performance constraints + +The correctness refactor must preserve an efficient steady state: + +- generation checks are O(1); +- READY views do not rebind without a generation change; +- binding refresh is batched; +- per-step pose/state access does not loop over entities in Python; +- arena transform tables are device-resident and rebuilt only when their + owning context changes; +- differentiable rollout buffers are deliberately owned and reused only when + doing so cannot invalidate an outstanding tape. + +Broad solver/render benchmarking is outside this specification. Focused +benchmarks may be added to demonstrate that generation-aware binding does not +regress steady-state batch access. + +## 17. Implementation and repository sequencing + +After this written specification is approved: + +1. Create DexSim branch `feature/embodichain-newton-contracts` from its current + `dev` branch. +2. Write a Stage 1 implementation plan spanning DexSim and EmbodiChain, with + tests preceding implementation changes. +3. Implement and verify the DexSim public contract and lifecycle first. +4. Update EmbodiChain to consume that contract and complete multi-env, + rebuild, articulation, and multi-manager parity. +5. Run the Stage 1 merge gate. +6. Write the dependent Stage 2 implementation plan. +7. Implement dynamics and kinematics sessions, bridge, environment, and + reference tasks. +8. Run the complete differentiable and repository merge gates. + +Stage 1 and Stage 2 remain reviewable as separate commit series even though +the first two previously proposed PR scopes are now one coordinated refactor. + +## 18. Accepted trade-offs + +- Transactional rebuild temporarily consumes more memory than clearing the + old runtime first. +- Explicit bindings and generation checks add types and lifecycle plumbing but + remove unsafe permanent IDs. +- Runtime articulation mutation requires upstream snapshot work rather than an + EmbodiChain-only workaround. +- Differentiable tape-owned buffers use more memory than mutating manager state + in place; this is required for correct backward ownership. +- Functor integration remains deliberately narrow in this iteration. +- Soft-body and cloth runtime mutation is explicitly postponed rather than + approximated with incomplete state preservation. diff --git a/embodichain/data_pipeline/engine/data.py b/embodichain/data_pipeline/engine/data.py index 069140a5e..0b1a590b9 100644 --- a/embodichain/data_pipeline/engine/data.py +++ b/embodichain/data_pipeline/engine/data.py @@ -66,6 +66,19 @@ class OnlineDataWorkerError(RuntimeError): """Fallback error for a worker exception that cannot be reconstructed.""" +def _add_exception_note(error: BaseException, note: str) -> None: + """Attach a PEP 678-style note on every supported Python version.""" + add_note = getattr(error, "add_note", None) + if add_note is not None: + add_note(note) + return + notes = getattr(error, "__notes__", None) + if notes is None: + notes = [] + error.__notes__ = notes + notes.append(note) + + def _forced_shutdown_error() -> OnlineDataWorkerError: """Build the error used when graceful worker durability is unknown.""" return OnlineDataWorkerError( @@ -287,7 +300,7 @@ def _run_sim_worker( env_cfg.max_episode_steps = int(shared_buffer.batch_size[1]) env_cfg.sim_cfg = SimulationManagerCfg( headless=gym_config.get("headless", True), - sim_device=gym_config.get("device", "cpu"), + device=gym_config.get("device", "cpu"), render_cfg=RenderCfg(renderer=gym_config.get("renderer", "hybrid")), gpu_id=gym_config.get("gpu_id", 0), ) @@ -699,8 +712,8 @@ def start(self) -> None: forced_shutdown = self._shutdown_worker() except BaseException as caught_cleanup_error: cleanup_error = caught_cleanup_error - error.add_note( - f"Worker cleanup also failed: {caught_cleanup_error}" + _add_exception_note( + error, f"Worker cleanup also failed: {caught_cleanup_error}" ) else: self._cleanup_complete = True @@ -712,9 +725,10 @@ def start(self) -> None: # primary, but never lose that late durability error. channel_error = self._receive_worker_error() if channel_error is not None and channel_error is not error: - error.add_note( + _add_exception_note( + error, "Worker also failed during cleanup: " - f"{type(channel_error).__name__}: {channel_error}" + f"{type(channel_error).__name__}: {channel_error}", ) if forced_shutdown: @@ -722,7 +736,7 @@ def start(self) -> None: if channel_error is None: self._record_worker_error(durability_error) channel_error = durability_error - error.add_note(str(durability_error)) + _add_exception_note(error, str(durability_error)) if ( stop_requested @@ -1226,12 +1240,12 @@ def stop(self) -> None: self._record_worker_error(durability_error) worker_error = durability_error else: - worker_error.add_note(str(durability_error)) + _add_exception_note(worker_error, str(durability_error)) self._set_state(OnlineDataEngineState.FAILED) self._lifecycle_condition.notify_all() if worker_error is not None: - worker_error.add_note( - f"Worker cleanup also failed: {cleanup_error}" + _add_exception_note( + worker_error, f"Worker cleanup also failed: {cleanup_error}" ) raise worker_error self._worker_error = cleanup_error @@ -1247,7 +1261,7 @@ def stop(self) -> None: self._record_worker_error(durability_error) worker_error = durability_error else: - worker_error.add_note(str(durability_error)) + _add_exception_note(worker_error, str(durability_error)) self._cleanup_complete = True if worker_error is not None: @@ -1276,9 +1290,10 @@ def __exit__( except BaseException as cleanup_error: if exc_value is None: raise - exc_value.add_note( + _add_exception_note( + exc_value, "OnlineDataEngine cleanup also failed: " - f"{type(cleanup_error).__name__}: {cleanup_error}" + f"{type(cleanup_error).__name__}: {cleanup_error}", ) return None diff --git a/embodichain/gen_sim/gradio_ui/app_articraft.py b/embodichain/gen_sim/gradio_ui/app_articraft.py index f39653df5..dfa6b92e2 100644 --- a/embodichain/gen_sim/gradio_ui/app_articraft.py +++ b/embodichain/gen_sim/gradio_ui/app_articraft.py @@ -742,7 +742,8 @@ def _start_remote_viser_preview(session_id: str, artifact: Path) -> str: str(artifact.resolve()), "--asset_type", "articulation", - "--use_usd_properties", + "--asset-physics-mode", + "preserve", "--viser", "--viser-host", "0.0.0.0", diff --git a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md index 81700e2db..8671993d3 100644 --- a/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md +++ b/embodichain/gen_sim/gradio_ui/gradio_visualization_architecture.md @@ -134,7 +134,7 @@ description + optional image → articraft view → Gradio iframe ``` -两种方式的本地产物均在 `ARTICRAFT_OUTPUT_ROOT` 下。Remote server 的 USDC sidecar 将单个 assembly 设为 `defaultPrim` 和 articulation root，Gradio 使用当前 Python 环境启动 `embodichain preview-asset --asset_type articulation --use_usd_properties --viser`，并把动态 Viser 端口嵌入页面；Viser 启动失败时仍保留成功 USDC，并回退到结果摘要。Local Codex 继续使用隔离的 Articraft Conda 环境和原生 USDZ Viewer。 +两种方式的本地产物均在 `ARTICRAFT_OUTPUT_ROOT` 下。Remote server 的 USDC sidecar 将单个 assembly 设为 `defaultPrim` 和 articulation root，Gradio 使用当前 Python 环境启动 `embodichain preview-asset --asset_type articulation --asset-physics-mode preserve --viser`，并把动态 Viser 端口嵌入页面；Viser 启动失败时仍保留成功 USDC，并回退到结果摘要。Local Codex 继续使用隔离的 Articraft Conda 环境和原生 USDZ Viewer。 `Reset Articulation` 会清空当前会话的描述、参考图、记录与下载结果，终止该会话的 Articraft 生成、Articraft/Viser Viewer 进程组，并请求取消仍在运行的远程任务。新请求替换旧请求时也执行相同的会话级取消。 diff --git a/embodichain/gen_sim/scene_engine/cli/preview.py b/embodichain/gen_sim/scene_engine/cli/preview.py index f2d19c263..bc78e3ed4 100644 --- a/embodichain/gen_sim/scene_engine/cli/preview.py +++ b/embodichain/gen_sim/scene_engine/cli/preview.py @@ -26,7 +26,7 @@ from typing import Any from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import LightCfg, MeshCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import LightCfg, MeshCfg, MeshCollisionCfg, RigidObjectCfg from embodichain.lab.visualization import ( VisualizationCfg, add_viser_args_to_parser, @@ -79,8 +79,6 @@ def preview_scene_export( ) ) try: - if sim.is_use_gpu_physics: - sim.init_gpu_physics() _add_lights(sim) _add_objects( sim=sim, @@ -94,6 +92,7 @@ def preview_scene_export( config_dir=config_path.parent, label="asset", ) + sim.prepare() is_viser = sim.sim_config.visualization.backend == "viser" if headless and not is_viser: @@ -185,19 +184,29 @@ def _add_objects( field_name=f"{uid}.body_scale", ) max_convex_hull_num = max(1, int(entry.get("max_convex_hull_num", 32))) + mesh_collision = MeshCollisionCfg(approximation="convex_hull") + if max_convex_hull_num > 1: + # The exported preview schema still carries the legacy hull budget; + # normalize it at this input boundary into the explicit Lab schema. + mesh_collision = MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=max_convex_hull_num, + acd_method="coacd", + ) sim.add_rigid_object( RigidObjectCfg( uid=uid, - shape=MeshCfg(fpath=str(mesh_path)), + shape=MeshCfg( + fpath=str(mesh_path), + collision=mesh_collision, + ), # Keep every preview body static: exported poses are already the # final gravity-settled poses and should not be simulated again. body_type="static", init_pos=tuple(init_pos), init_rot=tuple(init_rot), body_scale=tuple(body_scale), - max_convex_hull_num=max_convex_hull_num, - acd_method="vhacd", # Use vhacd by default. ) ) print(f"[{label}] {uid}: pos={init_pos} rot={init_rot} scale={body_scale}") diff --git a/embodichain/gen_sim/scene_engine/core/scene_object.py b/embodichain/gen_sim/scene_engine/core/scene_object.py index 2b868e3c3..d767ab20a 100644 --- a/embodichain/gen_sim/scene_engine/core/scene_object.py +++ b/embodichain/gen_sim/scene_engine/core/scene_object.py @@ -26,7 +26,7 @@ class ObjectPhysics: """Physics and collision settings shared by settling and scene export.""" body_type: Literal["dynamic", "kinematic"] # Runtime behaviour in simulation. - attrs: dict[str, float | int] # Rigid-body material and contact attributes. + attrs: dict[str, object] # Grouped rigid-body physics configuration. max_convex_hull_num: int # Collision-decomposition hull budget. def __post_init__(self) -> None: @@ -37,11 +37,8 @@ def __post_init__(self) -> None: raise ValueError("max_convex_hull_num must be positive.") if not self.attrs: raise ValueError("attrs must contain at least one physics attribute.") - if not all( - isinstance(name, str) and isinstance(value, (float, int)) - for name, value in self.attrs.items() - ): - raise ValueError("attrs must map strings to numeric physics values.") + if not all(isinstance(name, str) for name in self.attrs): + raise ValueError("attrs must use string configuration keys.") def to_dict(self) -> dict[str, object]: """Serialize the physics settings for scene debugging artifacts.""" diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py index b5078c75f..2799cf170 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/gravity_settler.py @@ -32,8 +32,8 @@ transform_matrix_to_layout_object, ) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg, RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.utils.logger import log_info @@ -181,6 +181,7 @@ def settle(self) -> dict[str, dict[str, list[float]]]: "dynamic" if asset_id in self.dynamic_asset_ids else "kinematic" ), ) + sim.prepare() sim.update(step=self.config.settle_steps) settled_pose_by_id: dict[str, dict[str, list[float]]] = {} @@ -273,7 +274,10 @@ def _add_sim_body( return sim.add_rigid_object( RigidObjectCfg( uid=object_id, - shape=MeshCfg(fpath=str(body_info["mesh_path"])), + shape=MeshCfg( + fpath=str(body_info["mesh_path"]), + collision=self._mesh_collision_cfg(physics), + ), init_pos=tuple( self._three_floats(rigid_layout.get("pos"), field_name="pos") ), @@ -283,24 +287,28 @@ def _add_sim_body( ), attrs=self._rigid_body_attrs(physics), body_type=body_type, - max_convex_hull_num=self._max_convex_hull_num(physics), - acd_method="vhacd", ) ) @staticmethod - def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyAttributesCfg: + def _rigid_body_attrs(physics: ObjectPhysics | None) -> RigidBodyPhysicsCfg: """Convert persisted collision material data into one Lab config.""" if physics is None: raise ValueError("Gravity settling requires SimReady physics settings.") - return RigidBodyAttributesCfg(**physics.attrs) + return RigidBodyPhysicsCfg.from_dict(physics.attrs) @staticmethod - def _max_convex_hull_num(physics: ObjectPhysics | None) -> int: - """Read the persisted collision-hull budget after validating physics.""" + def _mesh_collision_cfg(physics: ObjectPhysics | None) -> MeshCollisionCfg: + """Normalize the persisted legacy hull budget into the Lab schema.""" if physics is None: raise ValueError("Gravity settling requires SimReady physics settings.") - return physics.max_convex_hull_num + if physics.max_convex_hull_num == 1: + return MeshCollisionCfg(approximation="convex_hull") + return MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=physics.max_convex_hull_num, + acd_method="coacd", + ) @staticmethod def _require_body_layout_id(body: GravitySettleBody, *, name: str) -> str: diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py index e730b88cb..2dd082853 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/scene_importer.py @@ -312,7 +312,9 @@ def _scene_object_from_export_entry( support_optimization_rect_xy=support_optimization_rect_xy, physics=ObjectPhysics( body_type=str(entry.get("body_type", "dynamic")), # type: ignore[arg-type] - attrs=self._physics_attrs(entry.get("attrs", {"mass": 1.0})), + attrs=self._physics_attrs( + entry.get("attrs", {"mass_props": {"mass": 1.0}}) + ), max_convex_hull_num=max(1, int(entry.get("max_convex_hull_num", 32))), ), ) @@ -400,16 +402,14 @@ def _points2(cls, value: object, *, field_name: str) -> list[list[float]] | None ] @staticmethod - def _physics_attrs(value: object) -> dict[str, float | int]: + def _physics_attrs(value: object) -> dict[str, object]: """Validate exported physics attributes.""" if not isinstance(value, dict) or not value: raise ValueError("Scene object attrs must be a non-empty object.") - attrs: dict[str, float | int] = {} - for key, item in value.items(): - if not isinstance(key, str) or not isinstance(item, (float, int)): - raise ValueError("Scene object attrs must map strings to numbers.") - attrs[key] = item - return attrs + from embodichain.lab.sim.cfg.rigid import _rigid_body_physics_from_dict + + _rigid_body_physics_from_dict(value) + return dict(value) def import_scene_from_output_root(output_root: str | Path) -> Scene: diff --git a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py index 864e39a59..977920452 100644 --- a/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py +++ b/embodichain/gen_sim/scene_engine/pipeline/utils/simready_processor.py @@ -48,21 +48,34 @@ from embodichain.utils.logger import log_info _TABLE_PHYSICS_ATTRS = { - "mass": 10.0, # Keep the table heavy if a simulator treats it as movable. - "static_friction": 0.95, # Resist lateral sliding at table contacts. - "dynamic_friction": 0.9, # Maintain high friction during sliding contacts. - "restitution": 0.01, # Prevent a table contact from producing visible bounce. + "mass_props": { + "mass": 10.0, # Keep the table heavy if a simulator treats it as movable. + }, + "material_props": { + "static_friction": 0.95, # Resist lateral sliding at table contacts. + "dynamic_friction": 0.9, # Maintain high friction during sliding contacts. + "restitution": 0.01, # Prevent a table contact from producing visible bounce. + }, } _ASSET_PHYSICS_ATTRS = { - "mass": 0.01, # Use a lightweight default for unconstrained generated assets. - "contact_offset": 0.003, # Start contact detection slightly before mesh contact. - "rest_offset": 0.001, # Keep a small stable separation after contact resolution. - "restitution": 0.01, # Prevent generated assets from bouncing on the table. - "max_depenetration_velocity": 10.0, # Cap corrective separation speed. - "min_position_iters": 32, # Use extra position iterations for stable contacts. - "min_velocity_iters": 8, # Use extra velocity iterations for stable contacts. + "mass_props": { + "mass": 0.01, # Use a lightweight default for unconstrained generated assets. + }, + "collision_props": { + "contact_offset": 0.003, # Start contact detection slightly before mesh contact. + "rest_offset": 0.001, # Keep a small stable separation after contact resolution. + }, + "material_props": { + "restitution": 0.01, # Prevent generated assets from bouncing on the table. + }, + "rigid_props": { + "backend": "default", + "max_depenetration_velocity": 10.0, # Cap corrective separation speed. + "min_position_iters": 32, # Use extra position iterations for stable contacts. + "min_velocity_iters": 8, # Use extra velocity iterations for stable contacts. + }, } -_FIXED_MAX_CONVEX_HULL_NUM = 16 # Shared VHACD hull budget for settling and export. +_FIXED_MAX_CONVEX_HULL_NUM = 16 # Shared decomposition hull budget for settling/export. @dataclass(frozen=True) diff --git a/embodichain/lab/gym/envs/base_env.py b/embodichain/lab/gym/envs/base_env.py index 8e0c6207b..bdd9ac55b 100644 --- a/embodichain/lab/gym/envs/base_env.py +++ b/embodichain/lab/gym/envs/base_env.py @@ -160,21 +160,48 @@ def __init__( self._configure_timing() + # Phase 1 only declares scene topology. Spawn-backed assets intentionally + # remain metadata-light until the single prepare boundary below. self._setup_scene(**kwargs) # Keep the established env._profiler API while sharing the single # profiler instance owned by SimulationManager. self._profiler = self.sim.profiler - # TODO: To be removed. - if self.device.type == "cuda": - self.sim.init_gpu_physics() + # Materialize every physical declaration in one transaction. DexSim's + # articulation adapter parses each source while finalizing, then the + # resulting handles bind the existing EmbodiChain facades in place. + self.sim.prepare() + + # Phase 2 may now consume link/joint metadata, construct action spaces, + # and create render-only resources such as CameraGroup instances. + configured_robot = self._setup_robot(**kwargs) + if configured_robot is not None: + self.robot = configured_robot + + if self.robot is None: + logger.log_error( + f"The robot instance must be initialized in :meth:`_setup_robot` function." + ) + if len(self.active_joint_ids) == 0: + self.active_joint_ids = self.robot.active_joint_ids + if self.single_action_space is None: + logger.log_error( + f":attr:`single_action_space` must be defined in the :meth:`_setup_robot` function." + ) + + self.sensors = self._setup_sensors(**kwargs) + self._camera_group_ids = [ + sensor.group_id + for sensor in self.sensors.values() + if isinstance(sensor, Camera) + ] if not self.sim_cfg.headless: self.sim.open_window() 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. @@ -482,46 +509,47 @@ def add_camera_group_id(self, group_id: int) -> None: self._camera_group_ids.append(group_id) def _setup_scene(self, **kwargs): - # Init sim manager. - # we want to open gui window when the scene is setup, so init sim manager in headless mode first. + """Declare physical scene topology without consuming runtime metadata.""" + # Init sim manager. We want to open the GUI window after the scene is + # materialized, so construct the manager in headless mode first. headless = self.sim_cfg.headless self.sim_cfg.headless = True self.sim = SimulationManager(self.sim_cfg) 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) - if len(self.active_joint_ids) == 0: - self.active_joint_ids = self.robot.active_joint_ids - - if self.robot is None: - logger.log_error( - f"The robot instance must be initialized in :meth:`_setup_robot` function." - ) - if self.single_action_space is None: - logger.log_error( - f":attr:`single_action_space` must be defined in the :meth:`_setup_robot` function." - ) + # Config-driven environments can declare their robot here while + # deferring all link/joint queries until the post-prepare phase. Generic + # BaseEnv subclasses may keep returning None and add a runtime robot in + # _setup_robot() for backwards compatibility. + self.robot = self._declare_robot(**kwargs) self._prepare_scene(**kwargs) - self.sensors = self._setup_sensors(**kwargs) + def _declare_robot(self, **kwargs) -> Robot | None: + """Optionally declare a robot before the scene prepare boundary. + + Config-driven environments should override this hook and call + :meth:`SimulationManager.add_robot` without querying link/joint data. + The returned facade is bound in place by :meth:`SimulationManager.prepare`. - # Setup camera groups for rendering. - self._camera_group_ids: List[int] = [] - for sensor in self.sensors.values(): - if isinstance(sensor, Camera): - self._camera_group_ids.append(sensor.group_id) + Generic subclasses that only implement the historical + :meth:`_setup_robot` hook remain supported: their robot is added after + the initial prepare boundary and is prepared immediately by the manager. + """ + del kwargs + return None def _setup_robot(self, **kwargs) -> Robot: - """Load the robot agent, setup the controller and action space. + """Configure the bound robot, controller, and action space. Note: - 1. The fuction must return the robot instance. - 2. The self.single_action_space should be defined. + This hook runs after :meth:`SimulationManager.prepare`, so link, + joint, and limit metadata are available. It must return the robot + instance and define ``self.single_action_space``. """ # TODO: single_action_space may be configured in config? diff --git a/embodichain/lab/gym/envs/differentiable_env.py b/embodichain/lab/gym/envs/differentiable_env.py new file mode 100644 index 000000000..95dfbdad0 --- /dev/null +++ b/embodichain/lab/gym/envs/differentiable_env.py @@ -0,0 +1,250 @@ +# ---------------------------------------------------------------------------- +# 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_dynamics_action_kernel(self, action_wp, control, tape): ... + def _read_outputs(self, final_state) -> dict: ... +""" + +from __future__ import annotations + +from typing import Any, Callable, Literal + +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`. + + Dynamics subclasses must implement :meth:`_apply_dynamics_action_kernel` + and :meth:`_read_outputs`; the rest of the EmbodiedEnv contract (reset, + observation managers, reward functors) carries over. The default + ``dynamics`` route invokes the Newton solver through + :class:`NewtonStepFunc` using a detached trajectory-local control buffer. + Subclasses that intentionally use FK-only stepping must explicitly select + ``kinematics`` and implement :meth:`_make_kinematic_step_fn` together with + the legacy :meth:`_apply_action_kernel` hook. + """ + + differentiable_step_mode: Literal["dynamics", "kinematics"] = "dynamics" + """Stepping route used by :meth:`_build_sim_state_dict`.""" + + 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_dynamics_action_kernel( + self, + action_wp: Any, + control: Any, + tape: Any, + ) -> None: + """Write an action into a detached dynamics trajectory control buffer. + + 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 the supplied ``control``. + It is the isolated control owned by the active manager trajectory; do + not write ``self.sim.physics.newton_manager._control`` while the tape + is active. ``tape`` is the caller-owned active Warp tape for this + callback only; the bridge clears the per-step binding after tape exit. + """ + raise NotImplementedError( + "Dynamics subclasses of DifferentiableEmbodiedEnv must migrate " + "their legacy _apply_action_kernel(action_wp, tape) hook to " + "_apply_dynamics_action_kernel(action_wp, control, tape)." + ) + + def _apply_action_kernel(self, action_wp: Any, tape: Any) -> None: + """Write an action for the explicitly selected kinematics route. + + This legacy hook is deliberately reserved for + ``differentiable_step_mode = 'kinematics'``. It receives no detached + solver control because FK-only environments do not invoke Newton + solver dynamics. + """ + raise NotImplementedError( + "Kinematics 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_kinematic_step_fn(self) -> Callable[[], Any]: + """Return the explicitly selected FK-only stepping callback. + + Subclasses must override this hook only when they set + :attr:`differentiable_step_mode` to ``"kinematics"``. This keeps + kinematics distinct from the default solver-dynamics route. + + Raises: + NotImplementedError: If kinematics mode has no named FK hook. + """ + raise NotImplementedError( + "DifferentiableEmbodiedEnv in kinematics mode requires " + "_make_kinematic_step_fn()." + ) + + # -- gym surface ------------------------------------------------------ # + + def step(self, action: torch.Tensor): + """Advance one differentiable control step. + + Terminal environments are auto-reset only when the call cannot retain + a Warp tape for backward. A grad-tracked step returns terminal + observations unchanged and records ``deferred_reset_ids`` in ``info``; + callers must run backward before resetting those environments. + """ + if not isinstance(action, torch.Tensor): + action = torch.as_tensor(action, dtype=torch.float32) + retains_tape_for_backward = bool( + torch.is_grad_enabled() and action.requires_grad + ) + 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) + if retains_tape_for_backward: + info["requires_reset_after_backward"] = True + info["deferred_reset_ids"] = reset_ids.detach().clone() + else: + 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: + mode = self.differentiable_step_mode + if mode not in {"dynamics", "kinematics"}: + raise ValueError( + "differentiable_step_mode must be 'dynamics' or 'kinematics', " + f"got {mode!r}." + ) + + action_kernel, tape_binder = self._action_kernel_for_mode(mode) + sim_state = { + "manager": self.sim, + "step_mode": mode, + "substeps": self.cfg.sim_steps_per_control, + "action_to_control_kernel": action_kernel, + "kernel_args": (), + "obs_reward_fn": self._read_outputs, + "last_info": {}, + } + if tape_binder is not None: + sim_state["_bind_dynamics_tape"] = tape_binder + if mode == "kinematics": + sim_state["step_fn"] = self._make_kinematic_step_fn() + return sim_state + + def _action_kernel_for_mode( + self, + mode: str, + ) -> tuple[Callable[..., None], Callable[[Any | None], None] | None]: + """Build the mode-specific action callback consumed by NewtonStepFunc.""" + if mode == "dynamics": + dynamics_hook = getattr(self, "_apply_dynamics_action_kernel", None) + if ( + not callable(dynamics_hook) + or getattr(dynamics_hook, "__func__", None) + is DifferentiableEmbodiedEnv._apply_dynamics_action_kernel + ): + raise NotImplementedError( + "Dynamics environments using the legacy " + "_apply_action_kernel(action_wp, tape) must migrate to " + "_apply_dynamics_action_kernel(action_wp, control, tape)." + ) + return self._wrap_dynamics_action_kernel(dynamics_hook) + return self._wrap_kinematic_action_kernel(), None + + @staticmethod + def _wrap_dynamics_action_kernel( + dynamics_hook: Callable[..., None], + ) -> tuple[Callable[..., None], Callable[[Any | None], None]]: + """Expose a local-control hook with tape ownership scoped per step.""" + active_tape: list[Any | None] = [None] + + def _bind_tape(tape: Any | None) -> None: + active_tape[0] = tape + + def _inner(action_wp: Any, control: Any, *_: Any) -> None: + dynamics_hook(action_wp, control, tape=active_tape[0]) + + return _inner, _bind_tape + + def _wrap_kinematic_action_kernel(self): + """Expose the strict legacy action hook only for kinematics mode.""" + env = self + + def _inner(action_wp: Any, tape: Any, *_: Any) -> None: + env._apply_action_kernel(action_wp, tape=tape) + + return _inner diff --git a/embodichain/lab/gym/envs/embodied_env.py b/embodichain/lab/gym/envs/embodied_env.py index ba6fb168c..20e50c584 100644 --- a/embodichain/lab/gym/envs/embodied_env.py +++ b/embodichain/lab/gym/envs/embodied_env.py @@ -42,6 +42,7 @@ from embodichain.lab.sim.cfg import ( RobotCfg, + RobotPresetCfg, RigidObjectCfg, RigidObjectGroupCfg, ArticulationCfg, @@ -109,8 +110,9 @@ class EmbodiedEnvCfg(EnvCfg): instance as attributes during initialization. Key fields - - **robot**: `RobotCfg` (required) — the agent definition (URDF/MJCF, initial - state, control mode, etc.). + - **robot**: `RobotCfg | RobotPresetCfg` (required) — one portable robot + definition or replace-only complete alternatives selected by the active + physics backend. - **control_parts**: Optional[List[str]] — named robot parts to control. If `None`, all controllable joints are used. - **active_joint_ids**: List[int] — explicit joint indices to use for @@ -148,7 +150,7 @@ class EnvLightCfg: # TODO: support more types of indirect light in the future. indirect: dict[str, Any] | None = None - robot: RobotCfg = MISSING + robot: RobotCfg | RobotPresetCfg = MISSING control_parts: list[str] | None = None """List of robot parts to control. If None, all controllable joints will be used. @@ -853,9 +855,9 @@ def _extend_reward( return rewards def _prepare_scene(self, **kwargs) -> None: - self._setup_lights() self._setup_background() self._setup_interactive_objects() + self._setup_lights() def _update_sim_state(self, **kwargs) -> None: """Perform the simulation step and apply events if configured. @@ -1820,8 +1822,15 @@ def _postprocess_action(self, action): return self.action_manager.process_action(action, mode="post") return super()._postprocess_action(action) + def _declare_robot(self, **kwargs) -> Robot: + """Declare the configured robot without reading articulation metadata.""" + del kwargs + if self.cfg.robot is None: + logger.log_error("Robot configuration is not provided.") + return self.sim.add_robot(self.cfg.robot) + def _setup_robot(self, **kwargs) -> Robot: - """Setup the robot in the environment. + """Configure the finalized robot interface for the environment. Currently, only joint position control is supported. Would be extended to support joint velocity and torque control in the future. @@ -1829,11 +1838,10 @@ def _setup_robot(self, **kwargs) -> Robot: Returns: Robot: The robot instance added to the scene. """ - if self.cfg.robot is None: - logger.log_error("Robot configuration is not provided.") - - # Initialize the robot based on the configuration. - robot: Robot = self.sim.add_robot(self.cfg.robot) + del kwargs + robot = self.robot + if robot is None: + logger.log_error("Robot was not declared before simulation prepare.") # Setup active joints for robot to control. if self.cfg.control_parts: diff --git a/embodichain/lab/gym/envs/expert_program/cfg.py b/embodichain/lab/gym/envs/expert_program/cfg.py index a757f7a55..098894f50 100644 --- a/embodichain/lab/gym/envs/expert_program/cfg.py +++ b/embodichain/lab/gym/envs/expert_program/cfg.py @@ -231,33 +231,33 @@ def __post_init__(self) -> None: @configclass class PoseCfg: - """One declarative Cartesian pose using a WXYZ quaternion.""" + """One declarative Cartesian pose using an XYZW quaternion.""" position: tuple[float, float, float] = MISSING - quaternion_wxyz: tuple[float, float, float, float] = MISSING + quaternion_xyzw: tuple[float, float, float, float] = MISSING def __post_init__(self) -> None: """Validate pose shape, finiteness, and quaternion magnitude.""" if type(self.position) not in (list, tuple) or len(self.position) != 3: raise ValueError("position must contain exactly three numbers.") if ( - type(self.quaternion_wxyz) not in (list, tuple) - or len(self.quaternion_wxyz) != 4 + type(self.quaternion_xyzw) not in (list, tuple) + or len(self.quaternion_xyzw) != 4 ): - raise ValueError("quaternion_wxyz must contain exactly four numbers.") + raise ValueError("quaternion_xyzw must contain exactly four numbers.") position = tuple( _validate_number(value, field_name=f"position[{index}]") for index, value in enumerate(self.position) ) quaternion = tuple( - _validate_number(value, field_name=f"quaternion_wxyz[{index}]") - for index, value in enumerate(self.quaternion_wxyz) + _validate_number(value, field_name=f"quaternion_xyzw[{index}]") + for index, value in enumerate(self.quaternion_xyzw) ) norm = math.sqrt(sum(value * value for value in quaternion)) if norm <= 1.0e-12: - raise ValueError("quaternion_wxyz must have non-zero magnitude.") + raise ValueError("quaternion_xyzw must have non-zero magnitude.") self.position = position # type: ignore[assignment] - self.quaternion_wxyz = quaternion # type: ignore[assignment] + self.quaternion_xyzw = quaternion # type: ignore[assignment] @configclass diff --git a/embodichain/lab/gym/envs/expert_program/compiler.py b/embodichain/lab/gym/envs/expert_program/compiler.py index e05aa572e..5f66d4a6b 100644 --- a/embodichain/lab/gym/envs/expert_program/compiler.py +++ b/embodichain/lab/gym/envs/expert_program/compiler.py @@ -1675,7 +1675,7 @@ def _compile_targets( (*path, "values", index), "Target values must be exact PoseCfg values.", ) - poses.append(SemanticPose(pose.position, pose.quaternion_wxyz)) + poses.append(SemanticPose(pose.position, pose.quaternion_xyzw)) compiled[target_id] = tuple(poses) return MappingProxyType(compiled) diff --git a/embodichain/lab/gym/envs/expert_program/configured_runtime.py b/embodichain/lab/gym/envs/expert_program/configured_runtime.py index 078d2e036..c468a3331 100644 --- a/embodichain/lab/gym/envs/expert_program/configured_runtime.py +++ b/embodichain/lab/gym/envs/expert_program/configured_runtime.py @@ -1580,7 +1580,7 @@ def _decode_handover_pose_provider( { "kind", "final_position", - "final_quaternion_wxyz", + "final_quaternion_xyzw", } ), ) @@ -1596,9 +1596,9 @@ def _decode_handover_pose_provider( path=f"{path}.final_position", expected_length=3, ), - final_quaternion_wxyz=_finite_tuple( - config["final_quaternion_wxyz"], - path=f"{path}.final_quaternion_wxyz", + final_quaternion_xyzw=_finite_tuple( + config["final_quaternion_xyzw"], + path=f"{path}.final_quaternion_xyzw", expected_length=4, ), ) diff --git a/embodichain/lab/gym/envs/expert_program/decoder.py b/embodichain/lab/gym/envs/expert_program/decoder.py index fec6dfaef..0875eab1e 100644 --- a/embodichain/lab/gym/envs/expert_program/decoder.py +++ b/embodichain/lab/gym/envs/expert_program/decoder.py @@ -392,14 +392,14 @@ def _decode_pose(value: object, *, path: ConfigPath) -> PoseCfg: mapping = _expect_mapping(value, path=path) _validate_fields( mapping, - allowed=frozenset({"position", "quaternion_wxyz"}), - required=frozenset({"position", "quaternion_wxyz"}), + allowed=frozenset({"position", "quaternion_xyzw"}), + required=frozenset({"position", "quaternion_xyzw"}), path=path, ) position_values = _expect_list(mapping["position"], path=(*path, "position")) quaternion_values = _expect_list( - mapping["quaternion_wxyz"], - path=(*path, "quaternion_wxyz"), + mapping["quaternion_xyzw"], + path=(*path, "quaternion_xyzw"), ) if len(position_values) != 3: raise _error( @@ -410,12 +410,12 @@ def _decode_pose(value: object, *, path: ConfigPath) -> PoseCfg: if len(quaternion_values) != 4: raise _error( "invalid_pose_shape", - (*path, "quaternion_wxyz"), - "quaternion_wxyz must contain exactly four numbers.", + (*path, "quaternion_xyzw"), + "quaternion_xyzw must contain exactly four numbers.", ) for name, values in ( ("position", position_values), - ("quaternion_wxyz", quaternion_values), + ("quaternion_xyzw", quaternion_values), ): for index, number in enumerate(values): if type(number) not in (int, float): @@ -428,7 +428,7 @@ def _decode_pose(value: object, *, path: ConfigPath) -> PoseCfg: PoseCfg, path=path, position=tuple(position_values), - quaternion_wxyz=tuple(quaternion_values), + quaternion_xyzw=tuple(quaternion_values), ) # type: ignore[return-value] diff --git a/embodichain/lab/gym/envs/expert_program/simulation_handover.py b/embodichain/lab/gym/envs/expert_program/simulation_handover.py index 64b494746..be98a49d1 100644 --- a/embodichain/lab/gym/envs/expert_program/simulation_handover.py +++ b/embodichain/lab/gym/envs/expert_program/simulation_handover.py @@ -35,19 +35,19 @@ def _validated_pose( position: tuple[float, float, float], - quaternion_wxyz: tuple[float, float, float, float], + quaternion_xyzw: tuple[float, float, float, float], *, field_name: str, ) -> SemanticPose: """Build and validate one unbatched semantic pose declaration.""" if type(position) is not tuple or len(position) != 3: raise TypeError(f"{field_name}_position must be an exact 3-tuple.") - if type(quaternion_wxyz) is not tuple or len(quaternion_wxyz) != 4: - raise TypeError(f"{field_name}_quaternion_wxyz must be an exact 4-tuple.") + if type(quaternion_xyzw) is not tuple or len(quaternion_xyzw) != 4: + raise TypeError(f"{field_name}_quaternion_xyzw must be an exact 4-tuple.") try: return SemanticPose( position=position, - quaternion_wxyz=quaternion_wxyz, + quaternion_xyzw=quaternion_xyzw, ) except (TypeError, ValueError) as exc: raise type(exc)(f"Invalid {field_name} hand-over pose: {exc}") from exc @@ -65,18 +65,18 @@ class ConfiguredHandOverPoseProvider(HandOverPoseProvider): Args: final_position: World-frame object delivery position. - final_quaternion_wxyz: World-frame object delivery orientation. + final_quaternion_xyzw: World-frame object delivery orientation. """ provider_id: ClassVar[str] = "simulation.configured_handover_pose" final_position: tuple[float, float, float] - final_quaternion_wxyz: tuple[float, float, float, float] + final_quaternion_xyzw: tuple[float, float, float, float] def __post_init__(self) -> None: final = _validated_pose( self.final_position, - self.final_quaternion_wxyz, + self.final_quaternion_xyzw, field_name="final", ) object.__setattr__( @@ -86,8 +86,8 @@ def __post_init__(self) -> None: ) object.__setattr__( self, - "final_quaternion_wxyz", - tuple(float(value) for value in final.quaternion_wxyz.tolist()), + "final_quaternion_xyzw", + tuple(float(value) for value in final.quaternion_xyzw.tolist()), ) def resolve( @@ -112,7 +112,7 @@ def resolve( final=SemanticObjectTarget( pose=SemanticPose( position=self.final_position, - quaternion_wxyz=self.final_quaternion_wxyz, + quaternion_xyzw=self.final_quaternion_xyzw, ) ), ) diff --git a/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py b/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py index 458401b6c..48ff76c4b 100644 --- a/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py +++ b/embodichain/lab/gym/envs/managers/_event_functors/dynamic_settling.py @@ -124,8 +124,9 @@ def _get_dynamic_entity_catalog( def _is_dynamic_entity(kind: str, entity: _DynamicEntity) -> bool: """Return whether an entity participates in dynamic physics. - Articulation links are physics-backed even when ``fix_base`` constrains the - root link, so every non-robot articulation is a valid settle target. + Articulation links are physics-backed even when + ``root_props.fixed_base`` constrains the root link, so every + non-robot articulation is a valid settle target. """ if kind == "articulation": return True diff --git a/embodichain/lab/gym/envs/managers/actions.py b/embodichain/lab/gym/envs/managers/actions.py index 4f9430f9b..6882f998d 100644 --- a/embodichain/lab/gym/envs/managers/actions.py +++ b/embodichain/lab/gym/envs/managers/actions.py @@ -232,7 +232,7 @@ class EefPoseTerm(ActionTerm): Supports two pose representations: - 6D: position (3) + Euler angles (3) - - 7D: position (3) + quaternion (4) + - 7D: position (3) + quaternion in ``xyzw`` order (4) On IK failure, falls back to current_qpos for that env. Returns ``ik_success`` in the TensorDict so reward/observation @@ -248,7 +248,7 @@ class EefPoseTerm(ActionTerm): >>> # 7D: position (3) + quaternion (4) >>> action = torch.zeros(num_envs, 7) >>> action[:, :3] = 0.1 # target position - >>> action[:, 3] = 1.0 # quaternion w + >>> action[:, 6] = 1.0 # quaternion w (xyzw identity) >>> result = term.process_action(action) >>> # result["qpos"] = IK solution >>> # result["ik_success"] = bool tensor indicating IK success diff --git a/embodichain/lab/gym/envs/managers/events.py b/embodichain/lab/gym/envs/managers/events.py index afa1173e7..4c2ebfbbc 100644 --- a/embodichain/lab/gym/envs/managers/events.py +++ b/embodichain/lab/gym/envs/managers/events.py @@ -609,7 +609,7 @@ def drop_rigid_object_group_sequentially( .repeat(num_instance, 1) ) drop_pose = torch.zeros((num_instance, 7), device=env.device) - drop_pose[:, 3] = 1.0 # w component of quaternion + drop_pose[:, 6] = 1.0 # w component of xyzw quaternion drop_pose[:, :3] = drop_pos for i in range(num_objects): random_offset = sample_uniform( diff --git a/embodichain/lab/gym/envs/managers/observations.py b/embodichain/lab/gym/envs/managers/observations.py index 3729ec8d0..45374c9ef 100644 --- a/embodichain/lab/gym/envs/managers/observations.py +++ b/embodichain/lab/gym/envs/managers/observations.py @@ -50,7 +50,8 @@ def get_object_pose( env: The environment instance. obs: The observation dictionary. entity_cfg: The configuration of the scene entity. - to_matrix: Whether to return the pose as a 4x4 transformation matrix. If False, returns as (position, quaternion). + to_matrix: Whether to return the pose as a 4x4 transformation matrix. If + False, returns ``(x, y, z, qx, qy, qz, qw)``. Returns: A tensor of shape (num_envs, 7) or (num_envs, 4, 4) representing the world poses of the objects. @@ -90,7 +91,8 @@ def get_rigid_object_pose( env: The environment instance. obs: The observation dictionary. entity_cfg: The configuration of the scene entity. - to_matrix: Whether to return the pose as a 4x4 transformation matrix. If False, returns as (position, quaternion). + to_matrix: Whether to return the pose as a 4x4 transformation matrix. If + False, returns ``(x, y, z, qx, qy, qz, qw)``. Returns: A tensor of shape (num_envs, 7) or (num_envs, 4, 4) representing the world poses of the rigid objects. @@ -1157,14 +1159,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/managers/randomization/physics.py b/embodichain/lab/gym/envs/managers/randomization/physics.py index 1eea74e04..a0a7da9a7 100644 --- a/embodichain/lab/gym/envs/managers/randomization/physics.py +++ b/embodichain/lab/gym/envs/managers/randomization/physics.py @@ -35,6 +35,8 @@ def randomize_rigid_object_mass( entity_cfg: SceneEntityCfg, mass_range: tuple[float, float], relative: bool = False, + recompute_inertia: bool = True, + min_mass: float = 1e-6, ) -> None: """Randomize the mass of rigid objects in the environment. @@ -44,25 +46,54 @@ def randomize_rigid_object_mass( entity_cfg (SceneEntityCfg): The configuration for the scene entity. mass_range (tuple[float, float]): The range (min, max) to sample the mass from. relative (bool): Whether to apply the mass change relative to the initial mass. Defaults to False. + recompute_inertia (bool): Whether to scale the initial inertia by the sampled + mass ratio. Defaults to True. + min_mass (float): Minimum allowed sampled mass. Defaults to 1e-6. + + Raises: + ValueError: If ``min_mass`` is not positive or an initial mass is not positive. """ if entity_cfg.uid not in env.sim.get_rigid_object_uid_list(): return rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid) + if rigid_object.is_non_dynamic: + logger.log_warning( + f"Cannot randomize mass for non-dynamic rigid object '{entity_cfg.uid}'." + ) + return + if min_mass <= 0.0: + raise ValueError(f"min_mass must be positive, got {min_mass}.") + num_instance = len(env_ids) + index = torch.as_tensor(env_ids, dtype=torch.long, device=rigid_object.device) + body_data = rigid_object.body_data + if body_data is None: + return + default_masses = body_data.default_mass[index] + if torch.any(default_masses <= 0.0): + raise ValueError("Initial rigid-body masses must be positive.") sampled_masses = sample_uniform( - lower=mass_range[0], upper=mass_range[1], size=(num_instance,) + lower=mass_range[0], + upper=mass_range[1], + size=(num_instance,), + device=rigid_object.device, ) if relative: - init_mass = rigid_object.cfg.attrs.mass - init_mass = torch.full((sampled_masses.shape), init_mass, device=env.device) - sampled_masses = init_mass + sampled_masses + sampled_masses = default_masses + sampled_masses + + sampled_masses = sampled_masses.clamp_min(min_mass) rigid_object.set_mass(sampled_masses, env_ids=env_ids) + if recompute_inertia: + mass_ratios = sampled_masses / default_masses + sampled_inertia = body_data.default_inertia[index] * mass_ratios.unsqueeze(-1) + rigid_object.set_inertia(sampled_inertia, env_ids=env_ids) + def randomize_rigid_object_center_of_mass( env: EmbodiedEnv, @@ -111,6 +142,8 @@ def randomize_articulation_mass( mass_range: tuple[float, float] | dict[str, tuple[float, float]], link_names: str | list[str] | None = None, relative: bool = False, + recompute_inertia: bool = True, + min_mass: float = 1e-6, ) -> None: """Randomize the mass of articulation links in the environment. @@ -127,14 +160,23 @@ def randomize_articulation_mass( link_names (str | list[str] | None): A regex pattern or list of regex patterns to match link names. If None, all links are randomized. Ignored when ``mass_range`` is a dict. Defaults to None. - relative (bool): Whether to apply the mass change relative to the current mass. + relative (bool): Whether to apply the mass change relative to the initial mass. Defaults to False. + recompute_inertia (bool): Whether to scale initialization-time inertia by + the sampled mass ratio. Defaults to True. + min_mass (float): Minimum allowed sampled mass. Defaults to 1e-6. + + Raises: + ValueError: If ``min_mass`` or an initialization-time link mass is not + positive. """ if entity_cfg.uid not in env.sim.get_articulation_uid_list(): return articulation: Articulation = env.sim.get_articulation(entity_cfg.uid) + if min_mass <= 0.0: + raise ValueError(f"min_mass must be positive, got {min_mass}.") num_instance = len(env_ids) if isinstance(mass_range, dict): @@ -149,18 +191,18 @@ def randomize_articulation_mass( matched_link_names = list(mass_range.keys()) link_lower = torch.tensor( [mass_range[name][0] for name in matched_link_names], - device=env.device, + device=articulation.device, dtype=torch.float32, ) link_upper = torch.tensor( [mass_range[name][1] for name in matched_link_names], - device=env.device, + device=articulation.device, dtype=torch.float32, ) # Broadcast: (num_instance, num_links) sampled_masses = torch.rand( (num_instance, len(matched_link_names)), - device=env.device, + device=articulation.device, dtype=torch.float32, ) sampled_masses = link_lower + sampled_masses * (link_upper - link_lower) @@ -179,17 +221,39 @@ def randomize_articulation_mass( lower=mass_range[0], upper=mass_range[1], size=(num_instance, len(matched_link_names)), + device=articulation.device, + ) + + env_index = torch.as_tensor(env_ids, dtype=torch.long, device=articulation.device) + link_indices = torch.as_tensor( + [articulation.link_names.index(name) for name in matched_link_names], + dtype=torch.long, + device=articulation.device, + ) + default_masses = articulation.body_data.default_mass[ + env_index[:, None], link_indices[None, :] + ] + if torch.any(default_masses <= 0.0): + raise ValueError( + "Initialization-time articulation link masses must be positive." ) if relative: - link_indices = [ - articulation.link_names.index(name) for name in matched_link_names - ] - current_masses = articulation.default_link_masses.clone()[env_ids][ - :, link_indices - ] - sampled_masses = current_masses + sampled_masses + sampled_masses = default_masses + sampled_masses + + sampled_masses = sampled_masses.clamp_min(min_mass) articulation.set_mass( sampled_masses, link_names=matched_link_names, env_ids=env_ids ) + + if recompute_inertia: + default_inertia = articulation.body_data.default_inertia[ + env_index[:, None], link_indices[None, :] + ] + mass_ratios = sampled_masses / default_masses + articulation.set_inertia( + default_inertia * mass_ratios.unsqueeze(-1), + link_names=matched_link_names, + env_ids=env_ids, + ) diff --git a/embodichain/lab/gym/envs/managers/randomization/spatial.py b/embodichain/lab/gym/envs/managers/randomization/spatial.py index 384490d51..65cd83ea0 100644 --- a/embodichain/lab/gym/envs/managers/randomization/spatial.py +++ b/embodichain/lab/gym/envs/managers/randomization/spatial.py @@ -860,7 +860,7 @@ def _move_object_z( return # Both RigidObject and Articulation return (N, 7) by default: - # (x, y, z, qw, qx, qy, qz) + # (x, y, z, qx, qy, qz, qw) pose = obj.get_local_pose() # (N, 7) current_z = pose[env_ids, 2] if absolute: diff --git a/embodichain/lab/gym/utils/gym_utils.py b/embodichain/lab/gym/utils/gym_utils.py index 8636a3abe..b43661c5c 100644 --- a/embodichain/lab/gym/utils/gym_utils.py +++ b/embodichain/lab/gym/utils/gym_utils.py @@ -433,8 +433,8 @@ def config_to_cfg( RigidObjectGroupCfg, ArticulationCfg, LightCfg, - PhysicsCfg, RenderCfg, + physics_cfg_for_backend, ) from embodichain.lab.sim import SimulationManagerCfg from embodichain.lab.visualization import VisualizationCfg, ViserServerCfg @@ -532,6 +532,8 @@ class ComponentCfg: physics_config = deepcopy(config.get("physics_config", {})) if "gravity" in physics_config: physics_config["gravity"] = np.asarray(physics_config["gravity"]) + physics_cfg = physics_cfg_for_backend(config.get("physics", "default")) + physics_cfg = type(physics_cfg)(**physics_config) render_config = deepcopy(config.get("render_cfg", {})) if "renderer" in config: @@ -551,11 +553,11 @@ class ComponentCfg: env_cfg.sim_cfg = SimulationManagerCfg( headless=config.get("headless", False), - sim_device=config.get("device", "cpu"), + device=config.get("device", "cpu"), render_cfg=RenderCfg(**render_config), gpu_id=config.get("gpu_id", 0), arena_space=config.get("arena_space", 5.0), - physics_config=PhysicsCfg(**physics_config), + physics_cfg=physics_cfg, visualization=VisualizationCfg( **visualization_config, viser_server=viser_server, @@ -915,6 +917,7 @@ def add_env_launcher_args_to_parser( --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) @@ -961,6 +964,15 @@ def add_env_launcher_args_to_parser( "config, the configured render_cfg.renderer is used unless this option " "is provided.", ) + parser.add_argument( + "--physics", + type=str, + choices=["default", "newton"], + default=None if require_gym_config else "default", + help="Physics backend configuration to use for the simulation. When " + "loading a gym config, the configured backend is used unless this " + "option is provided.", + ) parser.add_argument( "--arena_space", help="The size of the arena space.", @@ -1065,6 +1077,8 @@ def merge_args_with_gym_config(args: argparse.Namespace, gym_config: dict) -> di merged_config["headless"] = args.headless or viser_enabled if args.renderer is not None: merged_config["renderer"] = args.renderer + if getattr(args, "physics", None) is not None: + merged_config["physics"] = args.physics merged_config["gpu_id"] = args.gpu_id merged_config["arena_space"] = args.arena_space if args.max_episodes is not None: diff --git a/embodichain/lab/gym/utils/trajectory_state.py b/embodichain/lab/gym/utils/trajectory_state.py index e165646d9..6e531a85e 100644 --- a/embodichain/lab/gym/utils/trajectory_state.py +++ b/embodichain/lab/gym/utils/trajectory_state.py @@ -14,7 +14,11 @@ # limitations under the License. # ---------------------------------------------------------------------------- -"""Shared simulation-state capture and restore helpers for trajectories.""" +"""Shared simulation-state capture and restore helpers for trajectories. + +All seven-element root and rigid-object poses use EmbodiChain's +``(x, y, z, qx, qy, qz, qw)`` convention. +""" from __future__ import annotations diff --git a/embodichain/lab/scripts/analyze_workspace.py b/embodichain/lab/scripts/analyze_workspace.py index 7f1649c14..b49af8a09 100644 --- a/embodichain/lab/scripts/analyze_workspace.py +++ b/embodichain/lab/scripts/analyze_workspace.py @@ -297,7 +297,7 @@ def _build_asset_robot_cfg( ValueError: If ``--ee-link`` is missing, or a USD/non-URDF asset is given without ``--urdf``. """ - from embodichain.lab.sim.cfg import RobotCfg + from embodichain.lab.sim.cfg import ArticulationRootPropertiesCfg, RobotCfg from embodichain.lab.sim.solvers import ( PinkSolverCfg, PinocchioSolverCfg, @@ -344,8 +344,10 @@ def _build_asset_robot_cfg( cfg.fpath = asset cfg.init_pos = tuple(args.init_pos) cfg.init_rot = tuple(args.init_rot) - cfg.fix_base = args.fix_base - cfg.use_usd_properties = args.use_usd_properties + cfg.root_props = ArticulationRootPropertiesCfg( + fixed_base=args.fix_base, + ) + cfg.asset_physics_mode = args.asset_physics_mode cfg.control_parts = {control_part: joints} cfg.solver_cfg = {control_part: solver_cfg} return cfg, control_part, solver_urdf @@ -652,6 +654,7 @@ def main(args: argparse.Namespace) -> None: if robot is None: log_error("Failed to load robot into the simulation.") return + sim.prepare() control_part = _resolve_control_part(robot, control_part) joints_desc = ( robot.control_parts.get(control_part) if control_part else "all joints" @@ -864,10 +867,13 @@ def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: help="Fix the robot base (default: fixed).", ) asset_opts.add_argument( - "--use-usd-properties", - action="store_true", - default=False, - help="Use physical properties from the USD file (USD assets only).", + "--asset-physics-mode", + choices=("preserve", "overlay"), + default="overlay", + help=( + "How asset physics is handled: preserve source-authored values or " + "overlay explicitly configured values (default: overlay for robots)." + ), ) # --- Analysis ----------------------------------------------------------- diff --git a/embodichain/lab/scripts/preview_asset.py b/embodichain/lab/scripts/preview_asset.py index 0bb4f1416..26cfa3d1a 100644 --- a/embodichain/lab/scripts/preview_asset.py +++ b/embodichain/lab/scripts/preview_asset.py @@ -59,6 +59,7 @@ from collections.abc import Sequence 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: @@ -78,14 +79,18 @@ def build_sim_cfg(args: argparse.Namespace) -> SimulationManagerCfg: 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 from embodichain.lab.visualization import visualization_cfg_from_args 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, visualization=visualization_cfg_from_args(args), ) @@ -108,6 +113,7 @@ def load_assets( """ from embodichain.lab.sim.cfg import ( ArticulationCfg, + ArticulationRootPropertiesCfg, LightCfg, RigidObjectCfg, ) @@ -117,6 +123,7 @@ def load_assets( init_pos = tuple(args.init_pos) init_rot = tuple(args.init_rot) spacing = float(args.asset_spacing) + asset_physics_mode = args.asset_physics_mode loaded_assets = [] for idx, asset_path in enumerate(asset_paths): @@ -155,8 +162,10 @@ def load_assets( fpath=asset_path, init_pos=asset_init_pos, init_rot=init_rot, - fix_base=args.fix_base, - use_usd_properties=args.use_usd_properties, + root_props=ArticulationRootPropertiesCfg( + fixed_base=args.fix_base, + ), + asset_physics_mode=asset_physics_mode, # The auxiliary pytorch-kinematics chain only accepts URDF XML. build_pk_chain=asset_suffix not in {".usd", ".usda", ".usdc"}, ) @@ -173,7 +182,7 @@ def load_assets( init_pos=asset_init_pos, init_rot=init_rot, body_type=args.body_type, - use_usd_properties=args.use_usd_properties, + asset_physics_mode=asset_physics_mode, ) loaded_assets.append(sim.add_rigid_object(cfg)) @@ -339,6 +348,7 @@ def main(args: argparse.Namespace) -> None: sim.set_indirect_lighting(args.env_map) assets = load_assets(sim, args) + sim.prepare() log_info(f"Loaded {len(assets)} asset(s) successfully.", color="green") joint_controller = _setup_viser_joint_control(sim, assets, args) _publish_loaded_assets(sim, args) @@ -354,6 +364,7 @@ def _create_parser() -> argparse.ArgumentParser: prog="embodichain preview-asset", description="Preview a USD or mesh asset in the EmbodiChain simulation.", ) + add_env_launcher_args_to_parser(parser) parser.add_argument( "--asset_path", @@ -408,10 +419,15 @@ def _create_parser() -> argparse.ArgumentParser: help="Body type for rigid objects (default: kinematic).", ) parser.add_argument( - "--use_usd_properties", - action="store_true", - default=False, - help="Use physical properties from the USD file instead of defaults.", + "--asset_physics_mode", + "--asset-physics-mode", + dest="asset_physics_mode", + choices=("preserve", "overlay"), + default="overlay", + help=( + "Preserve source-authored physics or overlay explicitly configured " + "values (default: overlay)." + ), ) parser.add_argument( "--fix_base", @@ -419,25 +435,6 @@ def _create_parser() -> argparse.ArgumentParser: default=True, help="Fix or unfix the base of articulations (default: fixed).", ) - 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, @@ -447,12 +444,6 @@ def _create_parser() -> argparse.ArgumentParser: "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.", - ) parser.add_argument( "--joint-control", action=argparse.BooleanOptionalAction, @@ -463,9 +454,6 @@ def _create_parser() -> argparse.ArgumentParser: ), ) - from embodichain.lab.visualization import add_viser_args_to_parser - - add_viser_args_to_parser(parser) return parser diff --git a/embodichain/lab/sim/cfg.py b/embodichain/lab/sim/cfg.py deleted file mode 100644 index c36fabbcd..000000000 --- a/embodichain/lab/sim/cfg.py +++ /dev/null @@ -1,1929 +0,0 @@ -# ---------------------------------------------------------------------------- -# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ---------------------------------------------------------------------------- - -from __future__ import annotations - -import enum -import json -import os - -import dexsim -import numpy as np -import torch - -from typing import Sequence, Dict, Literal, List, Any, Optional -from dataclasses import field, MISSING - -from dexsim.types import ( - DenoiserType, - Renderer, - ToneMappingType, - PhysicalAttr, - ActorType, - AxisArrowType, - AxisCornerType, - VoxelConfig, - SoftBodyAttr, - SoftBodyMaterialModel, - ClothBodyAttr, -) -from embodichain.utils import configclass, is_configclass -from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT -from embodichain.data import get_data_path -from embodichain.utils import logger -from embodichain.utils.utility import key_in_nested_dict - -from .shapes import ShapeCfg, MeshCfg -from .workspace.cfg import RobotWorkspaceCfg - -# Global default renderer settings for simulation. -# -# The sentinel value ``"auto"`` defers the choice to GPU-based auto-selection -# performed lazily when a :class:`SimulationManager` is constructed (see -# :func:`embodichain.lab.sim.utility.render_utils.select_default_renderer`). Assigning a -# concrete renderer here (e.g. in test fixtures) forces that renderer and takes -# precedence over auto-selection. -DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" - - -@configclass -class RenderCfg: - renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" - """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'rt'. - - Note: - - 'auto' selects a default renderer based on the detected GPU: RTX-series cards use - 'hybrid', while datacenter cards (A100/A800, H100/H800/H200/H20) use 'fast-rt'. - If no CUDA device is available or the GPU is unknown, it falls back to 'hybrid'. - - 'hybrid' uses ray tracing for shadows and reflections while keeping rasterization for primary rendering, - providing a balance between performance and visual quality. - - 'fast-rt' is a fully ray-traced renderer for maximum visual fidelity, but may have higher computational cost. - - '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'.""" - - tone_mapping_enabled: bool = False - """Whether to map HDR RGB output with the modified Reinhard curve.""" - - tone_mapping_exposure: float = 1.0 - """Fixed linear exposure multiplier applied before tone mapping.""" - - def __post_init__(self) -> None: - """Validate rendering parameters.""" - if self.spp < 1: - logger.log_error("RenderCfg.spp must be at least 1.", ValueError) - if self.tone_mapping_exposure < 0.0: - logger.log_error( - "RenderCfg.tone_mapping_exposure must be non-negative.", ValueError - ) - - def to_dexsim_flags(self) -> Renderer: - """Convert the renderer name to DexSim's renderer enum.""" - if self.renderer == "hybrid": - return Renderer.HYBRID - elif self.renderer == "fast-rt": - return Renderer.FASTRT - elif self.renderer == "rt": - return Renderer.OFFLINERT - elif self.renderer == "auto": - # 'auto' is normally resolved by the SimulationManager before this is - # called. If it reaches here (e.g. used standalone), fall back safely. - logger.log_warning( - "Renderer 'auto' was not resolved before converting to dexsim flags. " - "Falling back to 'hybrid'." - ) - return Renderer.HYBRID - else: - logger.log_error( - f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'rt'." - ) - - def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: - """Apply rendering settings to a DexSim world configuration. - - Args: - world_config: DexSim world configuration to update in place. - """ - world_config.renderer = self.to_dexsim_flags() - world_config.raytrace_config.render_iterations_per_frame = self.spp - world_config.raytrace_config.open_denoise = True - world_config.raytrace_config.denoiser_type = DenoiserType.OPTIX - world_config.postprocess_config.tone_mapping_enabled = self.tone_mapping_enabled - world_config.postprocess_config.tone_mapping_type = ( - ToneMappingType.MODIFIED_REINHARD - ) - world_config.postprocess_config.tone_mapping_exposure = ( - self.tone_mapping_exposure - ) - - -@configclass -class PhysicsCfg: - gravity: np.ndarray = field(default_factory=lambda: np.array([0, 0, -9.81])) - """Gravity vector for the simulation environment.""" - - bounce_threshold: float = 2.0 - """The speed threshold below which collisions will not produce bounce effects.""" - - enable_ccd: bool = False - """Enable continuous collision detection (CCD) for fast-moving objects.""" - - length_tolerance: float = 0.05 - """The length tolerance for the simulation. - - 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. - """ - - def to_dexsim_args(self) -> Dict[str, Any]: - """Convert to DexSim physics arguments. - - Solver implementation details that are not exposed by :class:`PhysicsCfg` - retain their established defaults here. - """ - args = { - "gravity": self.gravity.tolist(), - "bounce_threshold": self.bounce_threshold, - "enable_ccd": self.enable_ccd, - "enable_enhanced_determinism": False, - "enable_friction_every_iteration": True, - } - return args - - -@configclass -class MarkerCfg: - """Configuration for visual markers in the simulation. - - This class defines properties for creating visual markers such as coordinate frames, - lines, and points that can be used for debugging, visualization, or reference purposes - in the simulation environment. - """ - - name: str = "empty-mesh" - """Name of the marker for identification purposes.""" - - marker_type: Literal["axis", "line", "point"] = "axis" - """Type of marker to display. Can be 'axis' (3D coordinate frame), 'line', or 'point'. (only axis supported now)""" - - axis_xpos: torch.Tensor | None = None - """List of 4x4 transformation matrices defining the position and orientation of each axis marker.""" - - axis_size: float = 0.002 - """Thickness/size of the axis lines in meters.""" - - axis_len: float = 0.005 - """Length of each axis arm in meters.""" - - line_color: List[float] = [1, 1, 0, 1.0] - """RGBA color values for the marker lines. Values should be between 0.0 and 1.0.""" - - arrow_type: AxisArrowType = AxisArrowType.CONE - """Type of arrow head for axis markers (e.g., CONE, ARROW, etc.).""" - - corner_type: AxisCornerType = AxisCornerType.SPHERE - """Type of corner/joint visualization for axis markers (e.g., SPHERE, CUBE, etc.).""" - - arena_index: int = -1 - """Index of the arena where the marker should be placed. -1 means all arenas.""" - - -@configclass -class WindowRecordCfg: - """Configuration for interactive viewer window recording.""" - - enable_hotkey: bool = True - """Whether to register the ``r`` hotkey for viewer recording when the window opens.""" - - save_path: str | None = None - """Optional output path for viewer recordings. If None, use the default outputs directory.""" - - fps: int = 20 - """Frames per second for viewer recording.""" - - max_memory: int = 1024 - """Maximum buffered recording memory in MB before auto-stopping capture.""" - - video_prefix: str = "viewer_record" - """Video file prefix used when no explicit save path is provided.""" - - -@configclass -class WindowCameraPoseCfg: - """Configuration for printing the interactive viewer camera pose.""" - - enable_hotkey: bool = True - """Whether to register the ``p`` hotkey when the window opens.""" - - convert_to_look_at: bool = True - """Whether the hotkey prints a ``set_look_at`` call instead of a matrix.""" - - -@configclass -class GPUMemoryCfg: - """A gpu memory configuration dataclass that neatly holds all parameters that configure physics GPU memory for simulation""" - - temp_buffer_capacity: int = 2**24 - """Increase this if you get 'PxgPinnedHostLinearMemoryAllocator: overflowing initial allocation size, increase capacity to at least %.' """ - - max_rigid_contact_count: int = 2**19 - """Increase this if you get 'Contact buffer overflow detected'""" - - max_rigid_patch_count: int = ( - 2**18 - ) # 81920 is DexSim default but most tasks work with 2**18 - """Increase this if you get 'Patch buffer overflow detected'""" - - 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 RigidBodyAttributesCfg: - """Physical attributes for rigid bodies. - - There are three parts of attributes that can be set: - 1. The dynamic properties, such as mass, damping, etc. - 2. The collision properties. - 3. The physics material properties. - """ - - mass: float = 1.0 - """Mass of the rigid body in kilograms. - - Set to 0 will use density to calculate mass. - """ - - density: float = 1000.0 - """Density of the rigid body in kg/m^3.""" - - angular_damping: float = 0.7 - """Angular damping coefficient.""" - - linear_damping: float = 0.7 - """Linear damping coefficient.""" - - max_depenetration_velocity: float = 10.0 - """Maximum depenetration velocity.""" - - sleep_threshold: float = 0.001 - """Threshold below which the body can go to sleep.""" - - min_position_iters: int = 4 - """Minimum position iterations.""" - - min_velocity_iters: int = 1 - """Minimum velocity iterations.""" - - max_linear_velocity: float = 1e2 - """Maximum linear velocity.""" - - max_angular_velocity: float = 1e2 - """Maximum angular velocity.""" - - # collision properties. - enable_ccd: bool = False - """Enable continuous collision detection (CCD).""" - - contact_offset: float = 0.002 - """Contact offset for collision detection.""" - - rest_offset: float = 0.0 - """Rest offset for collision detection.""" - - enable_collision: bool = True - """Enable collision for the rigid body.""" - - # physics material properties. - restitution: float = 0.0 - """Restitution (bounciness) coefficient.""" - - dynamic_friction: float = 0.5 - """Dynamic friction coefficient.""" - - static_friction: float = 0.5 - """Static friction coefficient.""" - - def attr(self) -> PhysicalAttr: - """Convert to dexsim PhysicalAttr""" - attr = PhysicalAttr() - attr.mass = self.mass - attr.contact_offset = self.contact_offset - attr.rest_offset = self.rest_offset - attr.dynamic_friction = self.dynamic_friction - attr.static_friction = self.static_friction - attr.angular_damping = self.angular_damping - attr.linear_damping = self.linear_damping - attr.sleep_threshold = self.sleep_threshold - attr.restitution = self.restitution - attr.enable_ccd = self.enable_ccd - attr.max_linear_velocity = self.max_linear_velocity - attr.max_angular_velocity = self.max_angular_velocity - attr.max_depenetration_velocity = self.max_depenetration_velocity - attr.min_position_iters = self.min_position_iters - attr.min_velocity_iters = self.min_velocity_iters - return attr - - @classmethod - def from_dict( - cls, init_dict: Dict[str, str | float | int] - ) -> RigidBodyAttributesCfg: - """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 - - -@configclass -class RigidBodyAttributesOverrideCfg: - """Partial rigid-body attribute overrides for per-link physics configuration. - - Fields set to ``None`` are not applied and retain values from the base - :class:`RigidBodyAttributesCfg`. - """ - - mass: float | None = None - density: float | None = None - angular_damping: float | None = None - linear_damping: float | None = None - max_depenetration_velocity: float | None = None - sleep_threshold: float | None = None - min_position_iters: int | None = None - min_velocity_iters: int | None = None - max_linear_velocity: float | None = None - max_angular_velocity: float | None = None - enable_ccd: bool | None = None - contact_offset: float | None = None - rest_offset: float | None = None - enable_collision: bool | None = None - restitution: float | None = None - dynamic_friction: float | None = None - static_friction: float | None = None - - def merge_with(self, base: RigidBodyAttributesCfg) -> PhysicalAttr: - """Build a :class:`~dexsim.types.PhysicalAttr` from base values and overrides.""" - merged = RigidBodyAttributesCfg() - for field_name in merged.__dataclass_fields__: - 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() - - @classmethod - def from_dict( - cls, init_dict: Dict[str, str | float | int | bool] - ) -> RigidBodyAttributesOverrideCfg: - """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 - - -@configclass -class LinkPhysicsOverrideCfg: - """Per-link physics override matched by regex on articulation link names.""" - - link_names_expr: list[str] = MISSING - """Regex patterns matched against link names (full match).""" - - attrs: RigidBodyAttributesOverrideCfg = RigidBodyAttributesOverrideCfg() - """Partial attribute overrides applied on top of :attr:`ArticulationCfg.attrs`.""" - - replace_inertial: bool = False - """Whether to recompute inertia when mass is overridden (DexSim flag).""" - - @classmethod - def from_dict(cls, init_dict: Dict[str, Any]) -> LinkPhysicsOverrideCfg: - """Initialize the configuration from a dictionary.""" - cfg = cls() - for key, value in init_dict.items(): - if key == "attrs" and isinstance(value, dict): - setattr(cfg, key, RigidBodyAttributesOverrideCfg.from_dict(value)) - elif hasattr(cfg, key): - setattr(cfg, key, value) - else: - logger.log_warning( - f"Key '{key}' not found in {cfg.__class__.__name__}." - ) - return cfg - - -def link_attrs_from_dict( - value: dict[str, Any], -) -> dict[str, LinkPhysicsOverrideCfg]: - """Parse a ``link_attrs`` mapping from YAML/JSON-style dicts.""" - link_attrs: dict[str, LinkPhysicsOverrideCfg] = {} - for group_name, group_cfg in value.items(): - if isinstance(group_cfg, LinkPhysicsOverrideCfg): - link_attrs[group_name] = group_cfg - elif isinstance(group_cfg, dict): - link_attrs[group_name] = LinkPhysicsOverrideCfg.from_dict(group_cfg) - else: - raise TypeError( - f"link_attrs['{group_name}'] must be a dict or " - f"LinkPhysicsOverrideCfg, got {type(group_cfg)}." - ) - return link_attrs - - -@configclass -class SoftbodyVoxelAttributesCfg: - # voxel config - triangle_remesh_resolution: int = 8 - """Resolution to remesh the softbody mesh before building physics collision mesh.""" - - triangle_simplify_target: int = 0 - """Simplify mesh faces to target value. Do nothing if this value is zero.""" - - # TODO: this value will be automatically computed with simulation_mesh_resolution and mesh scale. - maximal_edge_length: float = 0 - # """To shorten edges that are too long, additional points get inserted at their center leading to a subdivision of the input mesh. Do nothing if this value is zero.""" - - simulation_mesh_resolution: int = 8 - """Resolution to build simulation voxelize textra mesh. This value must be greater than 0.""" - - simulation_mesh_output_obj: bool = False - """Whether to output the simulation mesh as an obj file for debugging.""" - - def attr(self) -> VoxelConfig: - """Convert to dexsim VoxelConfig""" - attr = VoxelConfig() - attr.triangle_remesh_resolution = self.triangle_remesh_resolution - attr.maximal_edge_length = self.maximal_edge_length - attr.simulation_mesh_resolution = self.simulation_mesh_resolution - attr.triangle_simplify_target = self.triangle_simplify_target - return attr - - -@configclass -class SoftbodyPhysicalAttributesCfg: - # material properties - youngs: float = 1e6 - """Young's modulus (higher = stiffer).""" - - poissons: float = 0.45 - """Poisson's ratio (higher = closer to incompressible).""" - - dynamic_friction: float = 0.0 - """Dynamic friction coefficient.""" - - elasticity_damping: float = 0.0 - """Elasticity damping factor.""" - - # soft body properties - material_model: SoftBodyMaterialModel = SoftBodyMaterialModel.CO_ROTATIONAL - """Material constitutive model.""" - - # --- Mode / collision switches --- - enable_kinematic: bool = False - """If True, (partially) kinematic behavior is enabled.""" - - enable_ccd: bool = False - """Enable continuous collision detection (CCD).""" - - enable_self_collision: bool = False - """Enable self-collision handling.""" - - has_gravity: bool = True - """Whether the soft body is affected by gravity.""" - - # --- Self-collision & simplification parameters --- - self_collision_stress_tolerance: float = 0.9 - """Stress tolerance threshold for self-collision constraints.""" - - collision_mesh_simplification: bool = True - """Whether to simplify the collision mesh for self-collision.""" - - self_collision_filter_distance: float = 0.1 - """Distance threshold below which vertex pairs may be filtered from self-collision checks.""" - - # --- Damping, sleep & settling --- - vertex_velocity_damping: float = 0.005 - """Per-vertex velocity damping.""" - - linear_damping: float = 0.0 - """Global linear damping applied to the soft body.""" - - sleep_threshold: float = 0.05 - """Velocity/energy threshold below which the soft body can go to sleep.""" - - settling_threshold: float = 0.1 - """Threshold used to decide convergence/settling state.""" - - settling_damping: float = 10.0 - """Additional damping applied during settling phase.""" - - # --- Mass / density & velocity limits --- - mass: float = -1.0 - """Total mass of the soft body. If set to a negative value, density will be used to compute mass.""" - - density: float = 1000.0 - """Material density in kg/m^3.""" - - max_depenetration_velocity: float = 1e6 - """Maximum velocity used to resolve penetrations. Must be larger than zero.""" - - max_velocity: float = 100 - """Clamp for linear (or vertex) velocity. If set to zero, the limit is ignored.""" - - # --- Solver iteration counts --- - min_position_iters: int = 4 - """Minimum solver iterations for position correction.""" - - min_velocity_iters: int = 1 - """Minimum solver iterations for velocity updates.""" - - def attr(self) -> SoftBodyAttr: - attr = SoftBodyAttr() - attr.youngs = self.youngs - attr.poissons = self.poissons - attr.dynamic_friction = self.dynamic_friction - attr.elasticity_damping = self.elasticity_damping - attr.material_model = self.material_model - attr.enable_kinematic = self.enable_kinematic - attr.enable_ccd = self.enable_ccd - attr.enable_self_collision = self.enable_self_collision - attr.has_gravity = self.has_gravity - attr.self_collision_stress_tolerance = self.self_collision_stress_tolerance - attr.collision_mesh_simplification = self.collision_mesh_simplification - attr.vertex_velocity_damping = self.vertex_velocity_damping - attr.mass = self.mass - attr.density = self.density - attr.max_depenetration_velocity = self.max_depenetration_velocity - attr.max_velocity = self.max_velocity - attr.self_collision_filter_distance = self.self_collision_filter_distance - attr.linear_damping = self.linear_damping - attr.sleep_threshold = self.sleep_threshold - attr.settling_threshold = self.settling_threshold - attr.settling_damping = self.settling_damping - attr.min_position_iters = self.min_position_iters - attr.min_velocity_iters = self.min_velocity_iters - return attr - - -@configclass -class ClothPhysicalAttributesCfg: - # material properties - youngs: float = 1e10 - """Young's modulus (higher = stiffer).""" - - poissons: float = 0.3 - """Poisson's ratio.""" - - dynamic_friction: float = 0.5 - """Dynamic friction coefficient.""" - - elasticity_damping: float = 0.0 - """Elasticity damping factor.""" - - thickness: float = 0.001 - """Cloth thickness (m).""" - - bending_stiffness: float = 0.00001 - """Bending stiffness.""" - - bending_damping: float = 0.0 - """Bending damping.""" - - # cloth body properties - enable_kinematic: bool = False - """If True, (partially) kinematic behavior is enabled.""" - - enable_ccd: bool = True - """Enable continuous collision detection (CCD).""" - - enable_self_collision: bool = False - """Enable self-collision handling.""" - - has_gravity: bool = True - """Whether the cloth is affected by gravity.""" - - self_collision_stress_tolerance: float = 0.9 - """Stress tolerance threshold for self-collision constraints.""" - - collision_mesh_simplification: bool = True - """Whether to simplify the collision mesh for self-collision.""" - - vertex_velocity_damping: float = 0.005 - """Per-vertex velocity damping.""" - - mass: float = -1.0 - """Total mass of the cloth. If negative, density is used to compute mass.""" - - density: float = 1.0 - """Material density in kg/m^3.""" - - max_depenetration_velocity: float = 1e6 - """Maximum velocity used to resolve penetrations.""" - - max_velocity: float = 100.0 - """Clamp for linear (or vertex) velocity.""" - - self_collision_filter_distance: float = 0.1 - """Distance threshold for filtering self-collision vertex pairs.""" - - linear_damping: float = 0.05 - """Global linear damping applied to the cloth.""" - - sleep_threshold: float = 0.05 - """Velocity/energy threshold below which the cloth can go to sleep.""" - - settling_threshold: float = 0.1 - """Threshold used to decide convergence/settling state.""" - - settling_damping: float = 10.0 - """Additional damping applied during settling phase.""" - - min_position_iters: int = 4 - """Minimum solver iterations for position correction.""" - - min_velocity_iters: int = 1 - """Minimum solver iterations for velocity updates.""" - - def attr(self) -> ClothBodyAttr: - """Convert to dexsim ClothBodyAttr.""" - attr = ClothBodyAttr() - attr.youngs = self.youngs - attr.poissons = self.poissons - attr.dynamic_friction = self.dynamic_friction - attr.elasticity_damping = self.elasticity_damping - attr.thickness = self.thickness - attr.bending_stiffness = self.bending_stiffness - attr.bending_damping = self.bending_damping - attr.enable_kinematic = self.enable_kinematic - attr.enable_ccd = self.enable_ccd - attr.enable_self_collision = self.enable_self_collision - attr.has_gravity = self.has_gravity - attr.self_collision_stress_tolerance = self.self_collision_stress_tolerance - attr.collision_mesh_simplification = self.collision_mesh_simplification - attr.vertex_velocity_damping = self.vertex_velocity_damping - attr.mass = self.mass - attr.density = self.density - attr.max_depenetration_velocity = self.max_depenetration_velocity - attr.max_velocity = self.max_velocity - attr.self_collision_filter_distance = self.self_collision_filter_distance - attr.linear_damping = self.linear_damping - attr.sleep_threshold = self.sleep_threshold - attr.settling_threshold = self.settling_threshold - attr.settling_damping = self.settling_damping - attr.min_position_iters = self.min_position_iters - attr.min_velocity_iters = self.min_velocity_iters - return attr - - -@configclass -class JointDrivePropertiesCfg: - """Properties to define the drive mechanism of a joint.""" - - drive_type: Literal["force", "acceleration", "none"] = "force" - """Joint drive type to apply. - - If the drive type is "force", then the joint is driven by a force and the acceleration is computed based on the force applied. - If the drive type is "acceleration", then the joint is driven by an acceleration and the force is computed based on the acceleration applied. - If the drive type is "none", then no force will be applied to joint. - """ - - stiffness: Dict[str, float] | float = 1e4 - """Stiffness of the joint drive. - - The unit depends on the joint model: - - * For linear joints, the unit is kg-m/s^2 (N/m). - * For angular joints, the unit is kg-m^2/s^2/rad (N-m/rad). - """ - - damping: Dict[str, float] | float = 1e3 - """Damping of the joint drive. - - The unit depends on the joint model: - - * For linear joints, the unit is kg-m/s (N-s/m). - * For angular joints, the unit is kg-m^2/s/rad (N-m-s/rad). - """ - - max_effort: Dict[str, float] | float = 1e10 - """Maximum effort that can be applied to the joint (in kg-m^2/s^2).""" - - max_velocity: Dict[str, float] | float = 1e10 - """Maximum velocity that the joint can reach (in rad/s or m/s). - - For linear joints, this is the maximum linear velocity with unit m/s. - For angular joints, this is the maximum angular velocity with unit rad/s. - """ - - friction: Dict[str, float] | float = 0.0 - """Friction coefficient of the joint""" - - armature: Dict[str, float] | float = 0.0 - """Joint armature added to joint-space spatial inertia. - - Units depend on the joint model: - - * For prismatic (linear) joints, the unit is mass [kg]. - * For revolute (angular) joints, the unit is mass * scene_length^2 [kg-m^2]. - """ - - @classmethod - def from_dict( - cls, - init_dict: Dict[str, str | float | int | Dict[str, float]], - *, - defaults: JointDrivePropertiesCfg | None = None, - ) -> JointDrivePropertiesCfg: - """Initialize the configuration from a dictionary. - - Args: - init_dict: Joint-drive properties to override. - defaults: Optional base properties whose unspecified values are - preserved. If omitted, the class defaults are used. - - Returns: - Parsed joint-drive properties. - """ - cfg = defaults.copy() if defaults is not None else 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 - - -@configclass -class ObjectBaseCfg: - """Base configuration for an asset in the simulation. - - This class defines the basic properties of an asset, such as its type, initial state, and collision group. - It is used as a base class for specific asset configurations. - """ - - uid: str | None = None - - init_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) - """Position of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" - - init_rot: tuple[float, float, float] = (0.0, 0.0, 0.0) - """Euler angles (in degree) of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" - - init_local_pose: np.ndarray | None = None - """4x4 transformation matrix of the root in local frame. If specified, it will override init_pos and init_rot.""" - - @classmethod - def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> ObjectBaseCfg: - """Initialize the configuration from a dictionary.""" - cfg = cls() # Create a new instance of the class (cls) - for key, value in init_dict.items(): - if hasattr(cfg, key): - attr = getattr(cfg, key) - if is_configclass(attr): - setattr( - cfg, key, attr.from_dict(value) - ) # Call from_dict on the attribute - else: - setattr(cfg, key, value) - else: - logger.log_warning( - f"Key '{key}' not found in {cfg.__class__.__name__}." - ) - - # Automatically infer init_local_pose if not provided - if cfg.init_local_pose is None: - # If only init_pos or init_rot are provided, generate the 4x4 pose matrix - from scipy.spatial.transform import Rotation as R - - T = np.eye(4) - T[:3, 3] = np.array(cfg.init_pos) - T[:3, :3] = R.from_euler("xyz", np.deg2rad(cfg.init_rot)).as_matrix() - cfg.init_local_pose = T - else: - # If only init_local_pose is provided, extract init_pos and init_rot - from scipy.spatial.transform import Rotation as R - - T = np.array(cfg.init_local_pose) - cfg.init_pos = tuple(T[:3, 3]) - cfg.init_rot = tuple(R.from_matrix(T[:3, :3]).as_euler("xyz", degrees=True)) - - return cfg - - -@configclass -class LightCfg(ObjectBaseCfg): - """Configuration for a light asset in the simulation. - - Supports six light types matching the dexsim rendering backend: - - - ``"point"``: Per-environment omnidirectional point light with position - and falloff radius. Created as a batched light (one per environment). - - ``"sun"``: Global directional sun light (infinite distance). Created as - a single scene-level instance. Uses direction only; position is ignored. - Sun-specific fields (``angular_radius``, ``halo_size``, ``halo_falloff``) - are reserved for future backend support. - - ``"direction"``: Global pure directional light at infinite distance. - Created as a single scene-level instance. Direction only; no position. - - ``"spot"``: Per-environment spotlight with position, direction, and - inner/outer cone angles. Created as a batched light. - - ``"rect"``: Per-environment rectangular area light with position, - direction, width, and height. Created as a batched light. - - ``"mesh"``: Per-environment mesh-based emissive light. Requires a - :class:`~dexsim.models.MeshObject` via - :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` - (not tensor-batched). Created as a batched light. - - .. attention:: - The ``angular_radius``, ``halo_size``, and ``halo_falloff`` fields are - reserved for future use. The dexsim Python bindings do not yet expose - setters for these sun-specific properties. - """ - - light_type: Literal["point", "sun", "direction", "spot", "rect", "mesh"] = "point" - """Light type. Supported: ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``.""" - - # ------------------------------------------------------------------ - # Universal properties (apply to all light types) - # ------------------------------------------------------------------ - - color: tuple[float, float, float] = (1.0, 1.0, 1.0) - """RGB color of the light source. Defaults to white ``(1.0, 1.0, 1.0)``.""" - - intensity: float = 30.0 - """Intensity of the light source in watts/m^2. Defaults to ``30.0``.""" - - enable_shadow: bool = True - """Whether the light casts shadows. Defaults to ``True``.""" - - # ------------------------------------------------------------------ - # Point light - # ------------------------------------------------------------------ - - radius: float = 10.0 - """Falloff radius for point lights. Only used when ``light_type="point"``. Defaults to ``10.0``.""" - - # ------------------------------------------------------------------ - # Directional properties (sun, direction, spot, rect, mesh) - # ------------------------------------------------------------------ - - direction: tuple[float, float, float] = (0.0, 0.0, -1.0) - """Direction vector for directional, spot, rect, and mesh lights. - Defaults to ``(0.0, 0.0, -1.0)`` (pointing down along -Z).""" - - # ------------------------------------------------------------------ - # Sun light (reserved — Python bindings not yet available) - # ------------------------------------------------------------------ - - angular_radius: float = 0.5 - """Angular radius of the sun disc in degrees. Reserved for future use.""" - - halo_size: float = 10.0 - """Halo size for sun light. Reserved for future use.""" - - halo_falloff: float = 3.0 - """Halo falloff for sun light. Reserved for future use.""" - - # ------------------------------------------------------------------ - # Spot light - # ------------------------------------------------------------------ - - spot_angle_inner: float = 30.0 - """Inner cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. - Defaults to ``30.0``.""" - - spot_angle_outer: float = 45.0 - """Outer cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. - Defaults to ``45.0``.""" - - # ------------------------------------------------------------------ - # Rect light - # ------------------------------------------------------------------ - - rect_width: float = 1.0 - """Width of the rectangular area light. Only used when ``light_type="rect"``. - Defaults to ``1.0``.""" - - rect_height: float = 1.0 - """Height of the rectangular area light. Only used when ``light_type="rect"``. - Defaults to ``1.0``.""" - - # ------------------------------------------------------------------ - # Mesh light - # ------------------------------------------------------------------ - - mesh_path: str = "" - """Asset path for mesh-based emissive lights. Only used when ``light_type="mesh"``. - The actual mesh assignment is done via - :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` which accepts a - :class:`dexsim.models.MeshObject`. This field stores the path for reference.""" - - -@configclass -class RigidObjectCfg(ObjectBaseCfg): - """Configuration for a rigid body asset in the simulation. - - This class extends the base asset configuration to include specific properties for rigid bodies, - such as physical attributes and collision group. - """ - - shape: ShapeCfg = ShapeCfg() - """Shape configuration for the rigid body. """ - - # TODO: supoort basic primitive shapes, such as box, sphere, etc cfg and spawn method. - - attrs: RigidBodyAttributesCfg = RigidBodyAttributesCfg() - - body_type: Literal["dynamic", "kinematic", "static"] = "dynamic" - - max_convex_hull_num: int = MISSING - """The maximum number of convex hulls that will be created for the rigid body. - - .. deprecated:: - Use :attr:`MeshCfg.max_convex_hull_num` instead. This field is kept for - backward compatibility and overrides the shape-level value when explicitly set. - - If set to larger than 1, the rigid body will be decomposed into multiple convex hulls - using the approximate convex decomposition method specified by :attr:`acd_method`. - Reference: https://github.com/SarahWeiii/CoACD - """ - - acd_method: str = MISSING - """The method used for approximate convex decomposition (ACD) of the mesh. - - .. deprecated:: - Use :attr:`MeshCfg.acd_method` instead. This field is kept for - backward compatibility and overrides the shape-level value when explicitly set. - - Currently, ``"coacd"`` and ``"vhacd"`` are supported. Only used when - :attr:`max_convex_hull_num` is set to larger than 1. - """ - - sdf_resolution: int = MISSING - """Resolution for the signed distance field (SDF) of the rigid body. - - .. deprecated:: - Use :attr:`MeshCfg.sdf_resolution` instead. This field is kept for - backward compatibility and overrides the shape-level value when explicitly set. - - The spacing of the uniformly sampled SDF is equal to the largest AABB extent - of the mesh, divided by the resolution. If ``sdf_resolution`` is set to larger - than 0, an SDF will be generated for collision detection. SDF will increase the - accuracy of collision, but also takes more time to initialize and simulate. - """ - - body_scale: tuple | list = (1.0, 1.0, 1.0) - """Scale of the rigid body in the simulation world frame.""" - - use_usd_properties: bool = False - """Whether to use physical properties from USD file instead of config. - - When True: Keep all physical properties (drive, physics attrs, etc.) from USD file. - When False (default): Override USD properties with config values. - Only effective for USD files. - """ - - def to_dexsim_body_type(self) -> ActorType: - """Convert the body type to dexsim ActorType.""" - if self.body_type == "dynamic": - return ActorType.DYNAMIC - elif self.body_type == "kinematic": - return ActorType.KINEMATIC - elif self.body_type == "static": - return ActorType.STATIC - else: - logger.log_error( - f"Invalid body type '{self.body_type}' specified. Must be one of 'dynamic', 'kinematic', or 'static'." - ) - - -@configclass -class SoftObjectCfg(ObjectBaseCfg): - """Configuration for a soft body asset in the simulation. - - This class extends the base asset configuration to include specific properties for soft bodies, - such as physical attributes and collision group. - """ - - voxel_attr: SoftbodyVoxelAttributesCfg = SoftbodyVoxelAttributesCfg() - """Tetra mesh voxelization attributes for the soft body.""" - - physical_attr: SoftbodyPhysicalAttributesCfg = SoftbodyPhysicalAttributesCfg() - """Physical attributes for the soft body.""" - - shape: MeshCfg = MeshCfg() - """Mesh configuration for the soft body.""" - - -@configclass -class ClothObjectCfg(ObjectBaseCfg): - """Configuration for a cloth body asset in the simulation. - - This class extends the base asset configuration to include specific properties for cloth bodies, - such as physical attributes and collision group. - """ - - physical_attr: ClothPhysicalAttributesCfg = ClothPhysicalAttributesCfg() - """Physical attributes for the cloth body.""" - - shape: MeshCfg = MeshCfg() - """Mesh configuration for the cloth body.""" - - -@configclass -class RigidObjectGroupCfg: - """Configuration for a rigid object group asset in the simulation. - - Rigid object groups can be initialized from multiple rigid object configurations specified in a folder. - If `folder_path` is specified, user should provide a RigidObjectCfg in `rigid_objects` as a template configuration for - all objects in the group. - - For example: - ```python - rigid_object_group: RigidObjectGroupCfg( - folder_path="path/to/folder", - max_num=5, - rigid_objects={ - "template_obj": RigidObjectCfg( - shape=MeshCfg( - fpath="", # fpath will be ignored when folder_path is specified - ), - body_type="dynamic", - ) - } - ) - """ - - uid: str | None = None - - rigid_objects: Dict[str, RigidObjectCfg] = MISSING - """Configuration for the rigid objects in the group.""" - - body_type: Literal["dynamic", "kinematic"] = "dynamic" - """Body type for all rigid objects in the group. """ - - folder_path: str | None = None - """Path to the folder containing the rigid object assets. - - This is used to initialize multiple rigid object configurations from a folder. - """ - - max_num: int = 1 - """Maximum number of rigid objects to initialize from the folder. - - This is only used when `folder_path` is specified. - """ - - ext: str = ".obj" - """File extension for the rigid object assets. - - This is only used when `folder_path` is specified. - """ - - @classmethod - def from_dict(cls, init_dict: Dict[str, Any]) -> RigidObjectGroupCfg: - """Initialize the configuration from a dictionary.""" - cfg = cls() - for key, value in init_dict.items(): - if hasattr(cfg, key): - attr = getattr(cfg, key) - if is_configclass(attr): - setattr( - cfg, key, attr.from_dict(value) - ) # Call from_dict on the attribute - elif key == "rigid_objects" and "folder_path" not in init_dict: - rigid_objects_cfg = {} - for obj_name, obj_cfg in value.items(): - rigid_objects_cfg[obj_name] = RigidObjectCfg.from_dict(obj_cfg) - setattr(cfg, key, rigid_objects_cfg) - elif key == "rigid_objects" and "folder_path" in init_dict: - folder_path = init_dict["folder_path"] - max_num = init_dict.get("max_num", 1) - rigid_objects_cfg = {} - if os.path.exists(folder_path) and os.path.isdir(folder_path): - files = os.listdir(folder_path) - files = [f for f in files if f.endswith(cfg.ext)] - # select files up to max_num - n_file = len(files) - select_files = [] - for i in range(max_num): - select_files.append(files[i % n_file]) - - for i, file_name in enumerate(select_files): - file_path = os.path.join(folder_path, file_name) - rigid_obj_cfg: RigidObjectCfg = RigidObjectCfg.from_dict( - list(init_dict["rigid_objects"].values())[0] - ) - rigid_obj_cfg.uid = f"{cfg.uid}_obj_{i}" - rigid_obj_cfg.shape.fpath = file_path - rigid_objects_cfg[rigid_obj_cfg.uid] = rigid_obj_cfg - setattr(cfg, "rigid_objects", rigid_objects_cfg) - else: - logger.log_error( - f"Folder '{folder_path}' does not exist or is not a directory." - ) - else: - setattr(cfg, key, value) - else: - logger.log_warning( - f"Key '{key}' not found in {cfg.__class__.__name__}." - ) - return cfg - - -@configclass -class RigidConstraintCfg: - """Configuration for a fixed constraint between two RigidObjects. - - The constraint binds rigid_object_a's entity[i] to rigid_object_b's entity[i] - within arena[i] (one constraint per arena). - - Args: - name: Base constraint name. Per-arena names are derived as ``f"{name}"`` - (single env) or ``f"{name}_{i}"`` (multi env). - rigid_object_a_uid: UID of the first RigidObject (must exist in the sim). - rigid_object_b_uid: UID of the second RigidObject (must exist in the sim). - local_frame_a: 4x4 joint frame in object A's local coordinates. - ``None`` -> identity (object A's origin). Accepts a single - ``(4, 4)`` matrix (shared by all envs) or an ``(N, 4, 4)`` array - (one frame per env). Defaults to None. - local_frame_b: 4x4 joint frame in object B's local coordinates. - ``None`` -> the frame is computed per env as ``inv(pose_B) @ pose_A`` - from the objects' current poses, so the constraint welds the objects - at their *current* relative pose (rather than pulling their origins - together). An explicit ``(4, 4)`` or ``(N, 4, 4)`` value is used - verbatim. Defaults to None. - constraint_type: Reserved for future typed constraints (prismatic, - revolute, spherical, d6). Only ``"fixed"`` is supported in v1. - - .. attention:: - Both objects must be :class:`RigidObject` instances and must share the - same number of arenas. - """ - - name: str = MISSING - """Base name of the constraint (per-arena names are derived from this).""" - - rigid_object_a_uid: str = MISSING - """UID of the first RigidObject.""" - - rigid_object_b_uid: str = MISSING - """UID of the second RigidObject.""" - - local_frame_a: np.ndarray | None = None - """Local joint frame on object A. None -> identity (object A's origin).""" - - local_frame_b: np.ndarray | None = None - """Local joint frame on object B. None -> ``inv(pose_B) @ pose_A`` per env - (weld at the objects' current relative pose).""" - - constraint_type: Literal["fixed"] = "fixed" - """Constraint type. Only ``"fixed"`` is supported in v1.""" - - -@configclass -class URDFCfg: - """Standalone configuration class for URDF assembly.""" - - components: Dict[str, Dict[str, str | Dict | np.ndarray]] = field( - default_factory=dict - ) - """Dictionary of robot components to be assembled.""" - - sensors: Dict[str, Dict[str, str | np.ndarray]] = field(default_factory=dict) - """Dictionary of sensors to be attached to the robot.""" - - use_signature_check: bool = True - """Whether to use signature check when merging URDFs.""" - - base_link_name: str = "base_link" - """Name of the base link in the assembled robot.""" - - fpath: str | None = None - """Full output file path for the assembled URDF. If specified, overrides fname and fpath_prefix.""" - - fname: str | None = None - """Name used for output file and directory. If not specified, auto-generated from component names.""" - - fpath_prefix: str = EMBODICHAIN_DEFAULT_DATA_ROOT + "/assembled" - """Output directory prefix for the assembled URDF file.""" - - component_prefix: List[tuple[str, str | None]] = field( - default_factory=lambda: [ - ("chassis", None), - ("legs", None), - ("torso", None), - ("head", None), - ("left_arm", "left_"), - ("right_arm", "right_"), - ("left_hand", "left_"), - ("right_hand", "right_"), - ("arm", None), - ("hand", None), - ] - ) - """Component name prefixes used during URDF assembly. - - Preferred form is a list of ``(component_name, prefix)`` tuples. For - convenience, a mapping ``{component_name: prefix}`` is also accepted when - constructing :class:`URDFCfg` and will be normalized internally. - """ - - name_case: dict[str, str] = field( - default_factory=lambda: { - "joint": "original", - "link": "original", - } - ) - """Case normalization policy applied to joint/link names during URDF assembly. - - Supported values per key are ``"upper"``, ``"lower"`` or ``"original"`` - (legacy alias ``"none"``). The default preserves source URDF casing. - """ - - def __init__( - self, - components: list[dict[str, str | np.ndarray]] | None = None, - sensors: dict[str, dict[str, str | np.ndarray]] | None = None, - fpath: str | None = None, - fname: str | None = None, - fpath_prefix: str = EMBODICHAIN_DEFAULT_DATA_ROOT + "/assembled", - use_signature_check: bool = True, - base_link_name: str = "base_link", - component_prefix: list[tuple[str, str | None]] | None = None, - name_case: dict[str, str] | None = None, - ): - """ - Initialize URDFCfg with optional list of components and output path settings. - - Args: - components (list[dict[str, str | np.ndarray]] | None): List of component configurations. Each dict should contain: - - 'component_type' (str): The type/name of the component (e.g., 'chassis', 'arm', 'hand'). - - 'urdf_path' (str): Path to the component's URDF file. - - 'transform' (np.ndarray | None): 4x4 transformation matrix (optional). - - Additional params can be included as extra keys. - sensors (dict[str, dict[str, str | np.ndarray]] | None): Sensor configurations for the robot. - fpath (str | None): Full output file path for the assembled URDF. If specified, overrides fname and fpath_prefix. - fname (str | None): Name used for output file and directory. If not specified, auto-generated from component names. - fpath_prefix (str): Output directory prefix for the assembled URDF file. - use_signature_check (bool): Whether to use signature check when merging URDFs. - base_link_name (str): Name of the base link in the assembled robot. - component_prefix (list[tuple[str, str | None]] | None): Optional - list of (component_type, prefix) pairs to override default - component name prefixes. - """ - self.components = {} - self.sensors = sensors or {} - self.fpath = fpath - self.use_signature_check = use_signature_check - self.base_link_name = base_link_name - self.fname = fname - self.fpath_prefix = fpath_prefix - - # Initialize component prefixes (patch-style mapping per component type) - if component_prefix is None: - # Use the same default as the dataclass field - self.component_prefix = [ - ("chassis", None), - ("legs", None), - ("torso", None), - ("head", None), - ("left_arm", "left_"), - ("right_arm", "right_"), - ("left_hand", "left_"), - ("right_hand", "right_"), - ("arm", None), - ("hand", None), - ] - elif isinstance(component_prefix, dict): - # Allow dict-style config: {"left_hand": "l_", ...} - self.component_prefix = list(component_prefix.items()) - else: - # Assume caller provided a list of (component_name, prefix) tuples - self.component_prefix = component_prefix - - if name_case is None: - self.name_case = { - "joint": "original", - "link": "original", - } - else: - self.name_case = name_case - - # Auto-add components if provided - if components: - for comp_config in components: - if not isinstance(comp_config, dict): - logger.log_error( - f"Component configuration must be a dict, got {type(comp_config)}" - ) - continue - - # Extract required fields - component_type = comp_config.get("component_type") - urdf_path = comp_config.get("urdf_path") - - if not component_type or not urdf_path: - logger.log_error( - f"Component configuration must contain 'component_type' and 'urdf_path', got {comp_config}" - ) - continue - - # Extract optional fields - transform = comp_config.get("transform", np.eye(4)) - - # Extract additional params (exclude known keys) - params = { - k: v - for k, v in comp_config.items() - if k not in ["component_type", "urdf_path", "transform"] - } - - # Add the component - self.add_component(component_type, urdf_path, transform, **params) - - if sensors is not None: - # Accept both list and dict; serialization round-trips an empty - # dict when no sensors are configured (the field default). - if isinstance(sensors, dict) and not sensors: - self.sensors = [] - elif not isinstance(sensors, (list, dict)): - logger.log_error( - f"sensors must be a list of dicts or a dict, got {type(sensors)}" - ) - self.sensors = [] - elif isinstance(sensors, dict): - # dict keyed by sensor_name -> config - self.sensors = list(sensors.values()) - else: - # Optionally check each sensor dict - valid_sensors = [] - for sensor_config in sensors: - if not isinstance(sensor_config, dict): - logger.log_error( - f"Sensor configuration must be a dict, got {type(sensor_config)}" - ) - continue - sensor_name = sensor_config.get("sensor_name") - if not sensor_name: - logger.log_error( - f"Sensor configuration must contain 'sensor_name', got {sensor_config}" - ) - continue - valid_sensors.append(sensor_config) - self.sensors = valid_sensors - - def set_urdf(self, urdf_path: str) -> "URDFCfg": - """Directly specify a single URDF file for the robot, compatible with the single-URDF robot case. - - Args: - urdf_path (str): Path to the robot's URDF file. - - Returns: - URDFCfg: Returns self to allow method chaining. - """ - self.components.clear() - urdf_file = os.path.splitext(os.path.basename(urdf_path))[0] - self.components[urdf_file] = { - "urdf_path": urdf_path, - "transform": None, - "params": {}, - } - self.fpath = urdf_path - return self - - def add_component( - self, - component_type: str, - urdf_path: str, - transform: np.ndarray | None = None, - **params, - ) -> URDFCfg: - """Add a robot component to the assembly configuration. - - Args: - component_type (str): The type/name of the component. Should be one of SUPPORTED_COMPONENTS - (e.g., 'chassis', 'torso', 'head', 'left_arm', 'right_hand', 'arm', 'hand', etc.). - urdf_path (str): Path to the component's URDF file. - transform (np.ndarray | None): 4x4 transformation matrix for the component in the robot frame (default: None). - **params: Additional keyword parameters for the component (e.g., color, material, etc.). - - Returns: - URDFCfg: Returns self to allow method chaining. - """ - if urdf_path: - if not os.path.exists(urdf_path): - urdf_path_candidate = get_data_path(urdf_path) - if os.path.exists(urdf_path_candidate): - urdf_path = urdf_path_candidate - else: - logger.log_error(f"URDF path '{urdf_path}' does not exist.") - raise FileNotFoundError(f"URDF path '{urdf_path}' does not exist.") - - if transform is None: - transform = np.eye(4) - - self.components[component_type] = { - "urdf_path": urdf_path, - "transform": np.array(transform), - "params": params, - } - - if self.fname: - self.fpath = f"{self.fpath_prefix}/{self.fname}/{self.fname}.urdf" - else: - # Update output_path to use all component urdf file names joined by underscores as directory - if len(self.components) == 1: - # Only one component, use its urdf file name - urdf_file = os.path.splitext(os.path.basename(urdf_path))[0] - name = urdf_file - else: - # Multiple components, join all urdf file names - urdf_files = [ - os.path.splitext(os.path.basename(v["urdf_path"]))[0] - for v in self.components.values() - ] - name = "_".join(urdf_files) - self.fpath = f"{self.fpath_prefix}/{name}/{name}.urdf" - - return self - - def add_sensor(self, sensor_name: str, **sensor_config) -> URDFCfg: - """Add a sensor to the robot configuration. - - Args: - sensor_name (str): The name of the sensor. - **sensor_config: Additional configuration parameters for the sensor. - - Returns: - URDFCfg: Returns self to allow method chaining. - """ - self.sensors.append({"sensor_name": sensor_name, **sensor_config}) - return self - - def assemble_urdf(self) -> str: - """Assemble URDF files for the robot based on the configuration. - - Returns: - str: The path to the resulting (possibly merged) URDF file. - """ - components = list(self.components.items()) - # If there is only one component, return its URDF path directly. - if len(components) == 1: - _, comp_config = components[0] - return comp_config["urdf_path"] - - from embodichain.toolkits.urdf_assembly import URDFAssemblyManager - - # If there are multiple components, merge them into a single URDF file. - manager = URDFAssemblyManager() - manager.base_link_name = self.base_link_name - - if self.component_prefix is None: - self.component_prefix = [ - ("left_arm", "left_"), - ("right_arm", "right_"), - ("left_hand", "left_"), - ("right_hand", "right_"), - ] - if isinstance(self.component_prefix, dict): - self.component_prefix = list(self.component_prefix.items()) - # Forward configured component prefixes to the assembly manager - manager.component_prefix = self.component_prefix - - if self.name_case is not None: - manager.name_case = self.name_case - - for comp_type, comp_config in components: - params = comp_config.get("params", {}) - success = manager.add_component( - comp_type, - comp_config["urdf_path"], - comp_config.get("transform"), - **params, - ) - if not success: - logger.log_error( - f"Failed to add component '{comp_type}' with config: {comp_config}" - ) - - for sensor in self.sensors: - manager.attach_sensor( - sensor_name=sensor.get("sensor_name"), - sensor_source=sensor.get("sensor_source"), - parent_component=sensor.get("parent_component"), - parent_link=sensor.get("parent_link"), - sensor_type=sensor.get("sensor_type"), - **{ - k: v - for k, v in sensor.items() - if k - not in [ - "sensor_name", - "sensor_source", - "parent_component", - "parent_link", - "sensor_type", - ] - }, - ) - - try: - # Merge all added components into a single URDF file at the specified output path. - merged_urdf_xml = manager.merge_urdfs(self.fpath, self.use_signature_check) - except Exception as e: - logger.log_error(f"URDF merge failed: {e}") - - return self.fpath - - @classmethod - def from_dict(cls, init_dict: Dict) -> "URDFCfg": - if isinstance(init_dict, cls): - return init_dict - components = init_dict.get("components", None) - if isinstance(components, dict): - components = [{"component_type": k, **v} for k, v in components.items()] - sensors = init_dict.get("sensors", None) - fpath = init_dict.get("fpath", None) - use_signature_check = init_dict.get("use_signature_check", True) - base_link_name = init_dict.get("base_link_name", "base_link") - component_prefix = init_dict.get("component_prefix", None) - name_case = init_dict.get("name_case", None) - return cls( - components=components, - sensors=sensors, - fpath=fpath, - use_signature_check=use_signature_check, - base_link_name=base_link_name, - component_prefix=component_prefix, - name_case=name_case, - ) - - -@configclass -class ArticulationCfg(ObjectBaseCfg): - """Configuration for an articulation asset in the simulation. - - This class extends the base asset configuration to include specific properties for articulations, - such as joint drive properties, physical attributes. - """ - - fpath: str = None - """Path to the articulation asset file.""" - - drive_pros: JointDrivePropertiesCfg = JointDrivePropertiesCfg(drive_type="none") - """Properties to define the drive mechanism of a joint.""" - - body_scale: tuple | list = (1.0, 1.0, 1.0) - """Scale of the articulation in the simulation world frame.""" - - attrs: RigidBodyAttributesCfg = RigidBodyAttributesCfg() - """Physical attributes for all links. We use default mass from the USD/URDF file if available. - The mass and density in attrs will only be used if specified. - """ - - link_attrs: dict[str, LinkPhysicsOverrideCfg] | None = None - """Named per-link physics override groups keyed by regex on link names. - - Each group applies :attr:`LinkPhysicsOverrideCfg.attrs` on top of :attr:`attrs` for - matched links only. A link must not match more than one group. - """ - - fix_base: bool = True - """Whether to fix the base of the articulation. - - Set to True for articulations that should not move, such as a fixed base robot arm or a door. - Set to False for articulations that should move freely, such as a mobile robot or a humanoid robot. - """ - - disable_self_collision: bool = True - """Whether to enable or disable self-collisions.""" - - init_qpos: torch.Tensor | np.ndarray | Sequence[float] = None - """Initial joint positions of the articulation. - - If None, the joint positions will be set to zero. - If provided, it should be a array of shape (num_joints,). - """ - - qpos_limits: ( - torch.Tensor | np.ndarray | Sequence[float] | Dict[str, List[float]] | None - ) = None - """Override joint position limits of the articulation. - - If None, the joint position limits from the asset file (URDF/USD) are used. - If provided as a tensor/array of shape (num_joints, 2), it is applied to all - joints in the order of ``joint_names``. - If provided as a dictionary, keys are joint names or regular expressions and - values are ``[min, max]`` limits. - - This field replaces the asset limits for the articulation and can be used to - either tighten or expand the allowed range. - """ - - sleep_threshold: float = 0.005 - """Energy below which the articulation may go to sleep. Range: [0, max_float32]""" - - min_position_iters: int = 4 - """Number of position iterations the solver should perform for this articulation. Range: [1,255].""" - - min_velocity_iters: int = 1 - """Number of velocity iterations the solver should perform for this articulation. Range: [0,255].""" - - build_pk_chain: bool = True - """Whether to build pytorch-kinematics chain for forward kinematics and jacobian computation.""" - - compute_uv: bool = False - """Whether to compute the UV mapping for the articulation link. - - Currently, the uv mapping is computed for each link with projection uv mapping method. - """ - - use_usd_properties: bool = False - """Whether to use physical properties from USD file instead of config. - - When True: Keep all physical properties (drive, physics attrs, etc.) from USD file. - When False (default): Override USD properties with config values (URDF behavior). - Only effective for USD files, ignored for URDF files. - """ - - @classmethod - def from_dict( - cls, init_dict: Dict[str, str | float | tuple | dict] - ) -> ArticulationCfg: - """Initialize the configuration from a dictionary.""" - cfg = cls() - for key, value in init_dict.items(): - if key == "link_attrs" and isinstance(value, dict): - cfg.link_attrs = link_attrs_from_dict(value) - elif hasattr(cfg, key): - attr = getattr(cfg, key) - if isinstance(attr, JointDrivePropertiesCfg) and isinstance( - value, dict - ): - setattr( - cfg, - key, - JointDrivePropertiesCfg.from_dict(value, defaults=attr), - ) - elif is_configclass(attr): - setattr(cfg, key, attr.from_dict(value)) - else: - setattr(cfg, key, value) - else: - logger.log_warning( - f"Key '{key}' not found in {cfg.__class__.__name__}." - ) - - if cfg.init_local_pose is None: - from scipy.spatial.transform import Rotation as R - - T = np.eye(4) - T[:3, 3] = np.array(cfg.init_pos) - T[:3, :3] = R.from_euler("xyz", np.deg2rad(cfg.init_rot)).as_matrix() - cfg.init_local_pose = T - else: - from scipy.spatial.transform import Rotation as R - - cfg.init_pos = tuple(cfg.init_local_pose[:3, 3]) - cfg.init_rot = tuple( - R.from_matrix(cfg.init_local_pose[:3, :3]).as_euler("xyz", degrees=True) - ) - - return cfg - - -@configclass -class RobotCfg(ArticulationCfg): - from embodichain.lab.sim.solvers import SolverCfg - - """Configuration for a robot asset in the simulation. - """ - - drive_pros: JointDrivePropertiesCfg = JointDrivePropertiesCfg(drive_type="force") - """Properties to define the drive mechanism of a joint.""" - - control_parts: Dict[str, List[str]] | None = None - """Control parts is the mapping from part name to joint names. - - For example, {'left_arm': ['joint1', 'joint2'], 'right_arm': ['joint3', 'joint4']} - If no control part is specified, the robot will use all joints as a single control part. - - Note: - - if `control_parts` is specified, `solver_cfg` must be a dict with part names as - keys corresponding to the control parts name. - - The joint names in the control parts support regular expressions, e.g., 'joint[1-6]'. - After initialization of robot, the names will be expanded to a list of full joint names. - - `Robot` is a derived class of `Articulation`, with control parts support. So the `drive_pros` - in `ArticulationCfg` can use control part as key to specify the corresponding joint drive properties, - which will be overridden if these joint names are already specified. - """ - - urdf_cfg: URDFCfg | None = None - """URDF assembly configuration which allows for assembling a robot from multiple URDF components. - """ - - # TODO: how to support one solver for multiple parts? - solver_cfg: SolverCfg | Dict[str, SolverCfg] | None = None - """Solver is used to compute forward and inverse kinematics for the robot. - """ - - workspace_cfg: Dict[str, RobotWorkspaceCfg] | None = None - """Runtime workspace cache configuration keyed by control-part name.""" - - @classmethod - def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> RobotCfg: - """Initialize the configuration from a dictionary.""" - if isinstance(init_dict, cls): - return init_dict - - import importlib - - solver_module = importlib.import_module("embodichain.lab.sim.solvers") - - cfg = cls() # Create a new instance of the class (cls) - for key, value in init_dict.items(): - if key == "link_attrs" and isinstance(value, dict): - cfg.link_attrs = link_attrs_from_dict(value) - elif hasattr(cfg, key): - attr = getattr(cfg, key) - if key == "urdf_cfg": - from embodichain.lab.sim.cfg import URDFCfg - - setattr(cfg, key, URDFCfg.from_dict(value)) - elif key == "workspace_cfg" and isinstance(value, dict): - setattr( - cfg, - key, - { - part: ( - part_cfg - if isinstance(part_cfg, RobotWorkspaceCfg) - else RobotWorkspaceCfg(**part_cfg) - ) - for part, part_cfg in value.items() - }, - ) - elif key == "fpath": - setattr(cfg, key, get_data_path(value)) - elif isinstance(attr, JointDrivePropertiesCfg) and isinstance( - value, dict - ): - setattr( - cfg, - key, - JointDrivePropertiesCfg.from_dict(value, defaults=attr), - ) - elif is_configclass(attr): - setattr( - cfg, key, attr.from_dict(value) - ) # Call from_dict on the attribute - elif isinstance(value, dict) and "class_type" in value: - setattr( - cfg, - key, - getattr(solver_module, f"{value['class_type']}Cfg").from_dict( - value - ), - ) - elif isinstance(value, dict) and key_in_nested_dict( - value, "class_type" - ): - setattr( - cfg, - key, - { - k: getattr( - solver_module, f"{v['class_type']}Cfg" - ).from_dict(v) - for k, v in value.items() - }, - ) - - else: - setattr(cfg, key, value) - else: - logger.log_warning( - f"Key '{key}' not found in {cfg.__class__.__name__}." - ) - return cfg - - def _build_defaults(self, init_dict: dict | None = None) -> None: - """Populate default config fields from ``init_dict``. - - Subclasses override this to read variant/version fields from - ``init_dict``, set them on ``self``, and populate ``urdf_cfg``, - ``control_parts``, ``solver_cfg``, ``drive_pros`` and ``attrs``. - The base implementation is a no-op. - - .. attention:: - Do NOT call :func:`merge_robot_cfg` from here -- the subclass - ``from_dict`` calls this hook first, then ``merge_robot_cfg``. - Calling ``merge_robot_cfg`` here would recurse, because - ``merge_robot_cfg`` itself calls ``RobotCfg.from_dict``. - - Args: - init_dict: The raw override dict passed to ``from_dict``. - """ - return None - - def to_dict(self): - """Serialize config to a plain dict (enums, numpy, nested configclass).""" - - def serialize(obj, _visited=None): - if _visited is None: - _visited = set() - if isinstance(obj, enum.Enum): - return obj.value - if isinstance(obj, (dict, object)) and not isinstance( - obj, (str, int, float, bool, type(None)) - ): - obj_id = id(obj) - if obj_id in _visited: - return None - _visited.add(obj_id) - - if isinstance(obj, np.ndarray): - return obj.tolist() - if isinstance(obj, dict): - return { - (k.value if isinstance(k, enum.Enum) else str(k)): serialize( - v, _visited - ) - for k, v in obj.items() - } - if isinstance(obj, (list, tuple)): - return [serialize(v, _visited) for v in obj] - if hasattr(obj, "to_dict") and obj is not self: - return serialize(obj.to_dict(), _visited) - if hasattr(obj, "__dict__"): - return { - k: serialize(v, _visited) - for k, v in obj.__dict__.items() - if v is not None - } - return obj - - return serialize(self) - - def to_string(self): - """Return config as a JSON string.""" - return json.dumps(self.to_dict(), indent=2) - - def save_to_file(self, filepath): - """Save config to a local file as JSON.""" - with open(filepath, "w") as f: - f.write(self.to_string()) - - def build_pk_serial_chain( - self, device: torch.device = torch.device("cpu"), **kwargs - ) -> Dict[str, "pk.SerialChain"]: - """Build the serial chain from the URDF file. - - Note: - This method is usually used in imitation dataset saving (compute eef pose from qpos using FK) - and model training (provide a differentiable FK layer or loss computation). - - Args: - device (torch.device): The device to which the chain will be moved. Defaults to CPU. - **kwargs: Additional arguments for building the serial chain. - - Returns: - Dict[str, pk.SerialChain]: The serial chain of the robot for specified control part. - """ - return {} diff --git a/embodichain/lab/sim/cfg/__init__.py b/embodichain/lab/sim/cfg/__init__.py new file mode 100644 index 000000000..a1c077a96 --- /dev/null +++ b/embodichain/lab/sim/cfg/__init__.py @@ -0,0 +1,136 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Public simulation-configuration facade. + +The implementation is split by domain while this package preserves the +historical ``embodichain.lab.sim.cfg`` import surface. +""" + +from __future__ import annotations + +from typing import Literal + +from embodichain.data import get_data_path + +from ..shapes import MeshCfg, MeshCollisionApproximation, MeshCollisionCfg, ShapeCfg +from ..workspace.cfg import RobotWorkspaceCfg +from .articulation import ( + ArticulationCfg, + ArticulationRootPropertiesCfg, + JointDrivePropertiesCfg, + LinkPhysicsOverrideCfg, + NewtonJointDrivePropertiesCfg, + _normalize_joint_target_mode, + _raise_removed_articulation_cfg_fields, + link_attrs_from_dict, +) +from .asset import AssetPhysicsMode, ObjectBaseCfg, _resolve_asset_physics_mode +from .deformable import ( + ClothObjectCfg, + ClothPhysicalAttributesCfg, + DeformableObjectCfg, + SoftObjectCfg, + SoftbodyPhysicalAttributesCfg, + SoftbodyVoxelAttributesCfg, + SurfaceDeformableObjectCfg, + VolumeDeformableObjectCfg, +) +from .rigid import ( + CollisionPropertiesCfg, + DefaultCollisionPropertiesCfg, + DefaultRigidBodyPropertiesCfg, + MassPropertiesCfg, + NewtonCollisionPropertiesCfg, + NewtonRigidBodyMaterialCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, +) +from .rigid_object import RigidObjectCfg, RigidObjectGroupCfg +from .scene import LightCfg, RigidConstraintCfg +from .simulation import ( + DefaultPhysicsCfg, + GPUMemoryCfg, + NewtonCollisionPipelineCfg, + NewtonPhysicsCfg, + PhysicsBackendCfg, + RenderCfg, + physics_backend_from_cfg, + physics_cfg_for_backend, + validate_physics_cfg, +) +from .urdf import URDFCfg +from .viewer import MarkerCfg, WindowCameraPoseCfg, WindowRecordCfg + +# The renderer selection code intentionally mutates this package-level value. +DEFAULT_RENDERER: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" + +# Robot imports are kept last because SolverCfg discovery imports simulation +# modules that themselves rely on the public facade above. +from .robot import RobotCfg, RobotPresetCfg # noqa: E402 + +__all__ = [ + "DEFAULT_RENDERER", + "AssetPhysicsMode", + "RenderCfg", + "GPUMemoryCfg", + "PhysicsBackendCfg", + "DefaultPhysicsCfg", + "NewtonCollisionPipelineCfg", + "NewtonPhysicsCfg", + "physics_cfg_for_backend", + "physics_backend_from_cfg", + "validate_physics_cfg", + "MarkerCfg", + "WindowRecordCfg", + "WindowCameraPoseCfg", + "ShapeCfg", + "MeshCfg", + "MeshCollisionApproximation", + "MeshCollisionCfg", + "MassPropertiesCfg", + "DefaultRigidBodyPropertiesCfg", + "CollisionPropertiesCfg", + "DefaultCollisionPropertiesCfg", + "NewtonCollisionPropertiesCfg", + "RigidBodyMaterialCfg", + "NewtonRigidBodyMaterialCfg", + "RigidBodyPhysicsCfg", + "ObjectBaseCfg", + "LightCfg", + "RigidObjectCfg", + "DeformableObjectCfg", + "VolumeDeformableObjectCfg", + "SoftObjectCfg", + "SurfaceDeformableObjectCfg", + "ClothObjectCfg", + "RigidObjectGroupCfg", + "RigidConstraintCfg", + "SoftbodyVoxelAttributesCfg", + "SoftbodyPhysicalAttributesCfg", + "ClothPhysicalAttributesCfg", + "ArticulationRootPropertiesCfg", + "LinkPhysicsOverrideCfg", + "link_attrs_from_dict", + "JointDrivePropertiesCfg", + "NewtonJointDrivePropertiesCfg", + "ArticulationCfg", + "URDFCfg", + "RobotCfg", + "RobotPresetCfg", + "RobotWorkspaceCfg", + "get_data_path", +] diff --git a/embodichain/lab/sim/cfg/articulation.py b/embodichain/lab/sim/cfg/articulation.py new file mode 100644 index 000000000..cd54c6440 --- /dev/null +++ b/embodichain/lab/sim/cfg/articulation.py @@ -0,0 +1,524 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Articulation-root, per-link, joint, and articulation configuration.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import MISSING, fields +import numbers +from typing import Any, Dict, List, Literal, Sequence + +import numpy as np +import torch + +from embodichain.utils import configclass, is_configclass, logger + +from .asset import AssetPhysicsMode, ObjectBaseCfg, _resolve_asset_physics_mode +from .rigid import ( + RigidBodyPhysicsCfg, + _rigid_body_physics_from_dict, +) + + +def _normalize_joint_target_mode(value: object) -> int: + """Normalize a portable joint target mode to its backend integer value.""" + if isinstance(value, str): + normalized = value.replace("-", "_").lower() + modes = { + "none": 0, + "position": 1, + "velocity": 2, + "position_velocity": 3, + "effort": 4, + } + if normalized not in modes: + raise ValueError( + f"Unsupported joint target mode {value!r}; expected one of " + f"{tuple(modes)}." + ) + return modes[normalized] + if isinstance(value, numbers.Integral) and not isinstance(value, bool): + mode = int(value) + if 0 <= mode <= 4: + return mode + raise ValueError("Joint target-mode integers must be in [0, 4].") + raise TypeError("Joint target mode must be a string or an integer in [0, 4].") + + +@configclass +class ArticulationRootPropertiesCfg: + """Articulation-root properties shared by robot definitions. + + ``fixed_base`` and ``self_collision_enabled`` are consumed by both + backends. ``sleep_threshold`` and the solver-iteration fields are supported + only by the Default backend and are ignored by Newton. By default, the + articulation root is fixed and self-collision is disabled. Explicit + ``None`` values preserve the source value or backend/import default. + """ + + fixed_base: bool | None = True + """Whether the articulation root is rigidly fixed to the world frame. + + Set to ``None`` to preserve the source value or backend/import default. + """ + + self_collision_enabled: bool | None = False + """Whether non-filtered link pairs in the articulation may self-collide. + + Newton may still apply source-authored or Spawn-owned filtering to adjacent + parent-child bodies. Set to ``None`` to preserve the source value or + backend/import default. + """ + + sleep_threshold: float | None = None + """Default-only articulation sleep threshold; Newton ignores this field.""" + + min_position_iters: int | None = None + """Default-only minimum root position-solver iterations (1 to 255).""" + + min_velocity_iters: int | None = None + """Default-only minimum root velocity-solver iterations (0 to 255).""" + + def __post_init__(self) -> None: + """Require the two values consumed by the atomic Default setter.""" + if (self.min_position_iters is None) != (self.min_velocity_iters is None): + raise ValueError( + "Articulation-root min_position_iters and min_velocity_iters " + "must be configured together." + ) + + @classmethod + def from_dict( + cls, + init_dict: Mapping[str, Any], + ) -> ArticulationRootPropertiesCfg: + """Parse articulation-root properties without a backend subtype.""" + return cls(**dict(init_dict)) + + +_REMOVED_ARTICULATION_CFG_FIELDS = { + "fix_base": "root_props.fixed_base", + "disable_self_collision": ( + "root_props.self_collision_enabled (invert the old boolean)" + ), + "sleep_threshold": "root_props.sleep_threshold", + "min_position_iters": "root_props.min_position_iters", + "min_velocity_iters": "root_props.min_velocity_iters", + "articulation_props": "root_props", + "drive_pros": "joint_drive_props", + "joint_props": "joint_drive_props", +} + + +def _raise_removed_articulation_cfg_fields(init_dict: Mapping[str, Any]) -> None: + """Reject removed flat articulation fields with actionable replacements.""" + removed = _REMOVED_ARTICULATION_CFG_FIELDS.keys() & init_dict.keys() + if not removed: + return + replacements = ", ".join( + f"{name} -> {_REMOVED_ARTICULATION_CFG_FIELDS[name]}" + for name in sorted(removed) + ) + raise ValueError(f"Removed ArticulationCfg fields: {replacements}.") + + +@configclass +class LinkPhysicsOverrideCfg: + """Partial physics overlay for a selected set of articulation links. + + Regex/control-group resolution happens before Spawn updates exact source + link names. A link may match only one override group. + """ + + link_names_expr: list[str] = MISSING + """Regular expressions matched against complete source link names.""" + + attrs: RigidBodyPhysicsCfg = RigidBodyPhysicsCfg() + """Partial grouped overlay for the selected links. + + Configure source-inertia recomputation through + :attr:`RigidBodyPhysicsCfg.mass_props`. + """ + + @classmethod + def from_dict(cls, init_dict: Dict[str, Any]) -> LinkPhysicsOverrideCfg: + """Initialize the configuration from a dictionary.""" + if "replace_inertial" in init_dict: + raise ValueError( + "LinkPhysicsOverrideCfg.replace_inertial was removed; use " + "attrs.mass_props.recompute_inertia instead." + ) + cfg = cls() + for key, value in init_dict.items(): + if key == "attrs" and isinstance(value, dict): + setattr(cfg, key, _rigid_body_physics_from_dict(value)) + elif hasattr(cfg, key): + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + return cfg + + +def link_attrs_from_dict( + value: dict[str, Any], +) -> dict[str, LinkPhysicsOverrideCfg]: + """Parse a ``link_attrs`` mapping from YAML/JSON-style dicts.""" + link_attrs: dict[str, LinkPhysicsOverrideCfg] = {} + for group_name, group_cfg in value.items(): + if isinstance(group_cfg, LinkPhysicsOverrideCfg): + link_attrs[group_name] = group_cfg + elif isinstance(group_cfg, dict): + link_attrs[group_name] = LinkPhysicsOverrideCfg.from_dict(group_cfg) + else: + raise TypeError( + f"link_attrs['{group_name}'] must be a dict or " + f"LinkPhysicsOverrideCfg, got {type(group_cfg)}." + ) + return link_attrs + + +@configclass +class JointDrivePropertiesCfg: + """Portable joint-drive and joint-dynamics properties. + + A scalar applies to every resolved joint. A dictionary maps exact joint + names, full-match regular expressions, or robot control-part names to + values; exact/regex rules override broader control-part rules. ``None`` + preserves source/backend ownership of a field. + + ``drive_type`` retains the Default drive response (force, acceleration, or + disabled), while ``target_mode`` selects the commanded target components. + Spawn resolves the two concepts before lowering them to the Default drive + descriptor and Newton ``JointDofConfig``. + + Effort and velocity limits, friction, and armature share the same matching + rules and descriptor compilation boundary as the actuator target and gains. + + Newton stores all fields in the model, but individual solvers may ignore + limits, friction, armature, or target modes; consult the `Newton solver + feature matrix + `_. + """ + + drive_type: Literal["force", "acceleration", "none"] | None = None + """Joint drive type to apply. + + On the Default backend, ``"force"`` applies a force/torque drive, + ``"acceleration"`` applies a mass-independent acceleration drive, and + ``"none"`` disables the drive. Newton has no acceleration-drive + equivalent. Unless :attr:`target_mode` is explicit, ``"force"`` and + ``"acceleration"`` select ``"position_velocity"`` while ``"none"`` + selects ``"none"``. + """ + + target_mode: ( + Literal[ + "none", + "position", + "velocity", + "position_velocity", + "effort", + ] + | Dict[ + str, + Literal[ + "none", + "position", + "velocity", + "position_velocity", + "effort", + ] + | int, + ] + | int + | None + ) = None + """Portable actuator target mode, as a scalar or joint-rule mapping. + + Accepted names and integer values are ``"none"``/``0`` (passive), + ``"position"``/``1``, ``"velocity"``/``2``, + ``"position_velocity"``/``3``, and ``"effort"``/``4``. Default emulates + these modes through its drive mode and effective gains. Newton authors the + corresponding ``JointTargetMode``; solvers without native target-mode + support use deterministic gain-based fallbacks where possible. + """ + + stiffness: Dict[str, float] | float | None = None + """Proportional position gain of the joint drive. + + The unit depends on the joint model: + + * For linear joints, the unit is kg-m/s^2 (N/m). + * For angular joints, the unit is kg-m^2/s^2/rad (N-m/rad). + """ + + damping: Dict[str, float] | float | None = None + """Derivative velocity gain of the joint drive. + + The unit depends on the joint model: + + * For linear joints, the unit is kg-m/s (N-s/m). + * For angular joints, the unit is kg-m^2/s/rad (N-m-s/rad). + """ + + max_effort: Dict[str, float] | float | None = None + """Maximum drive effort [N for prismatic, N*m for revolute joints]. + + The value is authored for both backends, but the selected Newton solver may + not enforce it. + """ + + max_velocity: Dict[str, float] | float | None = None + """Maximum joint speed [m/s for prismatic, rad/s for revolute joints]. + + The value is authored for both backends, but support is solver-dependent in + Newton. + """ + + friction: Dict[str, float] | float | None = None + """Passive friction value applied along the joint degree of freedom. + + Interpretation and enforcement are backend/solver-dependent. + """ + + armature: Dict[str, float] | float | None = None + """Artificial inertia added to the joint-space diagonal. + + Units depend on the joint model: + + * For prismatic (linear) joints, the unit is mass [kg]. + * For revolute (angular) joints, the unit is mass * scene_length^2 [kg-m^2]. + + Armature changes the physical model and should normally reflect actuator or + gearbox inertia. Newton solver support varies. + """ + + def _resolve_modes(self) -> tuple[object, str | None]: + """Resolve the target default implied by the original drive type.""" + target_mode = self.target_mode + drive_type = self.drive_type + if drive_type not in {None, "force", "acceleration", "none"}: + raise ValueError(f"Unsupported joint drive type {drive_type!r}.") + if target_mode is None: + target_mode = { + None: None, + "force": "position_velocity", + "acceleration": "position_velocity", + "none": "none", + }[drive_type] + return target_mode, drive_type + + @classmethod + def from_dict( + cls, + init_dict: Dict[str, Any], + *, + defaults: JointDrivePropertiesCfg | None = None, + ) -> JointDrivePropertiesCfg: + """Initialize the configuration from a dictionary. + + Args: + init_dict: Joint-drive properties to override. + defaults: Optional base properties whose unspecified values are + preserved. If omitted, the class defaults are used. + + Returns: + Parsed joint-drive properties. + """ + data = dict(init_dict) + backend = str(data.pop("backend", "common")).replace("-", "_").lower() + wants_newton = backend == "newton" + if backend not in {"common", "default", "newton"}: + raise ValueError( + "joint_drive_props.backend must be 'common', 'default', or 'newton', " + f"got {backend!r}." + ) + if wants_newton and not isinstance(defaults, NewtonJointDrivePropertiesCfg): + cfg = NewtonJointDrivePropertiesCfg() + if defaults is not None: + for item in fields(JointDrivePropertiesCfg): + setattr(cfg, item.name, getattr(defaults, item.name)) + else: + cfg = defaults.copy() if defaults is not None else cls() + for key, value in data.items(): + if hasattr(cfg, key): + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + return cfg + + def to_dict(self) -> dict[str, Any]: + """Serialize joint properties with their backend subtype.""" + data = {item.name: getattr(self, item.name) for item in fields(self)} + if isinstance(self, NewtonJointDrivePropertiesCfg): + data["backend"] = "newton" + return data + + +@configclass +class NewtonJointDrivePropertiesCfg(JointDrivePropertiesCfg): + """Compatibility subtype for serialized Newton joint-drive configs. + + ``target_mode`` is now portable and lives on + :class:`JointDrivePropertiesCfg`. The subtype remains so existing + ``backend="newton"`` dictionaries and round trips retain their type; new + robot definitions should use the common class. + """ + + +@configclass +class ArticulationCfg(ObjectBaseCfg): + """Configuration for an articulation asset in the simulation. + + This class extends the base asset configuration to include specific properties for articulations, + such as joint drive properties, physical attributes. + """ + + fpath: str = None + """Path to the articulation asset file.""" + + body_scale: tuple | list = (1.0, 1.0, 1.0) + """Scale of the articulation in the simulation world frame.""" + + compute_uv: bool = False + """Whether to compute the UV mapping for the articulation link. + + Currently, the uv mapping is computed for each link with projection uv mapping method. + """ + + asset_physics_mode: AssetPhysicsMode = "preserve" + """How source-authored articulation physics is handled. + + ``"preserve"`` keeps link, joint-drive, and joint-limit properties from + either USD or URDF. ``"overlay"`` applies only explicitly configured + values after the source has been resolved. + + Import policy such as root fixation and body scale remains controlled by + :attr:`root_props` and :attr:`body_scale`. + """ + + attrs: RigidBodyPhysicsCfg = RigidBodyPhysicsCfg() + """Physical attributes for all links. We use default mass from the USD/URDF file if available. + The mass and density in attrs will only be used if specified. + """ + + link_attrs: dict[str, LinkPhysicsOverrideCfg] | None = None + """Named per-link physics override groups keyed by regex on link names. + + Each group applies :attr:`LinkPhysicsOverrideCfg.attrs` on top of :attr:`attrs` for + matched links only. A link must not match more than one group. + """ + + root_props: ArticulationRootPropertiesCfg = ArticulationRootPropertiesCfg() + """Grouped articulation-root properties. + + Fixed-base and self-collision intent is portable. Root sleep and solver + iterations are Default-only fields and are ignored by Newton. The portable + fields default to a fixed base with self-collision disabled. Set either + field to ``None`` to preserve an authored USD/backend value; URDF imports + then use the established fixed-base, self-collision-off defaults. + """ + + joint_drive_props: JointDrivePropertiesCfg | None = None + """Optional joint-drive and joint-dynamics overrides. + + ``None`` preserves source drive properties. Individual ``None`` fields in + a provided config also preserve the corresponding source values. + """ + + init_qpos: torch.Tensor | np.ndarray | Sequence[float] = None + """Initial joint positions of the articulation. + + If None, the joint positions will be set to zero. + If provided, it should be an array of shape ``(num_dofs,)``. + """ + + qpos_limits: ( + torch.Tensor + | np.ndarray + | Sequence[Sequence[float]] + | Dict[str, List[float]] + | None + ) = None + """Override joint position limits of the articulation. + + If None, the joint position limits from the asset file (URDF/USD) are used. + If provided as a tensor/array of shape ``(num_dofs, 2)``, it is applied in + flattened source-resolved DOF order before the backend model is built. + If provided as a dictionary, keys are joint names or regular expressions and + values are ``[min, max]`` limits. + + This field replaces the asset limits for the articulation and can be used to + either tighten or expand the allowed range. + """ + + build_pk_chain: bool = True + """Whether to build pytorch-kinematics chain for forward kinematics and jacobian computation.""" + + def resolve_asset_physics_mode(self) -> AssetPhysicsMode: + """Return the effective file-backed physics policy.""" + return _resolve_asset_physics_mode(self.asset_physics_mode) + + @classmethod + def from_dict( + cls, init_dict: Dict[str, str | float | tuple | dict] + ) -> ArticulationCfg: + """Initialize the configuration from a dictionary.""" + _raise_removed_articulation_cfg_fields(init_dict) + cfg = cls() + for key, value in init_dict.items(): + if key == "link_attrs" and isinstance(value, dict): + cfg.link_attrs = link_attrs_from_dict(value) + elif key == "attrs" and isinstance(value, Mapping): + cfg.attrs = _rigid_body_physics_from_dict(value) + elif key == "joint_drive_props" and isinstance(value, Mapping): + cfg.joint_drive_props = JointDrivePropertiesCfg.from_dict( + dict(value), + defaults=cfg.joint_drive_props, + ) + elif hasattr(cfg, key): + attr = getattr(cfg, key) + if is_configclass(attr): + setattr(cfg, key, attr.from_dict(value)) + else: + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + + if cfg.init_local_pose is None: + from scipy.spatial.transform import Rotation as R + + T = np.eye(4) + T[:3, 3] = np.array(cfg.init_pos) + T[:3, :3] = R.from_euler("xyz", np.deg2rad(cfg.init_rot)).as_matrix() + cfg.init_local_pose = T + else: + from scipy.spatial.transform import Rotation as R + + cfg.init_pos = tuple(cfg.init_local_pose[:3, 3]) + cfg.init_rot = tuple( + R.from_matrix(cfg.init_local_pose[:3, :3]).as_euler("xyz", degrees=True) + ) + + return cfg diff --git a/embodichain/lab/sim/cfg/asset.py b/embodichain/lab/sim/cfg/asset.py new file mode 100644 index 000000000..36b77bf67 --- /dev/null +++ b/embodichain/lab/sim/cfg/asset.py @@ -0,0 +1,103 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Base asset configuration and file-backed physics policy.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Dict, Literal + +import numpy as np + +from embodichain.utils import configclass, is_configclass, logger + +AssetPhysicsMode = Literal["preserve", "overlay"] +"""Policy for applying EmbodiChain physics to a file-backed asset.""" + + +def _resolve_asset_physics_mode( + mode: AssetPhysicsMode, +) -> AssetPhysicsMode: + """Validate and return a source-agnostic asset-physics policy.""" + if mode not in ("preserve", "overlay"): + raise ValueError( + f"asset_physics_mode must be 'preserve' or 'overlay', got {mode!r}." + ) + return mode + + +@configclass +class ObjectBaseCfg: + """Base configuration for an asset in the simulation. + + This class defines the basic properties of an asset, such as its type, initial state, and collision group. + It is used as a base class for specific asset configurations. + """ + + uid: str | None = None + + init_pos: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Position of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" + + init_rot: tuple[float, float, float] = (0.0, 0.0, 0.0) + """Euler angles (in degree) of the root in simulation world frame. Defaults to (0.0, 0.0, 0.0).""" + + init_local_pose: np.ndarray | None = None + """4x4 transformation matrix of the root in local frame. If specified, it will override init_pos and init_rot.""" + + @classmethod + def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> ObjectBaseCfg: + """Initialize the configuration from a dictionary.""" + cfg = cls() # Create a new instance of the class (cls) + for key, value in init_dict.items(): + if hasattr(cfg, key): + attr = getattr(cfg, key) + if key == "attrs" and isinstance(value, Mapping): + # Keep the base module independent of rigid schemas at + # import time; only rigid-derived configs expose this key. + from .rigid import _rigid_body_physics_from_dict + + setattr(cfg, key, _rigid_body_physics_from_dict(value)) + elif is_configclass(attr): + setattr( + cfg, key, attr.from_dict(value) + ) # Call from_dict on the attribute + else: + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + + # Automatically infer init_local_pose if not provided + if cfg.init_local_pose is None: + # If only init_pos or init_rot are provided, generate the 4x4 pose matrix + from scipy.spatial.transform import Rotation as R + + T = np.eye(4) + T[:3, 3] = np.array(cfg.init_pos) + T[:3, :3] = R.from_euler("xyz", np.deg2rad(cfg.init_rot)).as_matrix() + cfg.init_local_pose = T + else: + # If only init_local_pose is provided, extract init_pos and init_rot + from scipy.spatial.transform import Rotation as R + + T = np.array(cfg.init_local_pose) + cfg.init_pos = tuple(T[:3, 3]) + cfg.init_rot = tuple(R.from_matrix(T[:3, :3]).as_euler("xyz", degrees=True)) + + return cfg diff --git a/embodichain/lab/sim/cfg/deformable.py b/embodichain/lab/sim/cfg/deformable.py new file mode 100644 index 000000000..e9a4de164 --- /dev/null +++ b/embodichain/lab/sim/cfg/deformable.py @@ -0,0 +1,328 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Deformable-body physical and object configuration.""" + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Literal + +from dexsim.types import ( + ClothBodyAttr, + SoftBodyAttr, + SoftBodyMaterialModel, + VoxelConfig, +) + +from embodichain.utils import configclass + +from ..shapes import MeshCfg +from .asset import ObjectBaseCfg + + +@configclass +class SoftbodyVoxelAttributesCfg: + # voxel config + triangle_remesh_resolution: int = 8 + """Resolution to remesh the softbody mesh before building physics collision mesh.""" + + triangle_simplify_target: int = 0 + """Simplify mesh faces to target value. Do nothing if this value is zero.""" + + # TODO: this value will be automatically computed with simulation_mesh_resolution and mesh scale. + maximal_edge_length: float = 0 + # """To shorten edges that are too long, additional points get inserted at their center leading to a subdivision of the input mesh. Do nothing if this value is zero.""" + + simulation_mesh_resolution: int = 8 + """Resolution to build simulation voxelize textra mesh. This value must be greater than 0.""" + + simulation_mesh_output_obj: bool = False + """Whether to output the simulation mesh as an obj file for debugging.""" + + def attr(self) -> VoxelConfig: + """Convert to dexsim VoxelConfig""" + attr = VoxelConfig() + attr.triangle_remesh_resolution = self.triangle_remesh_resolution + attr.maximal_edge_length = self.maximal_edge_length + attr.simulation_mesh_resolution = self.simulation_mesh_resolution + attr.triangle_simplify_target = self.triangle_simplify_target + return attr + + +@configclass +class SoftbodyPhysicalAttributesCfg: + # material properties + youngs: float = 1e6 + """Young's modulus (higher = stiffer).""" + + poissons: float = 0.45 + """Poisson's ratio (higher = closer to incompressible).""" + + dynamic_friction: float = 0.0 + """Dynamic friction coefficient.""" + + elasticity_damping: float = 0.0 + """Elasticity damping factor.""" + + # soft body properties + material_model: SoftBodyMaterialModel = SoftBodyMaterialModel.CO_ROTATIONAL + """Material constitutive model.""" + + # --- Mode / collision switches --- + enable_kinematic: bool = False + """If True, (partially) kinematic behavior is enabled.""" + + enable_ccd: bool = False + """Enable continuous collision detection (CCD).""" + + enable_self_collision: bool = False + """Enable self-collision handling.""" + + has_gravity: bool = True + """Whether the soft body is affected by gravity.""" + + # --- Self-collision & simplification parameters --- + self_collision_stress_tolerance: float = 0.9 + """Stress tolerance threshold for self-collision constraints.""" + + collision_mesh_simplification: bool = True + """Whether to simplify the collision mesh for self-collision.""" + + self_collision_filter_distance: float = 0.1 + """Distance threshold below which vertex pairs may be filtered from self-collision checks.""" + + # --- Damping, sleep & settling --- + vertex_velocity_damping: float = 0.005 + """Per-vertex velocity damping.""" + + linear_damping: float = 0.0 + """Global linear damping applied to the soft body.""" + + sleep_threshold: float = 0.05 + """Velocity/energy threshold below which the soft body can go to sleep.""" + + settling_threshold: float = 0.1 + """Threshold used to decide convergence/settling state.""" + + settling_damping: float = 10.0 + """Additional damping applied during settling phase.""" + + # --- Mass / density & velocity limits --- + mass: float = -1.0 + """Total mass of the soft body. If set to a negative value, density will be used to compute mass.""" + + density: float = 1000.0 + """Material density in kg/m^3.""" + + max_depenetration_velocity: float = 1e6 + """Maximum velocity used to resolve penetrations. Must be larger than zero.""" + + max_velocity: float = 100 + """Clamp for linear (or vertex) velocity. If set to zero, the limit is ignored.""" + + # --- Solver iteration counts --- + min_position_iters: int = 4 + """Minimum solver iterations for position correction.""" + + min_velocity_iters: int = 1 + """Minimum solver iterations for velocity updates.""" + + def attr(self) -> SoftBodyAttr: + attr = SoftBodyAttr() + attr.youngs = self.youngs + attr.poissons = self.poissons + attr.dynamic_friction = self.dynamic_friction + attr.elasticity_damping = self.elasticity_damping + attr.material_model = self.material_model + attr.enable_kinematic = self.enable_kinematic + attr.enable_ccd = self.enable_ccd + attr.enable_self_collision = self.enable_self_collision + attr.has_gravity = self.has_gravity + attr.self_collision_stress_tolerance = self.self_collision_stress_tolerance + attr.collision_mesh_simplification = self.collision_mesh_simplification + attr.vertex_velocity_damping = self.vertex_velocity_damping + attr.mass = self.mass + attr.density = self.density + attr.max_depenetration_velocity = self.max_depenetration_velocity + attr.max_velocity = self.max_velocity + attr.self_collision_filter_distance = self.self_collision_filter_distance + attr.linear_damping = self.linear_damping + attr.sleep_threshold = self.sleep_threshold + attr.settling_threshold = self.settling_threshold + attr.settling_damping = self.settling_damping + attr.min_position_iters = self.min_position_iters + attr.min_velocity_iters = self.min_velocity_iters + return attr + + +@configclass +class ClothPhysicalAttributesCfg: + # material properties + youngs: float = 1e10 + """Young's modulus (higher = stiffer).""" + + poissons: float = 0.3 + """Poisson's ratio.""" + + dynamic_friction: float = 0.5 + """Dynamic friction coefficient.""" + + elasticity_damping: float = 0.0 + """Elasticity damping factor.""" + + thickness: float = 0.001 + """Cloth thickness (m).""" + + bending_stiffness: float = 0.00001 + """Bending stiffness.""" + + bending_damping: float = 0.0 + """Bending damping.""" + + # cloth body properties + enable_kinematic: bool = False + """If True, (partially) kinematic behavior is enabled.""" + + enable_ccd: bool = True + """Enable continuous collision detection (CCD).""" + + enable_self_collision: bool = False + """Enable self-collision handling.""" + + has_gravity: bool = True + """Whether the cloth is affected by gravity.""" + + self_collision_stress_tolerance: float = 0.9 + """Stress tolerance threshold for self-collision constraints.""" + + collision_mesh_simplification: bool = True + """Whether to simplify the collision mesh for self-collision.""" + + vertex_velocity_damping: float = 0.005 + """Per-vertex velocity damping.""" + + mass: float = -1.0 + """Total mass of the cloth. If negative, density is used to compute mass.""" + + density: float = 1.0 + """Material density in kg/m^3.""" + + max_depenetration_velocity: float = 1e6 + """Maximum velocity used to resolve penetrations.""" + + max_velocity: float = 100.0 + """Clamp for linear (or vertex) velocity.""" + + self_collision_filter_distance: float = 0.1 + """Distance threshold for filtering self-collision vertex pairs.""" + + linear_damping: float = 0.05 + """Global linear damping applied to the cloth.""" + + sleep_threshold: float = 0.05 + """Velocity/energy threshold below which the cloth can go to sleep.""" + + settling_threshold: float = 0.1 + """Threshold used to decide convergence/settling state.""" + + settling_damping: float = 10.0 + """Additional damping applied during settling phase.""" + + min_position_iters: int = 4 + """Minimum solver iterations for position correction.""" + + min_velocity_iters: int = 1 + """Minimum solver iterations for velocity updates.""" + + def attr(self) -> ClothBodyAttr: + """Convert to dexsim ClothBodyAttr.""" + attr = ClothBodyAttr() + attr.youngs = self.youngs + attr.poissons = self.poissons + attr.dynamic_friction = self.dynamic_friction + attr.elasticity_damping = self.elasticity_damping + attr.thickness = self.thickness + attr.bending_stiffness = self.bending_stiffness + attr.bending_damping = self.bending_damping + attr.enable_kinematic = self.enable_kinematic + attr.enable_ccd = self.enable_ccd + attr.enable_self_collision = self.enable_self_collision + attr.has_gravity = self.has_gravity + attr.self_collision_stress_tolerance = self.self_collision_stress_tolerance + attr.collision_mesh_simplification = self.collision_mesh_simplification + attr.vertex_velocity_damping = self.vertex_velocity_damping + attr.mass = self.mass + attr.density = self.density + attr.max_depenetration_velocity = self.max_depenetration_velocity + attr.max_velocity = self.max_velocity + attr.self_collision_filter_distance = self.self_collision_filter_distance + attr.linear_damping = self.linear_damping + attr.sleep_threshold = self.sleep_threshold + attr.settling_threshold = self.settling_threshold + attr.settling_damping = self.settling_damping + attr.min_position_iters = self.min_position_iters + attr.min_velocity_iters = self.min_velocity_iters + return attr + + +@configclass +class DeformableObjectCfg(ObjectBaseCfg): + """Common configuration contract for one deformable asset. + + Concrete volume and surface configurations retain their native DexSim + properties. The discriminator is explicit so manager and visualization + code do not need to infer topology from a mesh or material type. + """ + + deformable_type: Literal["volume", "surface"] = MISSING + """Physical topology represented by the asset.""" + + shape: MeshCfg = MeshCfg() + """Render and source-mesh configuration.""" + + +@configclass +class VolumeDeformableObjectCfg(DeformableObjectCfg): + """Configuration for a volume deformable backed by DexSim ``SoftBody``.""" + + deformable_type: Literal["volume"] = "volume" + + voxel_attr: SoftbodyVoxelAttributesCfg = SoftbodyVoxelAttributesCfg() + """Tetrahedral simulation-mesh voxelization attributes.""" + + physical_attr: SoftbodyPhysicalAttributesCfg = SoftbodyPhysicalAttributesCfg() + """DexSim volume-deformable physical attributes.""" + + +@configclass +class SoftObjectCfg(VolumeDeformableObjectCfg): + """Compatibility name for :class:`VolumeDeformableObjectCfg`.""" + + +@configclass +class SurfaceDeformableObjectCfg(DeformableObjectCfg): + """Configuration for a surface deformable backed by DexSim ``ClothBody``.""" + + deformable_type: Literal["surface"] = "surface" + + physical_attr: ClothPhysicalAttributesCfg = ClothPhysicalAttributesCfg() + """DexSim surface-deformable physical attributes.""" + + +@configclass +class ClothObjectCfg(SurfaceDeformableObjectCfg): + """Compatibility name for :class:`SurfaceDeformableObjectCfg`.""" diff --git a/embodichain/lab/sim/cfg/rigid.py b/embodichain/lab/sim/cfg/rigid.py new file mode 100644 index 000000000..aec7aa184 --- /dev/null +++ b/embodichain/lab/sim/cfg/rigid.py @@ -0,0 +1,663 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Rigid-body mass, collision, material, and backend property schemas.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import fields +from typing import Any, Sequence + +import numpy as np +from dexsim.types import PhysicalAttr + +from embodichain.utils import configclass +from embodichain.utils.math import convert_quat + + +@configclass +class MassPropertiesCfg: + """Backend-neutral rigid-body mass properties. + + ``None`` means that the source asset or selected backend keeps ownership of + that value. For a non-static body, explicit inertia requires a positive + mass. A source-backed body retains authored inertia unless + :attr:`recompute_inertia` is enabled; procedural or recomputed bodies derive + inertia from collision geometry and the effective mass or density. Static + bodies omit all mass properties during Spawn compilation. + """ + + mass: float | None = None + """Rigid-body mass [kg]. + + A positive value takes precedence over :attr:`density`. Zero explicitly + selects density-based derivation and therefore requires a positive density. + Negative values are invalid. + """ + + density: float | None = None + """Uniform density used to derive mass properties from collision shapes [kg/m^3]. + + The value must be positive and is ignored when :attr:`mass` is positive. + """ + + inertia: Sequence[float] | np.ndarray | None = None + """Inertia about the center of mass [kg*m^2]. + + Supply either three positive principal moments or a symmetric, + positive-definite 3-by-3 tensor in the body frame. Explicit inertia is + accepted only together with a positive :attr:`mass`. For one definition + shared by both backends, prefer principal moments plus + :attr:`com_quaternion`; the current Default adapter consumes the principal- + moment representation, while Newton can retain a full tensor. + """ + + recompute_inertia: bool | None = None + """Whether collision geometry should replace source-authored inertia. + + ``True`` discards source inertia so the backend recomputes it from the + collision geometry and effective mass or density. ``False`` preserves the + source inertia. ``None`` inherits an outer rigid-body overlay and otherwise + behaves like ``False``. Explicit :attr:`inertia` cannot be combined with + recomputation. + """ + + com_position: Sequence[float] | np.ndarray | None = None + """Center-of-mass position expressed in the rigid body's local frame [m].""" + + com_quaternion: Sequence[float] | np.ndarray | None = None + """Orientation of the center-of-mass/inertia frame in ``xyzw`` order. + + Spawn normalizes the quaternion and converts it to the backend descriptor's + ``wxyz`` convention. A zero quaternion is invalid. + """ + + +@configclass +class DefaultRigidBodyPropertiesCfg: + """Rigid-body properties consumed only by the Default backend. + + Every field defaults to ``None`` so a partial overlay preserves an authored + USD/URDF value or the backend default. + """ + + linear_damping: float | None = None + """Non-negative damping coefficient applied to linear velocity.""" + + angular_damping: float | None = None + """Non-negative damping coefficient applied to angular velocity.""" + + has_gravity: bool | None = None + """Whether world gravity accelerates this body.""" + + max_linear_velocity: float | None = None + """Maximum rigid-body linear speed [m/s].""" + + max_angular_velocity: float | None = None + """Maximum rigid-body angular speed [rad/s].""" + + max_depenetration_velocity: float | None = None + """Maximum separation speed introduced to resolve penetration [m/s].""" + + retain_acceleration: bool | None = None + """Whether accumulated acceleration is retained across simulation steps.""" + + enable_ccd: bool | None = None + """Whether continuous collision detection is enabled for this body. + + Scene-level CCD must also be enabled through + :attr:`DefaultPhysicsCfg.enable_ccd`. + """ + + min_position_iters: int | None = None + """Minimum number of position-solver iterations for this body (1 to 255).""" + + min_velocity_iters: int | None = None + """Minimum number of velocity-solver iterations for this body (0 to 255).""" + + sleep_threshold: float | None = None + """Mass-normalized kinetic-energy threshold below which the body may sleep.""" + + +@configclass +class CollisionPropertiesCfg: + """Collision-shape properties with identical intent across both backends. + + ``None`` leaves the corresponding source/backend value unchanged. The + contact envelope is expressed once with Default-backend terminology and is + compiled to Newton's ``margin``/``gap`` representation at the Spawn + boundary. Mesh approximation and SDF cooking belong to + :class:`~embodichain.lab.sim.shapes.MeshCollisionCfg`. + """ + + collision_enabled: bool | None = None + """Whether the shape participates in rigid shape-shape collision. + + On Newton this maps to ``ShapeConfig.has_shape_collision``. ``None`` + preserves the source/backend value. + """ + + contact_offset: float | None = None + """Per-shape distance at which contact generation starts [m]. + + The pair threshold is the sum of both shapes' contact offsets. This value + must be non-negative and no smaller than :attr:`rest_offset`. Default + consumes it directly; Newton compiles it together with :attr:`rest_offset` + to ``gap = contact_offset - rest_offset``. + """ + + rest_offset: float | None = None + """Per-shape target separation at rest [m]. + + Pairwise rest separation is the sum of both shapes' values. Positive + values leave an air gap, zero targets touching surfaces, and negative + values permit limited penetration. Default consumes it directly; Newton + maps it to ``margin``. + """ + + +@configclass +class DefaultCollisionPropertiesCfg(CollisionPropertiesCfg): + """Collision-solver properties consumed only by the Default backend. + + ``contact_offset`` and ``rest_offset`` now live on + :class:`CollisionPropertiesCfg` because both backends consume their intent. + """ + + torsional_patch_radius: float | None = None + """Contact-patch radius used to approximate torsional friction [m].""" + + min_torsional_patch_radius: float | None = None + """Minimum contact-patch radius used for torsional friction [m].""" + + disable_strong_friction: bool | None = None + """Whether to disable Default-backend strong-friction contact anchoring.""" + + +@configclass +class NewtonCollisionPropertiesCfg(CollisionPropertiesCfg): + """Newton-native contact-envelope properties. + + Mesh construction belongs to ``MeshCfg.collision``; filtering, visual, and + semantic-site policies are deliberately not part of rigid-body physics. + + See `Newton Shape Configuration + `_. + """ + + margin: float | None = None + """Outward collision-surface offset [m]. + + Margins from both shapes are added. They determine where contact is placed + and also affect inertia/SDF handling for hollow shapes. + """ + + gap: float | None = None + """Additional contact-detection distance outside :attr:`margin` [m]. + + Gaps from both shapes are added. Broad phase expands each shape by + ``margin + gap``; increasing the gap detects approaching contact earlier. + """ + + +@configclass +class RigidBodyMaterialCfg: + """Common rigid-contact material intent. + + All fields use sparse-overlay semantics: ``None`` preserves the source or + backend default. The Default backend consumes all three values. Newton + has one Coulomb friction coefficient, so it maps :attr:`dynamic_friction` + to ``ShapeConfig.mu`` and currently has no separate static-friction input; + restitution is consumed only by Newton solvers that support it. + """ + + static_friction: float | None = None + """Static friction coefficient used before tangential slip begins. + + This is currently consumed only by the Default backend. + """ + + dynamic_friction: float | None = None + """Sliding friction coefficient. + + The Default backend uses it as dynamic friction; Newton uses it as its + single Coulomb friction coefficient ``mu``. + """ + + restitution: float | None = None + """Coefficient of restitution, where zero is inelastic and one is elastic. + + The active backend/solver may further restrict or ignore restitution. + """ + + +@configclass +class NewtonRigidBodyMaterialCfg(RigidBodyMaterialCfg): + """Newton contact-material extensions. + + Solver support differs by field. Semi-implicit and Featherstone consume + ``ke``, ``kd``, ``kf``, ``ka``, ``mu``, and ``kh``; MuJoCo Warp consumes + ``ke``, ``kd``, ``mu``, ``kh``, and the torsional/rolling coefficients; + XPBD consumes ``mu``, restitution, and torsional/rolling friction. DexSim + warns when an explicitly changed contact field is ignored by the selected + solver. + """ + + ke: float | None = None + """Elastic contact stiffness coefficient.""" + + kd: float | None = None + """Normal contact damping coefficient.""" + + kf: float | None = None + """Tangential/friction damping coefficient.""" + + ka: float | None = None + """Contact adhesion distance [m].""" + + kh: float | None = None + """Hydroelastic contact stiffness used when hydroelastic contact is enabled.""" + + torsional_friction: float | None = None + """Torsional friction coefficient resisting spin at a contact point.""" + + rolling_friction: float | None = None + """Rolling friction coefficient resisting rolling motion.""" + + +_RIGID_PHYSICS_GROUP_FIELDS = frozenset( + { + "mass_props", + "rigid_props", + "collision_props", + "material_props", + } +) + +_REMOVED_RIGID_PHYSICS_GROUP_FIELDS = { + "default_props": "the corresponding polymorphic property slot", + "newton_props": "the corresponding polymorphic property slot", + "mesh_collision_props": "MeshCfg.collision", +} + + +def _default_rigid_props_from_dict( + value: Mapping[str, Any] | object | None, +) -> DefaultRigidBodyPropertiesCfg | None: + """Parse the currently Default-only rigid-body property slot.""" + if value is None or isinstance(value, DefaultRigidBodyPropertiesCfg): + return value + if not isinstance(value, Mapping): + raise TypeError( + "rigid_props must be a mapping or DefaultRigidBodyPropertiesCfg." + ) + data = dict(value) + backend = str(data.pop("backend", "default")).replace("-", "_").lower() + if backend != "default": + raise ValueError( + "rigid_props.backend must be 'default'; Newton currently exposes no " + "body-level property config." + ) + try: + return DefaultRigidBodyPropertiesCfg(**data) + except TypeError as exc: + raise TypeError(f"Invalid rigid_props configuration: {exc}") from exc + + +def _physics_property_cfg_from_dict( + value: Mapping[str, Any] | object | None, + *, + common_type: type, + backend_types: Mapping[str, type], + field_name: str, +) -> object | None: + """Parse one polymorphic rigid-physics property slot.""" + if value is None: + return None + supported_types = (common_type, *backend_types.values()) + if isinstance(value, supported_types): + return value + if not isinstance(value, Mapping): + raise TypeError(f"{field_name} must be a mapping or {common_type.__name__}.") + data = dict(value) + configured_backend = data.pop("backend", None) + if configured_backend is None: + common_fields = {item.name for item in fields(common_type)} + matching_backends = [ + backend + for backend, config_type in backend_types.items() + if ( + {item.name for item in fields(config_type)} - common_fields + ).intersection(data) + ] + if len(matching_backends) > 1: + raise ValueError( + f"{field_name} mixes Default and Newton-only fields; select one " + "backend-specific property config." + ) + backend = matching_backends[0] if matching_backends else "common" + else: + backend = str(configured_backend).replace("-", "_").lower() + config_type = common_type if backend == "common" else backend_types.get(backend) + if config_type is None: + supported_backends = ("common", *backend_types) + raise ValueError( + f"{field_name}.backend must be one of {supported_backends}, got " + f"{backend!r}." + ) + try: + return config_type(**data) + except TypeError as exc: + raise TypeError(f"Invalid {field_name} configuration: {exc}") from exc + + +def _physics_property_cfg_to_dict( + value: object | None, + *, + common_type: type, + backend_types: Mapping[str, type], + field_name: str, +) -> dict[str, Any] | None: + """Serialize one polymorphic property slot with a stable discriminator.""" + if value is None: + return None + backend = next( + ( + name + for name, config_type in backend_types.items() + if isinstance(value, config_type) + ), + None, + ) + if backend is None and type(value) is not common_type: + raise TypeError( + f"Unsupported {field_name} config type {type(value).__name__!r}." + ) + data = dict(value.to_dict()) + if backend is not None: + data["backend"] = backend + return data + + +def _copy_dexsim_physical_attr(source: PhysicalAttr) -> PhysicalAttr: + """Copy a native ``PhysicalAttr`` without relying on pickle support. + + DexSim exposes ``PhysicalAttr`` through a pybind extension object, so + :func:`copy.deepcopy` cannot clone it. Copy scalar fields from the native + mapping and clone its array-valued mass/COM fields explicitly before a + sparse grouped overlay is applied. + """ + copied = PhysicalAttr() + for field_name, value in source.as_dict().items(): + setattr(copied, field_name, value) + for field_name in ("inertia", "com_position", "com_quaternion"): + value = getattr(source, field_name, None) + if value is not None: + setattr(copied, field_name, np.array(value, dtype=np.float32, copy=True)) + return copied + + +@configclass +class RigidBodyPhysicsCfg: + """Grouped rigid-body physics configuration used by Spawn. + + Every nested field defaults to ``None``. With + ``asset_physics_mode="overlay"``, Spawn therefore changes only explicitly + configured values and preserves all other USD/URDF or backend defaults. + Each physical concept has exactly one slot. Dict/YAML input selects a + backend subclass with a local discriminator, while a unique native field + may infer that subclass. Mesh collision construction belongs to + :class:`~embodichain.lab.sim.shapes.MeshCfg`, not this body-physics schema. + """ + + mass_props: MassPropertiesCfg | None = None + """Backend-neutral mass, inertia, COM, and recomputation overrides.""" + + rigid_props: DefaultRigidBodyPropertiesCfg | None = None + """Optional Default-native body properties. + + Newton currently exposes no body-level property group beyond common mass + properties, so there is no empty Newton marker config. + """ + + collision_props: CollisionPropertiesCfg | None = None + """Portable collision envelope plus one optional backend-specific subtype.""" + + material_props: RigidBodyMaterialCfg | None = None + """Portable contact material values plus optional backend-native coefficients.""" + + @classmethod + def from_dict(cls, init_dict: Mapping[str, Any]) -> RigidBodyPhysicsCfg: + """Parse grouped physics properties from a YAML/JSON-style mapping.""" + removed = _REMOVED_RIGID_PHYSICS_GROUP_FIELDS.keys() & init_dict.keys() + if removed: + replacements = ", ".join( + f"{name} -> {_REMOVED_RIGID_PHYSICS_GROUP_FIELDS[name]}" + for name in sorted(removed) + ) + raise ValueError(f"Removed RigidBodyPhysicsCfg fields: {replacements}.") + unknown = set(init_dict) - _RIGID_PHYSICS_GROUP_FIELDS + if unknown: + raise KeyError(f"Unknown RigidBodyPhysicsCfg fields: {sorted(unknown)}") + cfg = cls() + if "mass_props" in init_dict: + value = init_dict["mass_props"] + if value is not None: + if not isinstance(value, (MassPropertiesCfg, Mapping)): + raise TypeError( + "mass_props must be a mapping or MassPropertiesCfg." + ) + cfg.mass_props = ( + value + if isinstance(value, MassPropertiesCfg) + else MassPropertiesCfg(**value) + ) + if "rigid_props" in init_dict: + cfg.rigid_props = _default_rigid_props_from_dict(init_dict["rigid_props"]) + if "collision_props" in init_dict: + cfg.collision_props = _physics_property_cfg_from_dict( + init_dict["collision_props"], + common_type=CollisionPropertiesCfg, + backend_types={ + "default": DefaultCollisionPropertiesCfg, + "newton": NewtonCollisionPropertiesCfg, + }, + field_name="collision_props", + ) + if "material_props" in init_dict: + cfg.material_props = _physics_property_cfg_from_dict( + init_dict["material_props"], + common_type=RigidBodyMaterialCfg, + backend_types={"newton": NewtonRigidBodyMaterialCfg}, + field_name="material_props", + ) + return cfg + + def to_dict(self) -> dict[str, Any]: + """Serialize grouped properties without losing backend subclasses.""" + return { + "mass_props": ( + None if self.mass_props is None else self.mass_props.to_dict() + ), + "rigid_props": ( + None + if self.rigid_props is None + else {**self.rigid_props.to_dict(), "backend": "default"} + ), + "collision_props": _physics_property_cfg_to_dict( + self.collision_props, + common_type=CollisionPropertiesCfg, + backend_types={ + "default": DefaultCollisionPropertiesCfg, + "newton": NewtonCollisionPropertiesCfg, + }, + field_name="collision_props", + ), + "material_props": _physics_property_cfg_to_dict( + self.material_props, + common_type=RigidBodyMaterialCfg, + backend_types={"newton": NewtonRigidBodyMaterialCfg}, + field_name="material_props", + ), + } + + @property + def enable_collision(self) -> bool: + """Compatibility view used by legacy object initialization.""" + value = ( + None + if self.collision_props is None + else self.collision_props.collision_enabled + ) + return True if value is None else bool(value) + + def to_dexsim_physical_attr( + self, + *, + base: PhysicalAttr | None = None, + ) -> PhysicalAttr: + """Translate configured Default-compatible values to ``PhysicalAttr``. + + Args: + base: Optional native attributes to overlay. This is used by the + retained raw Default articulation path for sparse per-link + updates. + + Returns: + A DexSim physical-attribute object using its defaults for every + unconfigured grouped field. + """ + attr = PhysicalAttr() if base is None else _copy_dexsim_physical_attr(base) + configs = ( + (self.mass_props, {"recompute_inertia": None}), + (self.rigid_props, {}), + (self.collision_props, {"collision_enabled": "enable_collision"}), + (self.material_props, {}), + ) + for cfg, field_map in configs: + if cfg is None: + continue + for item in fields(cfg): + value = getattr(cfg, item.name) + target_name = field_map.get(item.name, item.name) + if ( + value is None + or target_name is None + or not hasattr(attr, target_name) + ): + continue + if target_name in {"inertia", "com_position"}: + value = np.asarray(value, dtype=np.float32) + elif target_name == "com_quaternion": + value = convert_quat(np.asarray(value, dtype=np.float32), to="wxyz") + setattr(attr, target_name, value) + return attr + + @classmethod + def from_dexsim_physical_attr( + cls, + attr: PhysicalAttr, + ) -> RigidBodyPhysicsCfg: + """Capture native Default attributes in the grouped configuration.""" + + def _array(name: str) -> np.ndarray | None: + value = getattr(attr, name, None) + return None if value is None else np.asarray(value, dtype=np.float32) + + com_quaternion = _array("com_quaternion") + if com_quaternion is not None: + com_quaternion = convert_quat(com_quaternion, to="xyzw") + return cls( + mass_props=MassPropertiesCfg( + mass=getattr(attr, "mass", None), + density=getattr(attr, "density", None), + inertia=_array("inertia"), + com_position=_array("com_position"), + com_quaternion=com_quaternion, + ), + rigid_props=DefaultRigidBodyPropertiesCfg( + angular_damping=getattr(attr, "angular_damping", None), + linear_damping=getattr(attr, "linear_damping", None), + max_depenetration_velocity=getattr( + attr, "max_depenetration_velocity", None + ), + sleep_threshold=getattr(attr, "sleep_threshold", None), + min_position_iters=getattr(attr, "min_position_iters", None), + min_velocity_iters=getattr(attr, "min_velocity_iters", None), + max_linear_velocity=getattr(attr, "max_linear_velocity", None), + max_angular_velocity=getattr(attr, "max_angular_velocity", None), + enable_ccd=getattr(attr, "enable_ccd", None), + ), + collision_props=DefaultCollisionPropertiesCfg( + collision_enabled=getattr(attr, "enable_collision", None), + contact_offset=getattr(attr, "contact_offset", None), + rest_offset=getattr(attr, "rest_offset", None), + torsional_patch_radius=getattr(attr, "torsional_patch_radius", None), + min_torsional_patch_radius=getattr( + attr, "min_torsional_patch_radius", None + ), + disable_strong_friction=getattr(attr, "disable_strong_friction", None), + ), + material_props=RigidBodyMaterialCfg( + restitution=getattr(attr, "restitution", None), + dynamic_friction=getattr(attr, "dynamic_friction", None), + static_friction=getattr(attr, "static_friction", None), + ), + ) + + +_REMOVED_FLAT_RIGID_BODY_FIELDS = frozenset( + { + "mass", + "density", + "inertia", + "com_position", + "com_quaternion", + "angular_damping", + "linear_damping", + "max_depenetration_velocity", + "sleep_threshold", + "min_position_iters", + "min_velocity_iters", + "max_linear_velocity", + "max_angular_velocity", + "enable_ccd", + "contact_offset", + "rest_offset", + "enable_collision", + "restitution", + "dynamic_friction", + "static_friction", + } +) + + +def _rigid_body_physics_from_dict(value: Mapping[str, Any]) -> RigidBodyPhysicsCfg: + """Parse the grouped rigid-body physics schema. + + Flat ``attrs`` fields and their compatibility configuration types were + removed. Reject them at the config boundary so no input silently changes + physical meaning. + """ + flat_fields = _REMOVED_FLAT_RIGID_BODY_FIELDS.intersection(value) + if flat_fields: + raise ValueError( + "Removed flat rigid-body attrs fields: " + f"{sorted(flat_fields)}. Use grouped mass_props, rigid_props, " + "collision_props, and material_props." + ) + return RigidBodyPhysicsCfg.from_dict(value) diff --git a/embodichain/lab/sim/cfg/rigid_object.py b/embodichain/lab/sim/cfg/rigid_object.py new file mode 100644 index 000000000..ab0bdfa6e --- /dev/null +++ b/embodichain/lab/sim/cfg/rigid_object.py @@ -0,0 +1,219 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Rigid object and rigid-object-group configuration.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import MISSING +import os +import warnings +from typing import Any, Dict, Literal + +from dexsim.types import ActorType + +from embodichain.utils import configclass, is_configclass, logger + +from ..shapes import ShapeCfg +from .asset import AssetPhysicsMode, ObjectBaseCfg, _resolve_asset_physics_mode +from .rigid import RigidBodyPhysicsCfg + + +@configclass +class RigidObjectCfg(ObjectBaseCfg): + """Configuration for a rigid body asset in the simulation. + + This class extends the base asset configuration to include specific properties for rigid bodies, + such as physical attributes and collision group. + """ + + shape: ShapeCfg = ShapeCfg() + """Shape configuration for the rigid body. """ + + # TODO: supoort basic primitive shapes, such as box, sphere, etc cfg and spawn method. + + attrs: RigidBodyPhysicsCfg = RigidBodyPhysicsCfg() + """Rigid-body physics. + + :class:`RigidBodyPhysicsCfg` groups portable and backend-native intent. + """ + + body_type: Literal["dynamic", "kinematic", "static"] = "dynamic" + + body_scale: tuple | list = (1.0, 1.0, 1.0) + """Scale of the rigid body in the simulation world frame.""" + + asset_physics_mode: AssetPhysicsMode = "preserve" + """How a file-backed asset's physical properties are handled. + + ``"preserve"`` keeps the USD-authored physics. ``"overlay"`` applies + configured properties on top of the parsed asset. Procedural shapes always + use config. + """ + + def resolve_asset_physics_mode(self) -> AssetPhysicsMode: + """Return the effective file-backed physics policy.""" + return _resolve_asset_physics_mode(self.asset_physics_mode) + + @classmethod + def from_dict(cls, init_dict: Dict[str, Any]) -> RigidObjectCfg: + """Parse a rigid object and normalize legacy mesh collision ownership.""" + data = dict(init_dict) + attrs_value = data.get("attrs") + if isinstance(attrs_value, Mapping) and "mesh_collision_props" in attrs_value: + shape_value = data.get("shape") + if not isinstance(shape_value, Mapping): + raise ValueError( + "Legacy attrs.mesh_collision_props requires a mapping-valued " + "MeshCfg shape so it can migrate to shape.collision." + ) + shape_data = dict(shape_value) + if shape_data.get("shape_type") != "Mesh": + raise ValueError( + "Legacy attrs.mesh_collision_props can migrate only to a " + "MeshCfg shape." + ) + if shape_data.get("collision") is not None: + raise ValueError( + "attrs.mesh_collision_props cannot be combined with " + "shape.collision." + ) + attrs_data = dict(attrs_value) + shape_data["collision"] = attrs_data.pop("mesh_collision_props") + data["shape"] = shape_data + data["attrs"] = attrs_data + warnings.warn( + "RigidBodyPhysicsCfg.mesh_collision_props is deprecated; use " + "MeshCfg.collision.", + DeprecationWarning, + stacklevel=2, + ) + return super().from_dict(data) + + def to_dexsim_body_type(self) -> ActorType: + """Convert the body type to dexsim ActorType.""" + if self.body_type == "dynamic": + return ActorType.DYNAMIC + elif self.body_type == "kinematic": + return ActorType.KINEMATIC + elif self.body_type == "static": + return ActorType.STATIC + else: + logger.log_error( + f"Invalid body type '{self.body_type}' specified. Must be one of 'dynamic', 'kinematic', or 'static'." + ) + + +@configclass +class RigidObjectGroupCfg: + """Configuration for a rigid object group asset in the simulation. + + Rigid object groups can be initialized from multiple rigid object configurations specified in a folder. + If `folder_path` is specified, user should provide a RigidObjectCfg in `rigid_objects` as a template configuration for + all objects in the group. + + For example: + ```python + rigid_object_group: RigidObjectGroupCfg( + folder_path="path/to/folder", + max_num=5, + rigid_objects={ + "template_obj": RigidObjectCfg( + shape=MeshCfg( + fpath="", # fpath will be ignored when folder_path is specified + ), + body_type="dynamic", + ) + } + ) + """ + + uid: str | None = None + + rigid_objects: Dict[str, RigidObjectCfg] = MISSING + """Configuration for the rigid objects in the group.""" + + body_type: Literal["dynamic", "kinematic"] = "dynamic" + """Body type for all rigid objects in the group. """ + + folder_path: str | None = None + """Path to the folder containing the rigid object assets. + + This is used to initialize multiple rigid object configurations from a folder. + """ + + max_num: int = 1 + """Maximum number of rigid objects to initialize from the folder. + + This is only used when `folder_path` is specified. + """ + + ext: str = ".obj" + """File extension for the rigid object assets. + + This is only used when `folder_path` is specified. + """ + + @classmethod + def from_dict(cls, init_dict: Dict[str, Any]) -> RigidObjectGroupCfg: + """Initialize the configuration from a dictionary.""" + cfg = cls() + for key, value in init_dict.items(): + if hasattr(cfg, key): + attr = getattr(cfg, key) + if is_configclass(attr): + setattr( + cfg, key, attr.from_dict(value) + ) # Call from_dict on the attribute + elif key == "rigid_objects" and "folder_path" not in init_dict: + rigid_objects_cfg = {} + for obj_name, obj_cfg in value.items(): + rigid_objects_cfg[obj_name] = RigidObjectCfg.from_dict(obj_cfg) + setattr(cfg, key, rigid_objects_cfg) + elif key == "rigid_objects" and "folder_path" in init_dict: + folder_path = init_dict["folder_path"] + max_num = init_dict.get("max_num", 1) + rigid_objects_cfg = {} + if os.path.exists(folder_path) and os.path.isdir(folder_path): + files = os.listdir(folder_path) + files = [f for f in files if f.endswith(cfg.ext)] + # select files up to max_num + n_file = len(files) + select_files = [] + for i in range(max_num): + select_files.append(files[i % n_file]) + + for i, file_name in enumerate(select_files): + file_path = os.path.join(folder_path, file_name) + rigid_obj_cfg: RigidObjectCfg = RigidObjectCfg.from_dict( + list(init_dict["rigid_objects"].values())[0] + ) + rigid_obj_cfg.uid = f"{cfg.uid}_obj_{i}" + rigid_obj_cfg.shape.fpath = file_path + rigid_objects_cfg[rigid_obj_cfg.uid] = rigid_obj_cfg + setattr(cfg, "rigid_objects", rigid_objects_cfg) + else: + logger.log_error( + f"Folder '{folder_path}' does not exist or is not a directory." + ) + else: + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + return cfg diff --git a/embodichain/lab/sim/cfg/robot.py b/embodichain/lab/sim/cfg/robot.py new file mode 100644 index 000000000..59de97e25 --- /dev/null +++ b/embodichain/lab/sim/cfg/robot.py @@ -0,0 +1,378 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Robot configuration, serialization, and backend preset selection.""" + +from __future__ import annotations + +from collections.abc import Mapping +from copy import deepcopy +from dataclasses import MISSING, fields +import enum +import json +from typing import Dict, List + +import numpy as np +import torch + +from embodichain.utils import configclass, is_configclass, logger +from embodichain.utils.utility import key_in_nested_dict + +from ..workspace.cfg import RobotWorkspaceCfg +from .articulation import ( + ArticulationCfg, + JointDrivePropertiesCfg, + _raise_removed_articulation_cfg_fields, + link_attrs_from_dict, +) +from .asset import AssetPhysicsMode +from .rigid import _rigid_body_physics_from_dict +from .simulation import ( + PhysicsBackendCfg, + _normalize_newton_solver_type, + physics_backend_from_cfg, +) +from .urdf import URDFCfg + + +def _get_data_path(path: str) -> str: + """Resolve data through the public facade for monkeypatch compatibility.""" + from . import get_data_path + + return get_data_path(path) + + +@configclass +class RobotCfg(ArticulationCfg): + from embodichain.lab.sim.solvers import SolverCfg + + """Configuration for a robot asset in the simulation. + """ + + joint_drive_props: JointDrivePropertiesCfg = JointDrivePropertiesCfg( + drive_type="force", + stiffness=1e4, + damping=1e3, + max_effort=1e10, + max_velocity=1e10, + friction=0.0, + armature=0.0, + ) + """Joint drive, limit, friction, and armature properties.""" + + asset_physics_mode: AssetPhysicsMode = "overlay" + """Apply configured robot physics on top of source-authored values.""" + + control_parts: Dict[str, List[str]] | None = None + """Control parts is the mapping from part name to joint names. + + For example, {'left_arm': ['joint1', 'joint2'], 'right_arm': ['joint3', 'joint4']} + If no control part is specified, the robot will use all joints as a single control part. + + Note: + - if `control_parts` is specified, `solver_cfg` must be a dict with part names as + keys corresponding to the control parts name. + - The joint names in the control parts support regular expressions, e.g., 'joint[1-6]'. + After initialization of robot, the names will be expanded to a list of full joint names. + - `Robot` is a derived class of `Articulation`, with control parts support. So the `joint_drive_props` + in `ArticulationCfg` can use control part as key to specify the corresponding joint drive properties, + which will be overridden if these joint names are already specified. + """ + + urdf_cfg: URDFCfg | None = None + """URDF assembly configuration which allows for assembling a robot from multiple URDF components. + """ + + # TODO: how to support one solver for multiple parts? + solver_cfg: SolverCfg | Dict[str, SolverCfg] | None = None + """Solver is used to compute forward and inverse kinematics for the robot. + """ + + workspace_cfg: Dict[str, RobotWorkspaceCfg] | None = None + """Runtime workspace cache configuration keyed by control-part name.""" + + @classmethod + def from_dict(cls, init_dict: Dict[str, str | float | tuple]) -> RobotCfg: + """Initialize the configuration from a dictionary.""" + if isinstance(init_dict, cls): + return init_dict + + _raise_removed_articulation_cfg_fields(init_dict) + + import importlib + + solver_module = importlib.import_module("embodichain.lab.sim.solvers") + + cfg = cls() # Create a new instance of the class (cls) + for key, value in init_dict.items(): + if key == "link_attrs" and isinstance(value, dict): + cfg.link_attrs = link_attrs_from_dict(value) + elif key == "attrs" and isinstance(value, Mapping): + cfg.attrs = _rigid_body_physics_from_dict(value) + elif hasattr(cfg, key): + attr = getattr(cfg, key) + if key == "urdf_cfg": + from embodichain.lab.sim.cfg import URDFCfg + + setattr(cfg, key, URDFCfg.from_dict(value)) + elif key == "workspace_cfg" and isinstance(value, dict): + setattr( + cfg, + key, + { + part: ( + part_cfg + if isinstance(part_cfg, RobotWorkspaceCfg) + else RobotWorkspaceCfg(**part_cfg) + ) + for part, part_cfg in value.items() + }, + ) + elif key == "fpath": + setattr(cfg, key, _get_data_path(value)) + elif isinstance(attr, JointDrivePropertiesCfg) and isinstance( + value, dict + ): + setattr( + cfg, + key, + JointDrivePropertiesCfg.from_dict(value, defaults=attr), + ) + elif is_configclass(attr): + setattr( + cfg, key, attr.from_dict(value) + ) # Call from_dict on the attribute + elif isinstance(value, dict) and "class_type" in value: + setattr( + cfg, + key, + getattr(solver_module, f"{value['class_type']}Cfg").from_dict( + value + ), + ) + elif isinstance(value, dict) and key_in_nested_dict( + value, "class_type" + ): + setattr( + cfg, + key, + { + k: getattr( + solver_module, f"{v['class_type']}Cfg" + ).from_dict(v) + for k, v in value.items() + }, + ) + + else: + setattr(cfg, key, value) + else: + logger.log_warning( + f"Key '{key}' not found in {cfg.__class__.__name__}." + ) + return cfg + + def _build_defaults(self, init_dict: dict | None = None) -> None: + """Populate default config fields from ``init_dict``. + + Subclasses override this to read variant/version fields from + ``init_dict``, set them on ``self``, and populate ``urdf_cfg``, + ``control_parts``, ``solver_cfg``, ``joint_drive_props`` and ``attrs``. + The base implementation is a no-op. + + .. attention:: + Do NOT call :func:`merge_robot_cfg` from here -- the subclass + ``from_dict`` calls this hook first, then ``merge_robot_cfg``. + Calling ``merge_robot_cfg`` here would recurse, because + ``merge_robot_cfg`` itself calls ``RobotCfg.from_dict``. + + Args: + init_dict: The raw override dict passed to ``from_dict``. + """ + return None + + def to_dict(self): + """Serialize config to a plain dict (enums, numpy, nested configclass).""" + + def serialize(obj, _visited=None): + if _visited is None: + _visited = set() + if isinstance(obj, enum.Enum): + return obj.value + tracked_id = None + if not isinstance(obj, (str, int, float, bool, type(None))): + tracked_id = id(obj) + if tracked_id in _visited: + return None + _visited.add(tracked_id) + + try: + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, dict): + return { + (k.value if isinstance(k, enum.Enum) else str(k)): serialize( + v, _visited + ) + for k, v in obj.items() + } + if isinstance(obj, (list, tuple)): + return [serialize(v, _visited) for v in obj] + if hasattr(obj, "to_dict") and obj is not self: + return serialize(obj.to_dict(), _visited) + if hasattr(obj, "__dict__"): + return { + k: serialize(v, _visited) + for k, v in obj.__dict__.items() + if v is not None + } + return obj + finally: + if tracked_id is not None: + _visited.remove(tracked_id) + + return serialize(self) + + def to_string(self): + """Return config as a JSON string.""" + return json.dumps(self.to_dict(), indent=2) + + def save_to_file(self, filepath): + """Save config to a local file as JSON.""" + with open(filepath, "w") as f: + f.write(self.to_string()) + + def build_pk_serial_chain( + self, device: torch.device = torch.device("cpu"), **kwargs + ) -> Dict[str, "pk.SerialChain"]: + """Build the serial chain from the URDF file. + + Note: + This method is usually used in imitation dataset saving (compute eef pose from qpos using FK) + and model training (provide a differentiable FK layer or loss computation). + + Args: + device (torch.device): The device to which the chain will be moved. Defaults to CPU. + **kwargs: Additional arguments for building the serial chain. + + Returns: + Dict[str, pk.SerialChain]: The serial chain of the robot for specified control part. + """ + return {} + + +@configclass +class RobotPresetCfg: + """Base class for replace-only robot configurations across physics backends. + + Subclasses declare complete :class:`RobotCfg` alternatives as fields. A + ``default`` field is required; optional fields use Newton backend or solver + profile names such as ``newton``, ``newton_mujoco_warp``, or + ``newton_mjwarp``. The active :class:`PhysicsBackendCfg` selects one + complete alternative at + :meth:`SimulationManager.add_robot`; alternatives are never field-merged. + + Portable robot properties should remain on one ordinary :class:`RobotCfg`. + Use this wrapper only when an asset, actuator model, or native physics value + genuinely requires a different complete robot definition. + + Example:: + + @configclass + class MyRobotPresetCfg(RobotPresetCfg): + default: RobotCfg = MyRobotCfg() + newton_mujoco_warp: RobotCfg = MyNewtonRobotCfg() + """ + + def resolve( + self, + physics_cfg: PhysicsBackendCfg, + *, + newton_solver_type: str | None = None, + ) -> RobotCfg: + """Return an isolated complete robot config for the active backend. + + Args: + physics_cfg: The scene's backend-selecting physics configuration. + newton_solver_type: Resolved Newton solver name when it is already + available from the runtime. If omitted, it is inferred from + ``physics_cfg``. + + Returns: + A deep copy of the highest-priority complete robot alternative. + + Raises: + TypeError: If a preset name is unsupported, ``default`` is + undeclared, or a selected alternative is not a + :class:`RobotCfg`. + ValueError: If no declared alternative can satisfy the backend. + """ + options = {item.name: getattr(self, item.name) for item in fields(self)} + invalid_names = { + name + for name in options + if name != "default" and name != "newton" and not name.startswith("newton_") + } + if invalid_names: + raise TypeError( + f"{type(self).__name__} uses unsupported preset name(s) " + f"{sorted(invalid_names)}; use 'default' or 'newton[_]'." + ) + if "default" not in options: + raise TypeError( + f"{type(self).__name__} must declare a 'default' RobotCfg preset." + ) + + backend = physics_backend_from_cfg(physics_cfg) + if backend == "default": + candidates = ("default",) + else: + solver_type = newton_solver_type + if solver_type is None: + solver_cfg = physics_cfg.solver_cfg + if solver_cfg is None: + solver_type = "auto" + elif isinstance(solver_cfg, Mapping): + solver_type = str( + solver_cfg.get("solver_type") + or solver_cfg.get("class_type") + or "auto" + ) + else: + solver_type = str(getattr(solver_cfg, "solver_type")) + solver_type = _normalize_newton_solver_type(solver_type) + solver_candidates = [] + if solver_type != "auto": + solver_candidates.append(f"newton_{solver_type}") + if solver_type == "mujoco_warp": + solver_candidates.append("newton_mjwarp") + candidates = (*solver_candidates, "newton", "default") + + for candidate in candidates: + selected = options.get(candidate) + if selected is None or selected is MISSING: + continue + if not isinstance(selected, RobotCfg): + raise TypeError( + f"{type(self).__name__}.{candidate} must be a RobotCfg, " + f"got {type(selected).__name__}." + ) + return deepcopy(selected) + + raise ValueError( + f"{type(self).__name__} has no usable preset for {candidates!r}; " + f"declared options are {sorted(options)}." + ) diff --git a/embodichain/lab/sim/cfg/scene.py b/embodichain/lab/sim/cfg/scene.py new file mode 100644 index 000000000..256978f74 --- /dev/null +++ b/embodichain/lab/sim/cfg/scene.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. +# ---------------------------------------------------------------------------- + +"""Light and inter-object constraint configuration.""" + +from __future__ import annotations + +from dataclasses import MISSING +from typing import Literal + +import numpy as np + +from embodichain.utils import configclass + +from .asset import ObjectBaseCfg + + +@configclass +class LightCfg(ObjectBaseCfg): + """Configuration for a light asset in the simulation. + + Supports six light types matching the dexsim rendering backend: + + - ``"point"``: Per-environment omnidirectional point light with position + and falloff radius. Created as a batched light (one per environment). + - ``"sun"``: Global directional sun light (infinite distance). Created as + a single scene-level instance. Uses direction only; position is ignored. + Sun-specific fields (``angular_radius``, ``halo_size``, ``halo_falloff``) + are reserved for future backend support. + - ``"direction"``: Global pure directional light at infinite distance. + Created as a single scene-level instance. Direction only; no position. + - ``"spot"``: Per-environment spotlight with position, direction, and + inner/outer cone angles. Created as a batched light. + - ``"rect"``: Per-environment rectangular area light with position, + direction, width, and height. Created as a batched light. + - ``"mesh"``: Per-environment mesh-based emissive light. Requires a + :class:`~dexsim.models.MeshObject` via + :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` + (not tensor-batched). Created as a batched light. + + .. attention:: + The ``angular_radius``, ``halo_size``, and ``halo_falloff`` fields are + reserved for future use. The dexsim Python bindings do not yet expose + setters for these sun-specific properties. + """ + + light_type: Literal["point", "sun", "direction", "spot", "rect", "mesh"] = "point" + """Light type. Supported: ``"point"``, ``"sun"``, ``"direction"``, ``"spot"``, ``"rect"``, ``"mesh"``.""" + + # ------------------------------------------------------------------ + # Universal properties (apply to all light types) + # ------------------------------------------------------------------ + + color: tuple[float, float, float] = (1.0, 1.0, 1.0) + """RGB color of the light source. Defaults to white ``(1.0, 1.0, 1.0)``.""" + + intensity: float = 30.0 + """Intensity of the light source in watts/m^2. Defaults to ``30.0``.""" + + enable_shadow: bool = True + """Whether the light casts shadows. Defaults to ``True``.""" + + # ------------------------------------------------------------------ + # Point light + # ------------------------------------------------------------------ + + radius: float = 10.0 + """Falloff radius for point lights. Only used when ``light_type="point"``. Defaults to ``10.0``.""" + + # ------------------------------------------------------------------ + # Directional properties (sun, direction, spot, rect, mesh) + # ------------------------------------------------------------------ + + direction: tuple[float, float, float] = (0.0, 0.0, -1.0) + """Direction vector for directional, spot, rect, and mesh lights. + Defaults to ``(0.0, 0.0, -1.0)`` (pointing down along -Z).""" + + # ------------------------------------------------------------------ + # Sun light (reserved — Python bindings not yet available) + # ------------------------------------------------------------------ + + angular_radius: float = 0.5 + """Angular radius of the sun disc in degrees. Reserved for future use.""" + + halo_size: float = 10.0 + """Halo size for sun light. Reserved for future use.""" + + halo_falloff: float = 3.0 + """Halo falloff for sun light. Reserved for future use.""" + + # ------------------------------------------------------------------ + # Spot light + # ------------------------------------------------------------------ + + spot_angle_inner: float = 30.0 + """Inner cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. + Defaults to ``30.0``.""" + + spot_angle_outer: float = 45.0 + """Outer cone angle of the spotlight in degrees. Only used when ``light_type="spot"``. + Defaults to ``45.0``.""" + + # ------------------------------------------------------------------ + # Rect light + # ------------------------------------------------------------------ + + rect_width: float = 1.0 + """Width of the rectangular area light. Only used when ``light_type="rect"``. + Defaults to ``1.0``.""" + + rect_height: float = 1.0 + """Height of the rectangular area light. Only used when ``light_type="rect"``. + Defaults to ``1.0``.""" + + # ------------------------------------------------------------------ + # Mesh light + # ------------------------------------------------------------------ + + mesh_path: str = "" + """Asset path for mesh-based emissive lights. Only used when ``light_type="mesh"``. + The actual mesh assignment is done via + :meth:`embodichain.lab.sim.objects.light.Light.set_mesh` which accepts a + :class:`dexsim.models.MeshObject`. This field stores the path for reference.""" + + +@configclass +class RigidConstraintCfg: + """Configuration for a fixed constraint between two RigidObjects. + + The constraint binds rigid_object_a's entity[i] to rigid_object_b's entity[i] + within arena[i] (one constraint per arena). + + Args: + name: Base constraint name. Per-arena names are derived as ``f"{name}"`` + (single env) or ``f"{name}_{i}"`` (multi env). + rigid_object_a_uid: UID of the first RigidObject (must exist in the sim). + rigid_object_b_uid: UID of the second RigidObject (must exist in the sim). + local_frame_a: 4x4 joint frame in object A's local coordinates. + ``None`` -> identity (object A's origin). Accepts a single + ``(4, 4)`` matrix (shared by all envs) or an ``(N, 4, 4)`` array + (one frame per env). Defaults to None. + local_frame_b: 4x4 joint frame in object B's local coordinates. + ``None`` -> the frame is computed per env as ``inv(pose_B) @ pose_A`` + from the objects' current poses, so the constraint welds the objects + at their *current* relative pose (rather than pulling their origins + together). An explicit ``(4, 4)`` or ``(N, 4, 4)`` value is used + verbatim. Defaults to None. + constraint_type: Reserved for future typed constraints (prismatic, + revolute, spherical, d6). Only ``"fixed"`` is supported in v1. + + .. attention:: + Both objects must be :class:`RigidObject` instances and must share the + same number of arenas. + """ + + name: str = MISSING + """Base name of the constraint (per-arena names are derived from this).""" + + rigid_object_a_uid: str = MISSING + """UID of the first RigidObject.""" + + rigid_object_b_uid: str = MISSING + """UID of the second RigidObject.""" + + local_frame_a: np.ndarray | None = None + """Local joint frame on object A. None -> identity (object A's origin).""" + + local_frame_b: np.ndarray | None = None + """Local joint frame on object B. None -> ``inv(pose_B) @ pose_A`` per env + (weld at the objects' current relative pose).""" + + constraint_type: Literal["fixed"] = "fixed" + """Constraint type. Only ``"fixed"`` is supported in v1.""" diff --git a/embodichain/lab/sim/cfg/simulation.py b/embodichain/lab/sim/cfg/simulation.py new file mode 100644 index 000000000..476077936 --- /dev/null +++ b/embodichain/lab/sim/cfg/simulation.py @@ -0,0 +1,563 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""World-level rendering and physics-backend configuration.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import field, fields +from typing import Any, Literal, Sequence, TYPE_CHECKING + +import dexsim +import numpy as np +import torch +from dexsim.types import Renderer, ToneMappingType + +from embodichain.utils import configclass, logger + +if TYPE_CHECKING: + from dexsim.engine.newton_physics import NewtonCfg + from dexsim.engine.newton_physics.solvers_cfg import NewtonSolverCfg + + +@configclass +class RenderCfg: + renderer: Literal["auto", "hybrid", "fast-rt", "rt"] = "auto" + """Renderer backend to use for the simulation. Options are 'auto', 'hybrid', 'fast-rt', and 'rt'. + + Note: + - 'auto' selects a default renderer based on the detected GPU: RTX-series cards use + 'hybrid', while datacenter cards (A100/A800, H100/H800/H200/H20) use 'fast-rt'. + If no CUDA device is available or the GPU is unknown, it falls back to 'hybrid'. + - 'hybrid' uses ray tracing for shadows and reflections while keeping rasterization for primary rendering, + providing a balance between performance and visual quality. + - 'fast-rt' is a fully ray-traced renderer for maximum visual fidelity, but may have higher computational cost. + - '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' or 'fast-rt' and enable_denoiser is False.""" + + tone_mapping_enabled: bool = False + """Whether to map HDR RGB output with the modified Reinhard curve.""" + + tone_mapping_exposure: float = 1.0 + """Fixed linear exposure multiplier applied before tone mapping.""" + + def __post_init__(self) -> None: + """Validate rendering parameters.""" + if self.spp < 1: + logger.log_error("RenderCfg.spp must be at least 1.", ValueError) + if self.tone_mapping_exposure < 0.0: + logger.log_error( + "RenderCfg.tone_mapping_exposure must be non-negative.", ValueError + ) + + def to_dexsim_flags(self) -> Renderer: + """Convert the renderer name to DexSim's renderer enum.""" + if self.renderer == "hybrid": + return Renderer.HYBRID + elif self.renderer == "fast-rt": + return Renderer.FASTRT + elif self.renderer == "rt": + return Renderer.OFFLINERT + elif self.renderer == "auto": + # 'auto' is normally resolved by the SimulationManager before this is + # called. If it reaches here (e.g. used standalone), fall back safely. + logger.log_warning( + "Renderer 'auto' was not resolved before converting to dexsim flags. " + "Falling back to 'hybrid'." + ) + return Renderer.HYBRID + else: + logger.log_error( + f"Invalid renderer type '{self.renderer}' specified. Must be one of 'auto', 'hybrid', 'fast-rt', or 'rt'." + ) + + def apply_to_dexsim_config(self, world_config: dexsim.WorldConfig) -> None: + """Apply rendering settings to a DexSim world configuration. + + Args: + world_config: DexSim world configuration to update in place. + """ + world_config.renderer = self.to_dexsim_flags() + world_config.raytrace_config.render_iterations_per_frame = self.spp + world_config.raytrace_config.open_denoise = True + world_config.postprocess_config.tone_mapping_enabled = self.tone_mapping_enabled + world_config.postprocess_config.tone_mapping_type = ( + ToneMappingType.MODIFIED_REINHARD + ) + world_config.postprocess_config.tone_mapping_exposure = ( + self.tone_mapping_exposure + ) + + +@configclass +class GPUMemoryCfg: + """GPU buffer capacities for the Default backend's GPU dynamics pipeline. + + Default-backend GPU buffers cannot all grow dynamically. Values that are + too small may therefore produce overflow warnings, dropped contacts, or an + invalid simulation. These settings are applied only when the Default + backend runs on CUDA; they have no effect on Default CPU or Newton. + """ + + temp_buffer_capacity: int = 2**24 + """Temporary pinned-host buffer capacity in bytes. + + Increase this when the Default backend reports a pinned-host linear + allocator overflow. + """ + + max_rigid_contact_count: int = 2**19 + """Maximum number of rigid-contact records in the GPU contact stream. + + Increase this when the Default backend reports + ``Contact buffer overflow detected``. + """ + + max_rigid_patch_count: int = ( + 2**18 + ) # 81920 is DexSim default but most tasks work with 2**18 + """Maximum number of rigid-contact patches in the GPU patch stream. + + A patch groups nearby contact points that share a contact normal. Increase + this when the Default backend reports ``Patch buffer overflow detected``. + """ + + heap_capacity: int = 2**26 + """Initial capacity in bytes of the GPU and pinned-host memory heaps.""" + + found_lost_pairs_capacity: int = ( + 2**25 + ) # 262144 is DexSim default but most tasks work with 2**25 + """Capacity of broad-phase found/lost pair records.""" + + found_lost_aggregate_pairs_capacity: int = 2**10 + """Capacity of found/lost pair records generated by aggregates.""" + + total_aggregate_pairs_capacity: int = 2**10 + """Capacity of all aggregate-pair records in the GPU pipeline.""" + + +def _gravity_vector( + gravity: Sequence[float] | np.ndarray, +) -> list[float]: + """Validate and normalize a backend-neutral gravity vector.""" + values = np.asarray(gravity, dtype=np.float64).reshape(-1) + if values.size != 3 or not np.all(np.isfinite(values)): + raise ValueError("Gravity must contain three finite values.") + return values.tolist() + + +@configclass +class PhysicsBackendCfg: + """Backend-neutral simulation timing, device, and gravity configuration. + + Concrete backend configs inherit this class. The config type selects the + backend; no independent backend string can disagree with it. + """ + + physics_dt: float = 1.0 / 100.0 + """Duration of one physics step in seconds. + + Environment control steps may contain multiple physics steps. For Newton, + this interval is further divided by :attr:`NewtonPhysicsCfg.num_substeps`. + """ + + device: str | torch.device = "cpu" + """Compute device used to build and step the selected physics backend.""" + + gravity: Sequence[float] | np.ndarray = field( + default_factory=lambda: np.array([0.0, 0.0, -9.81]) + ) + """World-frame gravity vector in meters per second squared.""" + + +@configclass +class DefaultPhysicsCfg(PhysicsBackendCfg): + """Configuration selector for the Default physics backend.""" + + bounce_threshold: float = 2.0 + """Relative normal-speed threshold below which contacts do not bounce [m/s].""" + + enable_ccd: bool = False + """Whether to enable scene-level continuous collision detection (CCD). + + A rigid body must also set :attr:`DefaultRigidBodyPropertiesCfg.enable_ccd` + for CCD to be used on that body. + """ + + length_tolerance: float = 0.05 + """Representative scene length used by the Default backend's tolerance scale [m]. + + Set this near the characteristic size of simulated objects. It is a scene + scale, not an accuracy knob, and must be configured before world creation. + """ + + speed_tolerance: float = 0.25 + """Representative scene speed used by the Default backend's tolerance scale [m/s]. + + The backend derives several internal thresholds from this value and + :attr:`length_tolerance`. + """ + + gpu_memory: GPUMemoryCfg = field(default_factory=GPUMemoryCfg) + """Fixed-capacity GPU buffers used by Default-backend CUDA simulation.""" + + def to_dexsim_args(self) -> dict[str, Any]: + """Convert to DexSim physics arguments. + + Solver implementation details that are not exposed by + :class:`DefaultPhysicsCfg` retain their established defaults here. + """ + args = { + "gravity": _gravity_vector(self.gravity), + "bounce_threshold": self.bounce_threshold, + "enable_ccd": self.enable_ccd, + "enable_enhanced_determinism": False, + "enable_friction_every_iteration": True, + } + return args + + +@configclass +class NewtonCollisionPipelineCfg: + """Newton collision-pipeline settings owned at scene scope. + + These values map to DexSim's ``NewtonCollisionPipelineCfg``. Per-shape + contact and SDF values belong to :class:`NewtonCollisionPropertiesCfg` + instead. The pipeline performs broad-phase pair selection, narrow-phase + contact generation, and optional contact reduction for the complete scene. + + See the `Newton collision guide + `_ + for the native pipeline semantics. + """ + + reduce_contacts: bool = True + """Whether to reduce dense mesh contacts to a representative subset. + + Reduction lowers contact count and usually improves performance and solver + stability for mesh-heavy scenes. + """ + + rigid_contact_max: int | None = None + """Maximum number of allocated rigid contacts. + + ``None`` uses the model-provided capacity when available and otherwise lets + Newton estimate it from the scene's shapes and candidate pairs. + """ + + max_triangle_pairs: int = 4_000_000 + """Maximum triangle-pair candidates allocated by the narrow phase. + + Increase this only when complex meshes or heightfields report triangle-pair + overflow. EmbodiChain intentionally uses a larger default than upstream + Newton for mesh-heavy robotics scenes. + """ + + soft_contact_max: int | None = None + """Maximum number of allocated particle/soft contacts. + + ``None`` lets Newton derive the capacity from shape and particle counts. + """ + + soft_contact_margin: float = 0.01 + """Distance margin used to generate particle/soft contacts [m].""" + + broad_phase: Literal["nxn", "sap", "explicit"] | Any | None = None + """Built-in broad-phase mode or a prebuilt Newton broad-phase object. + + ``"explicit"`` tests precomputed pairs, ``"nxn"`` performs an all-pairs + test, and ``"sap"`` uses sweep-and-prune. ``None`` keeps Newton's default. + A prebuilt object is an expert path and must be compatible with + :attr:`narrow_phase`. + """ + + shape_pairs_filtered: Any | None = None + """Optional precomputed pairs for ``"explicit"`` broad phase. + + When provided, this must be a Warp array of shape-index pairs with + ``dtype=wp.vec2i``. ``None`` uses the model's contact-pair list. + """ + + narrow_phase: Any | None = None + """Optional prebuilt Newton narrow-phase object for expert pipelines.""" + + sdf_hydroelastic_config: Any | None = None + """Optional Newton ``HydroelasticSDF.Config``-compatible object. + + ``None`` disables the hydroelastic pipeline. Individual participating + procedural meshes must also opt in through + :attr:`~embodichain.lab.sim.shapes.MeshCollisionCfg.is_hydroelastic`. + """ + + +@configclass +class NewtonPhysicsCfg(PhysicsBackendCfg): + """Configuration selector for the Newton physics backend. + + DexSim wraps and extends Newton for EmbodiChain. The selected solver and + collision pipeline are scene-wide. Shape, contact, material, and joint + values are configured separately on object and articulation configs and + compiled into DexSim Spawn descriptors. + """ + + device: str | torch.device = "cuda:0" + """Warp device used to build and step Newton, for example ``"cuda:0"``.""" + + num_substeps: int = 10 + """Number of Newton solver substeps per EmbodiChain physics step. + + The effective solver interval is ``physics_dt / num_substeps``. + """ + + requires_grad: bool = False + """Whether to finalize the Newton model with differentiable state enabled. + + EmbodiChain currently requires the Semi-implicit solver for this mode and + disables CUDA graph capture when gradients are enabled. + """ + + use_cuda_graph: bool = True + """Whether to capture Newton stepping in a CUDA graph when supported. + + This is ignored for gradient mode and is unavailable on a CPU device. + """ + + debug_mode: bool = False + """Whether to enable additional Newton runtime diagnostics.""" + + suppress_warp_kernel_logs: bool = True + """Whether to hide Warp startup and kernel compile/load messages. + + Genuine Newton/Warp warnings and errors are not suppressed. + """ + + 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, EmbodiChain + preserves DexSim's scene-aware ``AutoSolverCfg`` default. A DexSim build + exporting ``AutoSolverCfg`` is required; no concrete-solver fallback is used. + """ + + collision_cfg: NewtonCollisionPipelineCfg | Mapping[str, Any] = field( + default_factory=NewtonCollisionPipelineCfg + ) + """Scene-level Newton collision-pipeline configuration.""" + + enable_collision_pipeline: bool = True + """Whether Newton generates rigid contacts before each solver substep. + + Disable this only for a solver/workflow that deliberately obtains contacts + elsewhere; ordinary rigid-body scenes require it. + """ + + broad_phase: Literal["nxn", "sap", "explicit"] | None = None + """Deprecated shortcut for ``collision_cfg.broad_phase``. + + If both are set, ``collision_cfg.broad_phase`` wins. + """ + + visualizer_enabled: bool = False + """Whether to enable DexSim Newton's optional diagnostic visualizer.""" + + def __post_init__(self) -> None: + """Normalize dictionary collision settings at the config boundary.""" + if isinstance(self.collision_cfg, Mapping): + self.collision_cfg = NewtonCollisionPipelineCfg(**self.collision_cfg) + + def to_dexsim_cfg( + self, + gpu_id: int, + ) -> NewtonCfg: + """Convert this config to ``dexsim.engine.newton_physics.NewtonCfg``.""" + from dexsim.engine.newton_physics import ( + AutoSolverCfg, + 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: dict[str, type] = { + "auto": AutoSolverCfg, + "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 is None or solver_cfg.solver_type != "semi_implicit" + ): + logger.log_error( + "Newton gradient mode requires an explicit " + "solver_type='semi_implicit'; AutoSolver does not select a " + "differentiable solver." + ) + + collision_values = { + item.name: getattr(self.collision_cfg, item.name) + for item in fields(self.collision_cfg) + } + if collision_values["broad_phase"] is None: + collision_values["broad_phase"] = self.broad_phase + collision_values["requires_grad"] = self.requires_grad + + newton_cfg_args: dict[str, Any] = { + "dt": self.physics_dt, + "num_substeps": self.num_substeps, + "device": device, + "gravity": _gravity_vector(self.gravity), + "debug_mode": self.debug_mode, + "requires_grad": self.requires_grad, + "suppress_warp_kernel_logs": self.suppress_warp_kernel_logs, + "collision_pipeline_cfg": NewtonCollisionPipelineCfg(**collision_values), + "enable_collision_pipeline": self.enable_collision_pipeline, + "sync_to_dexsim": True, + } + if solver_cfg is not None: + newton_cfg_args["solver_cfg"] = solver_cfg + + cfg = NewtonCfg( + **newton_cfg_args, + ) + 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 = { + "auto": "auto", + "autosolver": "auto", + "autosolvercfg": "auto", + "auto_solver": "auto", + "auto_solver_cfg": "auto", + "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 'auto', '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 | None: + """Convert EmbodiChain Newton solver config input to a DexSim config.""" + if solver_cfg is None: + return None + + 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 "auto" + ) + normalized_solver_type = _normalize_newton_solver_type(str(configured_solver_type)) + solver_cfg_type = solver_cfg_map[normalized_solver_type] + return solver_cfg_type(**solver_cfg_data) + + +def physics_cfg_for_backend( + backend: Literal["default", "newton"], +) -> PhysicsBackendCfg: + """Return a default physics configuration instance for the given backend.""" + if backend == "newton": + return NewtonPhysicsCfg() + if backend == "default": + return DefaultPhysicsCfg() + raise ValueError( + f"Unsupported physics backend {backend!r}; expected 'default' or 'newton'." + ) + + +def physics_backend_from_cfg( + physics_cfg: PhysicsBackendCfg, +) -> 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: PhysicsBackendCfg) -> None: + """Validate that ``physics_cfg`` is a supported backend configuration.""" + physics_backend_from_cfg(physics_cfg) diff --git a/embodichain/lab/sim/cfg/urdf.py b/embodichain/lab/sim/cfg/urdf.py new file mode 100644 index 000000000..49a82f4d8 --- /dev/null +++ b/embodichain/lab/sim/cfg/urdf.py @@ -0,0 +1,414 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""URDF assembly configuration.""" + +from __future__ import annotations + +from dataclasses import field +import os +from typing import Any, Dict, List + +import numpy as np + +from embodichain.data.constants import EMBODICHAIN_DEFAULT_DATA_ROOT +from embodichain.utils import configclass, logger + + +def _get_data_path(path: str) -> str: + """Resolve data through the public facade for monkeypatch compatibility.""" + from . import get_data_path + + return get_data_path(path) + + +@configclass +class URDFCfg: + """Standalone configuration class for URDF assembly.""" + + components: Dict[str, Dict[str, str | Dict | np.ndarray]] = field( + default_factory=dict + ) + """Dictionary of robot components to be assembled.""" + + sensors: Dict[str, Dict[str, str | np.ndarray]] = field(default_factory=dict) + """Dictionary of sensors to be attached to the robot.""" + + use_signature_check: bool = True + """Whether to use signature check when merging URDFs.""" + + base_link_name: str = "base_link" + """Name of the base link in the assembled robot.""" + + fpath: str | None = None + """Full output file path for the assembled URDF. If specified, overrides fname and fpath_prefix.""" + + fname: str | None = None + """Name used for output file and directory. If not specified, auto-generated from component names.""" + + fpath_prefix: str = EMBODICHAIN_DEFAULT_DATA_ROOT + "/assembled" + """Output directory prefix for the assembled URDF file.""" + + component_prefix: List[tuple[str, str | None]] = field( + default_factory=lambda: [ + ("chassis", None), + ("legs", None), + ("torso", None), + ("head", None), + ("left_arm", "left_"), + ("right_arm", "right_"), + ("left_hand", "left_"), + ("right_hand", "right_"), + ("arm", None), + ("hand", None), + ] + ) + """Component name prefixes used during URDF assembly. + + Preferred form is a list of ``(component_name, prefix)`` tuples. For + convenience, a mapping ``{component_name: prefix}`` is also accepted when + constructing :class:`URDFCfg` and will be normalized internally. + """ + + name_case: dict[str, str] = field( + default_factory=lambda: { + "joint": "original", + "link": "original", + } + ) + """Case normalization policy applied to joint/link names during URDF assembly. + + Supported values per key are ``"upper"``, ``"lower"`` or ``"original"`` + (legacy alias ``"none"``). The default preserves source URDF casing. + """ + + def __init__( + self, + components: list[dict[str, str | np.ndarray]] | None = None, + sensors: dict[str, dict[str, str | np.ndarray]] | None = None, + fpath: str | None = None, + fname: str | None = None, + fpath_prefix: str = EMBODICHAIN_DEFAULT_DATA_ROOT + "/assembled", + use_signature_check: bool = True, + base_link_name: str = "base_link", + component_prefix: list[tuple[str, str | None]] | None = None, + name_case: dict[str, str] | None = None, + ): + """ + Initialize URDFCfg with optional list of components and output path settings. + + Args: + components (list[dict[str, str | np.ndarray]] | None): List of component configurations. Each dict should contain: + - 'component_type' (str): The type/name of the component (e.g., 'chassis', 'arm', 'hand'). + - 'urdf_path' (str): Path to the component's URDF file. + - 'transform' (np.ndarray | None): 4x4 transformation matrix (optional). + - Additional params can be included as extra keys. + sensors (dict[str, dict[str, str | np.ndarray]] | None): Sensor configurations for the robot. + fpath (str | None): Full output file path for the assembled URDF. If specified, overrides fname and fpath_prefix. + fname (str | None): Name used for output file and directory. If not specified, auto-generated from component names. + fpath_prefix (str): Output directory prefix for the assembled URDF file. + use_signature_check (bool): Whether to use signature check when merging URDFs. + base_link_name (str): Name of the base link in the assembled robot. + component_prefix (list[tuple[str, str | None]] | None): Optional + list of (component_type, prefix) pairs to override default + component name prefixes. + """ + self.components = {} + self.sensors = sensors or {} + self.fpath = fpath + self.use_signature_check = use_signature_check + self.base_link_name = base_link_name + self.fname = fname + self.fpath_prefix = fpath_prefix + + # Initialize component prefixes (patch-style mapping per component type) + if component_prefix is None: + # Use the same default as the dataclass field + self.component_prefix = [ + ("chassis", None), + ("legs", None), + ("torso", None), + ("head", None), + ("left_arm", "left_"), + ("right_arm", "right_"), + ("left_hand", "left_"), + ("right_hand", "right_"), + ("arm", None), + ("hand", None), + ] + elif isinstance(component_prefix, dict): + # Allow dict-style config: {"left_hand": "l_", ...} + self.component_prefix = list(component_prefix.items()) + else: + # Assume caller provided a list of (component_name, prefix) tuples + self.component_prefix = component_prefix + + if name_case is None: + self.name_case = { + "joint": "original", + "link": "original", + } + else: + self.name_case = name_case + + # Auto-add components if provided + if components: + for comp_config in components: + if not isinstance(comp_config, dict): + logger.log_error( + f"Component configuration must be a dict, got {type(comp_config)}" + ) + continue + + # Extract required fields + component_type = comp_config.get("component_type") + urdf_path = comp_config.get("urdf_path") + + if not component_type or not urdf_path: + logger.log_error( + f"Component configuration must contain 'component_type' and 'urdf_path', got {comp_config}" + ) + continue + + # Extract optional fields + transform = comp_config.get("transform", np.eye(4)) + + # Extract additional params (exclude known keys) + params = { + k: v + for k, v in comp_config.items() + if k not in ["component_type", "urdf_path", "transform"] + } + + # Add the component + self.add_component(component_type, urdf_path, transform, **params) + + if sensors is not None: + # Accept both list and dict; serialization round-trips an empty + # dict when no sensors are configured (the field default). + if isinstance(sensors, dict) and not sensors: + self.sensors = [] + elif not isinstance(sensors, (list, dict)): + logger.log_error( + f"sensors must be a list of dicts or a dict, got {type(sensors)}" + ) + self.sensors = [] + elif isinstance(sensors, dict): + # dict keyed by sensor_name -> config + self.sensors = list(sensors.values()) + else: + # Optionally check each sensor dict + valid_sensors = [] + for sensor_config in sensors: + if not isinstance(sensor_config, dict): + logger.log_error( + f"Sensor configuration must be a dict, got {type(sensor_config)}" + ) + continue + sensor_name = sensor_config.get("sensor_name") + if not sensor_name: + logger.log_error( + f"Sensor configuration must contain 'sensor_name', got {sensor_config}" + ) + continue + valid_sensors.append(sensor_config) + self.sensors = valid_sensors + + def set_urdf(self, urdf_path: str) -> "URDFCfg": + """Directly specify a single URDF file for the robot, compatible with the single-URDF robot case. + + Args: + urdf_path (str): Path to the robot's URDF file. + + Returns: + URDFCfg: Returns self to allow method chaining. + """ + self.components.clear() + urdf_file = os.path.splitext(os.path.basename(urdf_path))[0] + self.components[urdf_file] = { + "urdf_path": urdf_path, + "transform": None, + "params": {}, + } + self.fpath = urdf_path + return self + + def add_component( + self, + component_type: str, + urdf_path: str, + transform: np.ndarray | None = None, + **params, + ) -> URDFCfg: + """Add a robot component to the assembly configuration. + + Args: + component_type (str): The type/name of the component. Should be one of SUPPORTED_COMPONENTS + (e.g., 'chassis', 'torso', 'head', 'left_arm', 'right_hand', 'arm', 'hand', etc.). + urdf_path (str): Path to the component's URDF file. + transform (np.ndarray | None): 4x4 transformation matrix for the component in the robot frame (default: None). + **params: Additional keyword parameters for the component (e.g., color, material, etc.). + + Returns: + URDFCfg: Returns self to allow method chaining. + """ + if urdf_path: + if not os.path.exists(urdf_path): + urdf_path_candidate = _get_data_path(urdf_path) + if os.path.exists(urdf_path_candidate): + urdf_path = urdf_path_candidate + else: + logger.log_error(f"URDF path '{urdf_path}' does not exist.") + raise FileNotFoundError(f"URDF path '{urdf_path}' does not exist.") + + if transform is None: + transform = np.eye(4) + + self.components[component_type] = { + "urdf_path": urdf_path, + "transform": np.array(transform), + "params": params, + } + + if self.fname: + self.fpath = f"{self.fpath_prefix}/{self.fname}/{self.fname}.urdf" + else: + # Update output_path to use all component urdf file names joined by underscores as directory + if len(self.components) == 1: + # Only one component, use its urdf file name + urdf_file = os.path.splitext(os.path.basename(urdf_path))[0] + name = urdf_file + else: + # Multiple components, join all urdf file names + urdf_files = [ + os.path.splitext(os.path.basename(v["urdf_path"]))[0] + for v in self.components.values() + ] + name = "_".join(urdf_files) + self.fpath = f"{self.fpath_prefix}/{name}/{name}.urdf" + + return self + + def add_sensor(self, sensor_name: str, **sensor_config) -> URDFCfg: + """Add a sensor to the robot configuration. + + Args: + sensor_name (str): The name of the sensor. + **sensor_config: Additional configuration parameters for the sensor. + + Returns: + URDFCfg: Returns self to allow method chaining. + """ + self.sensors.append({"sensor_name": sensor_name, **sensor_config}) + return self + + def assemble_urdf(self) -> str: + """Assemble URDF files for the robot based on the configuration. + + Returns: + str: The path to the resulting (possibly merged) URDF file. + """ + components = list(self.components.items()) + # If there is only one component, return its URDF path directly. + if len(components) == 1: + _, comp_config = components[0] + return comp_config["urdf_path"] + + from embodichain.toolkits.urdf_assembly import URDFAssemblyManager + + # If there are multiple components, merge them into a single URDF file. + manager = URDFAssemblyManager() + manager.base_link_name = self.base_link_name + + if self.component_prefix is None: + self.component_prefix = [ + ("left_arm", "left_"), + ("right_arm", "right_"), + ("left_hand", "left_"), + ("right_hand", "right_"), + ] + if isinstance(self.component_prefix, dict): + self.component_prefix = list(self.component_prefix.items()) + # Forward configured component prefixes to the assembly manager + manager.component_prefix = self.component_prefix + + if self.name_case is not None: + manager.name_case = self.name_case + + for comp_type, comp_config in components: + params = comp_config.get("params", {}) + success = manager.add_component( + comp_type, + comp_config["urdf_path"], + comp_config.get("transform"), + **params, + ) + if not success: + logger.log_error( + f"Failed to add component '{comp_type}' with config: {comp_config}" + ) + + for sensor in self.sensors: + manager.attach_sensor( + sensor_name=sensor.get("sensor_name"), + sensor_source=sensor.get("sensor_source"), + parent_component=sensor.get("parent_component"), + parent_link=sensor.get("parent_link"), + sensor_type=sensor.get("sensor_type"), + **{ + k: v + for k, v in sensor.items() + if k + not in [ + "sensor_name", + "sensor_source", + "parent_component", + "parent_link", + "sensor_type", + ] + }, + ) + + try: + # Merge all added components into a single URDF file at the specified output path. + merged_urdf_xml = manager.merge_urdfs(self.fpath, self.use_signature_check) + except Exception as e: + logger.log_error(f"URDF merge failed: {e}") + + return self.fpath + + @classmethod + def from_dict(cls, init_dict: Dict) -> "URDFCfg": + if isinstance(init_dict, cls): + return init_dict + components = init_dict.get("components", None) + if isinstance(components, dict): + components = [{"component_type": k, **v} for k, v in components.items()] + sensors = init_dict.get("sensors", None) + fpath = init_dict.get("fpath", None) + use_signature_check = init_dict.get("use_signature_check", True) + base_link_name = init_dict.get("base_link_name", "base_link") + component_prefix = init_dict.get("component_prefix", None) + name_case = init_dict.get("name_case", None) + return cls( + components=components, + sensors=sensors, + fpath=fpath, + use_signature_check=use_signature_check, + base_link_name=base_link_name, + component_prefix=component_prefix, + name_case=name_case, + ) diff --git a/embodichain/lab/sim/cfg/viewer.py b/embodichain/lab/sim/cfg/viewer.py new file mode 100644 index 000000000..35710cda4 --- /dev/null +++ b/embodichain/lab/sim/cfg/viewer.py @@ -0,0 +1,94 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Interactive viewer, marker, and recording configuration.""" + +from __future__ import annotations + +from typing import List, Literal + +import torch +from dexsim.types import AxisArrowType, AxisCornerType + +from embodichain.utils import configclass + + +@configclass +class MarkerCfg: + """Configuration for visual markers in the simulation. + + This class defines properties for creating visual markers such as coordinate frames, + lines, and points that can be used for debugging, visualization, or reference purposes + in the simulation environment. + """ + + name: str = "empty-mesh" + """Name of the marker for identification purposes.""" + + marker_type: Literal["axis", "line", "point"] = "axis" + """Type of marker to display. Can be 'axis' (3D coordinate frame), 'line', or 'point'. (only axis supported now)""" + + axis_xpos: torch.Tensor | None = None + """List of 4x4 transformation matrices defining the position and orientation of each axis marker.""" + + axis_size: float = 0.002 + """Thickness/size of the axis lines in meters.""" + + axis_len: float = 0.005 + """Length of each axis arm in meters.""" + + line_color: List[float] = [1, 1, 0, 1.0] + """RGBA color values for the marker lines. Values should be between 0.0 and 1.0.""" + + arrow_type: AxisArrowType = AxisArrowType.CONE + """Type of arrow head for axis markers (e.g., CONE, ARROW, etc.).""" + + corner_type: AxisCornerType = AxisCornerType.SPHERE + """Type of corner/joint visualization for axis markers (e.g., SPHERE, CUBE, etc.).""" + + arena_index: int = -1 + """Index of the arena where the marker should be placed. -1 means all arenas.""" + + +@configclass +class WindowRecordCfg: + """Configuration for interactive viewer window recording.""" + + enable_hotkey: bool = True + """Whether to register the ``r`` hotkey for viewer recording when the window opens.""" + + save_path: str | None = None + """Optional output path for viewer recordings. If None, use the default outputs directory.""" + + fps: int = 20 + """Frames per second for viewer recording.""" + + max_memory: int = 1024 + """Maximum buffered recording memory in MB before auto-stopping capture.""" + + video_prefix: str = "viewer_record" + """Video file prefix used when no explicit save path is provided.""" + + +@configclass +class WindowCameraPoseCfg: + """Configuration for printing the interactive viewer camera pose.""" + + enable_hotkey: bool = True + """Whether to register the ``p`` hotkey when the window opens.""" + + convert_to_look_at: bool = True + """Whether the hotkey prints a ``set_look_at`` call instead of a matrix.""" diff --git a/embodichain/lab/sim/common.py b/embodichain/lab/sim/common.py index f1380ed6b..a578fb9c7 100644 --- a/embodichain/lab/sim/common.py +++ b/embodichain/lab/sim/common.py @@ -66,8 +66,6 @@ def __init__( self._entities = entities self.device = device - 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 000000000..ad84e89a7 --- /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 manager-owned differentiable trajectory transaction 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 000000000..29d0fe5e0 --- /dev/null +++ b/embodichain/lab/sim/diff/bridge.py @@ -0,0 +1,371 @@ +# ---------------------------------------------------------------------------- +# 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 +import math +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"] + + +def _differentiable_runtime(manager: Any) -> Any: + """Resolve Spawn's runtime while retaining lightweight test compatibility.""" + runtime = getattr(manager, "differentiable_runtime", None) + if runtime is not None: + return runtime + return manager.physics.newton_manager + + +def _physics_dt(nm: Any, sim_state: dict[str, Any]) -> float: + """Resolve the outer Newton step duration represented by one control step.""" + physics_dt = sim_state.get("physics_dt") + if physics_dt is None: + physics_dt = float(nm.solver_dt) * int(nm.num_substeps) + try: + physics_dt = float(physics_dt) + except (TypeError, ValueError) as exc: + raise TypeError("physics_dt must be a positive finite float.") from exc + if not math.isfinite(physics_dt) or physics_dt <= 0.0: + raise ValueError("physics_dt must be a positive finite float.") + return physics_dt + + +def _resolve_step_mode(sim_state: dict[str, Any]) -> tuple[str, Callable | None]: + """Validate the explicit dynamics-versus-kinematics bridge contract.""" + step_mode = sim_state.get("step_mode", "dynamics") + if step_mode not in {"dynamics", "kinematics"}: + raise ValueError( + "step_mode must be 'dynamics' or 'kinematics', " f"got {step_mode!r}." + ) + + step_fn = sim_state.get("step_fn") + if step_mode == "dynamics" and step_fn is not None: + raise ValueError( + "step_fn is only supported when step_mode='kinematics'; " + "the dynamics route always uses Newton solver dynamics." + ) + if step_mode == "kinematics" and step_fn is None: + raise ValueError("step_mode='kinematics' requires a named step_fn.") + return step_mode, step_fn + + +def _reset_tape_then_release(tape: wp.Tape | None, trajectory: Any | None) -> None: + """End tape ownership before releasing the trajectory's model lease.""" + try: + if tape is not None: + tape.reset() + finally: + if trajectory is not None: + trajectory.release() + + +def _abort_forward( + tape: wp.Tape | None, + trajectory: Any | None, +) -> None: + """Best-effort cleanup which never masks the original forward failure.""" + try: + _reset_tape_then_release(tape, trajectory) + except BaseException: + # The active trajectory must not mask the action/solver/output failure + # which caused the abort. DexSim release is idempotent and this path is + # only entered to preserve the original exception. + pass + + +@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, Any], None], + substeps: int, + dt: float | None = None, +) -> dict[str, Any]: + """Run a low-level manager-owned Newton trajectory inside a Warp tape. + + ``substeps`` remains a legacy solver-step count. It must therefore divide + evenly into whole Newton physics steps; the public trajectory transaction + owns every detached state, contact, control, and generation lease. + + The returned tape and trajectory remain active for the caller to use in a + custom backward pass. After ``tape.backward()`` (or when abandoning the + result), callers must invoke ``tape.reset()`` and then + ``trajectory.release()`` in a ``finally`` block. The helper releases both + automatically only when forward construction itself fails. + + Args: + manager: The owning :class:`SimulationManager` (must be Newton). + apply_control_fn: Callable that writes the trajectory-local joint/body + control targets inside the tape. It receives ``(tape, control)`` + and must launch Warp kernels targeting ``control``, never the + manager's shared control buffer. + substeps: Number of solver substeps to run. + dt: Solver dt; defaults to the manager's configured solver dt. + + Returns: + A dict carrying the tape, trajectory, and detached final state for the + caller to retain through backward before resetting/releasing it. + """ + if not manager.is_newton_backend: + raise RuntimeError("differentiable_step requires the Newton backend.") + nm = _differentiable_runtime(manager) + if isinstance(substeps, bool) or int(substeps) != substeps or substeps <= 0: + raise ValueError("substeps must be a positive integer.") + substeps = int(substeps) + num_substeps = int(nm.num_substeps) + if num_substeps <= 0: + raise ValueError("Newton num_substeps must be positive.") + if substeps % num_substeps != 0: + raise ValueError( + "substeps must be divisible by Newton num_substeps so the " + "trajectory represents whole physics steps." + ) + dt_val = float(nm.solver_dt if dt is None else dt) + if not math.isfinite(dt_val) or dt_val <= 0.0: + raise ValueError("dt must be a positive finite solver time step.") + + trajectory = None + tape = None + try: + trajectory = nm.create_differentiable_trajectory( + physics_steps=substeps // num_substeps, + physics_dt=dt_val * num_substeps, + ) + tape = wp.Tape() + with tape: + apply_control_fn(tape, trajectory.control) + final_state = trajectory.step() + nm.commit_differentiable_trajectory(trajectory) + except BaseException: + _abort_forward(tape, trajectory) + raise + + return { + "tape": tape, + "trajectory": trajectory, + "final_state": final_state, + "states": trajectory.states, + "contacts": trajectory.contacts, + "control": trajectory.control, + } + + +class NewtonStepFunc(torch.autograd.Function): + """torch.autograd.Function bridging Warp tape autodiff to PyTorch. + + Forward: validates an explicit step mode before creating a tape. The + default ``dynamics`` route allocates a manager-owned detached trajectory, + launches the action-to-local-control Warp kernel, records its solver + horizon, and commits it only after tape exit. The explicitly selected + ``kinematics`` route retains its named FK ``step_fn`` escape hatch. + Observation/reward kernels run inside the tape so 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 control-level physics updates (used by the default + solver-based step route) + step_mode: ``"dynamics"`` (default) or explicit ``"kinematics"`` + action_to_control_kernel: dynamics callable + ``(action_wp, trajectory_control, *kernel_args)``; kinematics + retains ``(action_wp, tape, *kernel_args)`` + kernel_args: tuple consumed by action_to_control_kernel + obs_reward_fn: callable(final_state) -> dict with torch outputs + physics_dt: optional outer Newton step duration (defaults to + ``solver_dt * num_substeps``) + step_fn: required only when ``step_mode == "kinematics"`` + + 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`` + """ + + @classmethod + def apply(cls, action_torch: torch.Tensor, sim_state: dict[str, Any]) -> Any: + """Capture the caller's grad mode before PyTorch enters ``forward``. + + ``torch.autograd.Function.forward`` always executes with grad mode + disabled, and ``ctx.needs_input_grad`` alone remains true when a + requires-grad action is passed through an outer ``torch.no_grad()`` + block. Passing the ambient mode as a non-differentiable argument lets + the bridge synchronously reset/release no-grad trajectories instead of + retaining an unreachable manager lease. + """ + return super().apply(action_torch, sim_state, torch.is_grad_enabled()) + + @staticmethod + def forward( + ctx: Any, + action_torch: torch.Tensor, + sim_state: dict[str, Any], + outer_grad_enabled: bool, + ) -> tuple[torch.Tensor, ...]: + 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_mode, step_fn = _resolve_step_mode(sim_state) + tape_binder = ( + sim_state.get("_bind_dynamics_tape") if step_mode == "dynamics" else None + ) + + # Save the original action shape so backward can reshape the gradient. + ctx.saved_action_shape = action_torch.shape + + nm = _differentiable_runtime(manager) + + action_flat = action_torch.detach().clone().reshape(-1).contiguous() + needs_action_grad = bool(outer_grad_enabled and ctx.needs_input_grad[0]) + action_wp = wp.from_torch( + action_flat, + dtype=wp.float32, + requires_grad=needs_action_grad, + ) + + trajectory = None + tape = None + try: + if step_mode == "dynamics": + if substeps <= 0: + raise ValueError("substeps must be a positive integer.") + trajectory = nm.create_differentiable_trajectory( + physics_steps=substeps, + physics_dt=_physics_dt(nm, sim_state), + ) + + tape = wp.Tape() + try: + with tape: + if tape_binder is not None: + tape_binder(tape) + if step_mode == "dynamics": + kernel(action_wp, trajectory.control, *kernel_args) + final_state = trajectory.step() + else: + # The explicit FK route keeps the historical callback + # shape and receives the open tape, but never detached + # solver control. + kernel(action_wp, tape, *kernel_args) + final_state = step_fn() + + # Validate and materialize outputs inside the tape. A malformed + # output dictionary is a forward failure and must not publish a + # detached dynamics trajectory. + outputs = obs_reward_fn(final_state) + outputs_order = tuple(outputs["_order"]) + output_values = tuple(outputs[name] for name in outputs_order) + outputs_grad_track = outputs.get("_grad_track", {}) + finally: + if tape_binder is not None: + tape_binder(None) + + if trajectory is not None: + nm.commit_differentiable_trajectory(trajectory) + except BaseException: + _abort_forward(tape, trajectory) + raise + + if not needs_action_grad: + _reset_tape_then_release(tape, trajectory) + return output_values + + ctx.tape = tape + ctx.trajectory = trajectory + ctx.action_wp = action_wp + ctx.outputs_order = outputs_order + ctx.outputs_grad_track = outputs_grad_track + ctx._bridge_released = False + return output_values + + @staticmethod + def backward( + ctx: Any, + *grad_outputs: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, None, None]: + if getattr(ctx, "_bridge_released", False): + raise RuntimeError( + "NewtonStepFunc backward was already consumed; create a new " + "differentiable trajectory for another backward pass." + ) + + action_grad = None + try: + # 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_wp_grad = getattr(ctx.action_wp, "grad", None) + if action_wp_grad is not None: + # Capture the action gradient before reset invalidates tape + # storage, then terminate tape ownership before releasing the + # trajectory's active manager token. + action_grad = wp.to_torch(action_wp_grad).clone() + finally: + try: + _reset_tape_then_release(ctx.tape, ctx.trajectory) + finally: + ctx._bridge_released = True + + if action_grad is None: + return None, None, None + # Reshape to the original action layout; metadata inputs have no + # gradient. + return action_grad.reshape(ctx.saved_action_shape), None, None diff --git a/embodichain/lab/sim/diff/runtime.py b/embodichain/lab/sim/diff/runtime.py new file mode 100644 index 000000000..c6c46bdfb --- /dev/null +++ b/embodichain/lab/sim/diff/runtime.py @@ -0,0 +1,344 @@ +# ---------------------------------------------------------------------------- +# 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 transactions over a Spawn-owned Newton runtime.""" + +from __future__ import annotations + +import math +from typing import Any, Callable + +__all__ = ["NewtonDifferentiableRuntime"] + + +class NewtonDifferentiableTrajectory: + """Own the detached buffers for one differentiable Newton trajectory.""" + + def __init__( + self, + runtime: "NewtonDifferentiableRuntime", + *, + physics_steps: int, + physics_dt: float, + ) -> None: + self._runtime = runtime + self._backend = runtime._validated_backend() + self.physics_steps = int(physics_steps) + self.physics_dt = float(physics_dt) + self.total_solver_steps = self.physics_steps * runtime.num_substeps + self.solver_dt = self.physics_dt / runtime.num_substeps + + model = self._backend.model + self.states = [model.state() for _ in range(self.total_solver_steps + 1)] + self.states[0].assign(self._backend.runtime.current_state) + self.control = model.control() + self.contacts = [ + self._backend.collision_pipeline.contacts() + for _ in range(self.total_solver_steps) + ] + self._stepped = False + self._committed = False + self._released = False + + @property + def final_state(self) -> Any: + """Return the terminal state owned by this trajectory.""" + return self.states[-1] + + def step(self) -> Any: + """Run the complete trajectory inside the caller's active Warp tape.""" + if self._released: + raise RuntimeError("Cannot step a released differentiable trajectory.") + if self._stepped: + raise RuntimeError("A differentiable trajectory can only be stepped once.") + if self._runtime._backend() is not self._backend: + raise RuntimeError( + "The Spawn-owned Newton backend changed while a differentiable " + "trajectory was active. Release it and create a fresh trajectory." + ) + + backend = self._backend + apply_external_wrenches = backend.runtime.has_external_wrenches + for index, (state_in, state_out, contacts) in enumerate( + zip(self.states, self.states[1:], self.contacts) + ): + state_in.clear_forces() + if apply_external_wrenches and index < self._runtime.num_substeps: + backend.runtime.apply_external_wrenches(state_in) + if backend.cfg.enable_collision_pipeline: + backend.collision_pipeline.collide(state_in, contacts) + backend.solver.step( + state_in, + state_out, + self.control, + contacts, + self.solver_dt, + ) + self._stepped = True + return self.final_state + + def release(self) -> None: + """Release the runtime lease after the owning Warp tape is reset.""" + if self._released: + return + self._runtime._release_differentiable_trajectory(self) + self._released = True + + +class NewtonDifferentiableRuntime: + """Adapt the current Spawn-owned Newton backend to the autograd bridge. + + The provider is resolved for every public operation so a scene rebuild + cannot silently publish a trajectory into a replaced Newton backend. + """ + + def __init__(self, backend_provider: Callable[[], Any]) -> None: + self._backend_provider = backend_provider + self._active_trajectory: NewtonDifferentiableTrajectory | None = None + + def _backend(self) -> Any: + backend = self._backend_provider() + if backend is None: + raise RuntimeError( + "The Spawn-owned Newton backend is unavailable. Call " + "SimulationManager.prepare() before using differentiable physics." + ) + return backend + + def _validated_backend(self) -> Any: + backend = self._backend() + if backend.model is None: + raise RuntimeError( + "The Spawn-owned Newton model is not finalized. Call " + "SimulationManager.prepare() first." + ) + if not bool(backend.cfg.requires_grad): + raise RuntimeError( + "Differentiable Newton physics requires requires_grad=True." + ) + if backend.cfg.solver_cfg.solver_type != "semi_implicit": + raise RuntimeError( + "Differentiable Newton physics requires " "solver_type='semi_implicit'." + ) + if backend.collision_pipeline is None: + raise RuntimeError( + "Differentiable Newton physics requires a collision pipeline." + ) + if getattr(backend, "_runtime_controls", ()): + raise RuntimeError( + "Differentiable trajectories do not support Spawn runtime " + "controls yet. Remove them before finalizing the scene." + ) + return backend + + @property + def model(self) -> Any: + """Return the finalized Newton model for expert Warp operations.""" + return self._validated_backend().model + + @property + def current_state(self) -> Any: + """Return the live state currently selected by the Spawn runtime.""" + return self._validated_backend().runtime.current_state + + @property + def live_states(self) -> tuple[Any, Any]: + """Return both live ping-pong states owned by the Spawn backend.""" + backend = self._validated_backend() + return backend.state_0, backend.state_1 + + @property + def control(self) -> Any: + """Return the live Spawn control buffer.""" + return self._validated_backend().control + + @property + def num_substeps(self) -> int: + """Return the number of Newton solver substeps per physics step.""" + return max(int(self._validated_backend().cfg.num_substeps), 1) + + @property + def physics_dt(self) -> float: + """Return the configured outer physics-step duration.""" + return float(self._validated_backend().cfg.dt) + + @property + def solver_dt(self) -> float: + """Return the configured Newton solver substep duration.""" + return self.physics_dt / self.num_substeps + + # Compatibility aliases consumed by DexSim's low-level differentiable + # stepper/rollout helpers. They borrow, but never own, Spawn resources. + @property + def _model(self) -> Any: + return self.model + + @property + def _state_0(self) -> Any: + return self._validated_backend().state_0 + + @property + def _state_1(self) -> Any: + return self._validated_backend().state_1 + + @property + def _control(self) -> Any: + return self.control + + @property + def _solver(self) -> Any: + return self._validated_backend().solver + + @property + def _collision_pipeline(self) -> Any: + return self._validated_backend().collision_pipeline + + @property + def _external_forces(self) -> Any: + return self._validated_backend().runtime.external_wrenches + + def _ensure_external_force_buffers(self) -> None: + self._validated_backend() + + def clear_external_forces(self) -> None: + """Clear pending Spawn runtime wrenches.""" + self._validated_backend().runtime.clear_external_wrenches() + + def create_differentiable_trajectory( + self, + *, + physics_steps: int, + physics_dt: float, + ) -> NewtonDifferentiableTrajectory: + """Allocate one detached trajectory and acquire the runtime lease.""" + if isinstance(physics_steps, bool) or int(physics_steps) != physics_steps: + raise TypeError("physics_steps must be a positive integer.") + physics_steps = int(physics_steps) + if physics_steps <= 0: + raise ValueError("physics_steps must be a positive integer.") + try: + physics_dt = float(physics_dt) + except (TypeError, ValueError) as exc: + raise TypeError("physics_dt must be a positive finite float.") from exc + if not math.isfinite(physics_dt) or physics_dt <= 0.0: + raise ValueError("physics_dt must be a positive finite float.") + if self._active_trajectory is not None: + raise RuntimeError( + "A differentiable trajectory is still active; release it after " + "backward before creating another trajectory." + ) + + trajectory = NewtonDifferentiableTrajectory( + self, + physics_steps=physics_steps, + physics_dt=physics_dt, + ) + self._active_trajectory = trajectory + return trajectory + + def commit_differentiable_trajectory( + self, + trajectory: NewtonDifferentiableTrajectory, + ) -> None: + """Publish one detached terminal state back to the live Spawn runtime.""" + if self._active_trajectory is not trajectory: + raise RuntimeError( + "The differentiable trajectory is not active on this runtime." + ) + if trajectory._released: + raise RuntimeError("Cannot commit a released differentiable trajectory.") + if trajectory._committed: + raise RuntimeError( + "A differentiable trajectory can only be committed once." + ) + if not trajectory._stepped: + raise RuntimeError( + "Step the differentiable trajectory before committing it." + ) + + backend = self._validated_backend() + if backend is not trajectory._backend: + raise RuntimeError( + "The Spawn-owned Newton backend changed before trajectory commit." + ) + backend.state_0.assign(trajectory.final_state) + backend.state_1.assign(trajectory.final_state) + backend.runtime.set_current_state(backend.state_0) + backend.runtime.clear_external_wrenches() + backend.set_sim_time( + backend.sim_time + trajectory.physics_steps * trajectory.physics_dt, + backend.step_index + trajectory.physics_steps, + ) + trajectory._committed = True + + def _release_differentiable_trajectory( + self, + trajectory: NewtonDifferentiableTrajectory, + ) -> None: + if self._active_trajectory is not trajectory: + raise RuntimeError( + "The differentiable trajectory is not active on this runtime." + ) + self._active_trajectory = None + + def create_differentiable_stepper(self) -> Any: + """Create DexSim's low-level differentiable Newton step primitive.""" + self._validated_backend() + from dexsim.engine.newton_physics.differentiable_stepper import ( + DifferentiableStepper, + ) + + return DifferentiableStepper(self) + + def create_gradient_rollout( + self, + record_steps: int, + substeps_per_record: int | None = None, + record_dt: float | None = None, + ) -> Any: + """Create DexSim's standalone gradient-rollout buffers.""" + backend = self._validated_backend() + record_steps = int(record_steps) + if record_steps <= 0: + raise ValueError("record_steps must be positive.") + substeps = ( + self.num_substeps + if substeps_per_record is None + else int(substeps_per_record) + ) + if substeps <= 0: + raise ValueError("substeps_per_record must be positive.") + duration = self.physics_dt if record_dt is None else float(record_dt) + if not math.isfinite(duration) or duration <= 0.0: + raise ValueError("record_dt must be a positive finite float.") + + from dexsim.engine.newton_physics.gradient_rollout import GradientRollout + + total_substeps = record_steps * substeps + states = [backend.model.state() for _ in range(total_substeps + 1)] + states[0].assign(backend.runtime.current_state) + contacts = [ + backend.collision_pipeline.contacts() for _ in range(total_substeps) + ] + return GradientRollout( + self, + record_steps=record_steps, + substeps_per_record=substeps, + record_dt=duration, + states=states, + control=backend.model.control(), + contacts=contacts, + stepper=self.create_differentiable_stepper(), + ) diff --git a/embodichain/lab/sim/objects/__init__.py b/embodichain/lab/sim/objects/__init__.py index f74767b3c..6a608a1ef 100644 --- a/embodichain/lab/sim/objects/__init__.py +++ b/embodichain/lab/sim/objects/__init__.py @@ -19,6 +19,8 @@ Covers lights, rigid bodies (and groups), articulations, robots, deformables (soft/cloth), gizmos, and rigid constraints; every object derives from ``BatchEntity``. """ +from __future__ import annotations + from ..common import BatchEntity from .rigid_object import RigidObject, RigidBodyData, RigidObjectCfg from .rigid_object_group import ( @@ -26,8 +28,25 @@ RigidBodyGroupData, RigidObjectGroupCfg, ) -from .soft_object import SoftObject, SoftBodyData, SoftObjectCfg -from .cloth_object import ClothObject, ClothBodyData, ClothObjectCfg +from .deformable import ( + ClothBodyData, + ClothObject, + DeformableObject, + DeformableObjectData, + SoftBodyData, + SoftObject, + SurfaceDeformableData, + SurfaceDeformableObject, + VolumeDeformableData, + VolumeDeformableObject, +) +from ..cfg import ( + ClothObjectCfg, + DeformableObjectCfg, + SoftObjectCfg, + SurfaceDeformableObjectCfg, + VolumeDeformableObjectCfg, +) from .articulation import ( Articulation, ArticulationData, diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index 6499dfa15..4dd01be97 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/embodichain/lab/sim/objects/articulation.py @@ -22,9 +22,10 @@ import dexsim import numpy as np +from copy import deepcopy from dataclasses import dataclass from functools import cached_property -from typing import List, Sequence, Dict, Union, Tuple, Optional +from typing import TYPE_CHECKING, List, Sequence, Dict, Union, Tuple, Optional from dexsim.engine import Articulation as _Articulation from dexsim.types import ( @@ -43,10 +44,10 @@ _wrap_first_render_material, ) from embodichain.lab.sim.cfg import ( + _normalize_joint_target_mode, ArticulationCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, - RigidBodyAttributesOverrideCfg, + RigidBodyPhysicsCfg, ) from dexsim.types import PhysicalAttr from embodichain.utils.string import ( @@ -54,15 +55,26 @@ resolve_matching_names_values, ) from embodichain.lab.sim.common import BatchEntity +from embodichain.lab.sim.physics.newton import is_newton_gradient_mode +from embodichain.lab.sim.objects.backends import ( + DefaultArticulationView, + NewtonArticulationView, + SpawnArticulationView, + is_newton_scene, +) +from embodichain.lab.sim.objects.backends.base import ArticulationViewBase +from embodichain.lab.sim.objects.backends.newton import ( + _configure_newton_mimic_compliance, +) from embodichain.utils.math import ( + convert_quat, matrix_from_quat, quat_from_matrix, - convert_quat, matrix_from_euler, ) from embodichain.lab.sim.utility.sim_utils import ( + _apply_default_articulation_root_properties, get_dexsim_drive_type, - set_dexsim_articulation_cfg, ) from embodichain.lab.sim.utility.solver_utils import ( create_pk_chain, @@ -70,6 +82,19 @@ ) from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedArticulation + + +@dataclass(frozen=True, slots=True) +class _MimicInfo: + """Mimic metadata expressed in the backing state-buffer index domain.""" + + mimic_id: np.ndarray + mimic_parent: np.ndarray + mimic_multiplier: np.ndarray + mimic_offset: np.ndarray + @dataclass(frozen=True, slots=True, eq=False) class ArticulationJointKinematics: @@ -132,7 +157,11 @@ class ArticulationData: """GPU data manager for articulation.""" def __init__( - self, entities: List[_Articulation], ps: PhysicsScene, device: torch.device + self, + entities: Sequence[_Articulation | SpawnedArticulation], + ps: PhysicsScene | None, + device: torch.device, + articulation_view: ArticulationViewBase | None = None, ) -> None: """Initialize the ArticulationData. @@ -145,22 +174,28 @@ 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 articulation_view is not None: + self.articulation_view = articulation_view + elif is_newton_scene(ps): + self.articulation_view = NewtonArticulationView( + entities=entities, scene=ps, device=device + ) + else: + self.articulation_view = DefaultArticulationView( + entities=entities, ps=ps, device=device ) - if self.device.type == "cuda" - else None - ) - 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() + # Backward-compatible alias for callers that use GPU/articulation ids. + self.gpu_indices = self.articulation_view.articulation_ids_tensor + + if isinstance(self.articulation_view, SpawnArticulationView): + self.dof = self.articulation_view.dof + self.num_links = self.articulation_view.num_links + self.link_names = self.articulation_view.link_names + else: + self.dof = self.entities[0].get_dof() + self.num_links = self.entities[0].get_links_num() + self.link_names = self.entities[0].get_link_names() self._root_pose = torch.zeros( (self.num_instances, 7), dtype=torch.float32, device=self.device @@ -172,11 +207,13 @@ def __init__( (self.num_instances, 3), dtype=torch.float32, device=self.device ) - max_num_links = ( - self.ps.gpu_get_articulation_max_link_count() - if self.device.type == "cuda" - else self.num_links - ) + max_num_links = self.num_links + if ( + articulation_view is None + and self.device.type == "cuda" + and not self.is_newton_backend + ): + max_num_links = self.ps.gpu_get_articulation_max_link_count() self._body_link_pose = torch.zeros( (self.num_instances, max_num_links, 7), dtype=torch.float32, @@ -199,11 +236,35 @@ def __init__( device=self.device, ) - max_dof = ( - self.ps.gpu_get_articulation_max_dof() - if self.device.type == "cuda" - else self.dof + # Current link mass-property buffers use the public articulation link + # ordering. Initialization snapshots are captured after backend + # materialization and remain unchanged by runtime writes. + self._mass = torch.zeros( + (self.num_instances, self.num_links), + dtype=torch.float32, + device=self.device, + ) + self._inertia = torch.zeros( + (self.num_instances, self.num_links, 3), + dtype=torch.float32, + device=self.device, ) + self._com_pose = torch.zeros( + (self.num_instances, self.num_links, 7), + dtype=torch.float32, + device=self.device, + ) + self._default_mass: torch.Tensor | None = None + self._default_inertia: torch.Tensor | None = None + self._default_com_pose: torch.Tensor | None = None + + max_dof = self.dof + if ( + articulation_view is None + and self.device.type == "cuda" + and not self.is_newton_backend + ): + max_dof = self.ps.gpu_get_articulation_max_dof() self._target_qpos = torch.zeros( (self.num_instances, max_dof), dtype=torch.float32, device=self.device ) @@ -238,31 +299,23 @@ def __init__( 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. Returns: - torch.Tensor: The root pose of the articulation with shape of (num_instances, 7). + torch.Tensor: Root poses with shape ``(num_instances, 7)`` in + ``(x, y, z, qx, qy, qz, qw)`` order. """ - 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: @@ -271,22 +324,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: @@ -295,22 +333,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: @@ -328,22 +351,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: @@ -352,22 +360,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: @@ -376,20 +369,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: @@ -397,22 +377,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: @@ -421,20 +386,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: @@ -443,56 +395,17 @@ 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: """Get the pose of all links in the articulation. Returns: - torch.Tensor: The poses of the links in the articulation with shape (N, num_links, 7). + torch.Tensor: Link poses with shape ``(N, num_links, 7)`` in + ``(x, y, z, qx, qy, qz, qw)`` order. """ - 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: @@ -501,26 +414,169 @@ 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.articulation_view.fetch_link_velocity( + self._body_link_vel, + self._body_link_lin_vel, + self._body_link_ang_vel, + ) + + def _entity_link_name(self, entity: object, link_name: str) -> str: + """Resolve one public link name to an entity-local backend name.""" + resolver = getattr(self.articulation_view, "entity_link_name", None) + if resolver is not None: + return resolver(entity, link_name) + return link_name + + def _entity_drive_properties(self, entity: object) -> tuple[object, ...]: + """Read drive values without conflating backend target semantics.""" + if ( + isinstance(self.articulation_view, SpawnArticulationView) + and self.is_newton_backend + ): + return tuple(entity.get_newton_drive()) + return tuple(entity.get_drive()) + + def _entity_link_properties(self, entity: object, link_name: str) -> object: + """Read native mass properties through the active backend contract.""" + if ( + isinstance(self.articulation_view, SpawnArticulationView) + and self.is_newton_backend + ): + return entity.get_newton_link_properties(link_name) + return entity.get_physical_attr(link_name) + + def read_physical_properties( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Refresh current mass, inertia diagonal, and local COM pose buffers. + + COM poses use the EmbodiChain convention ``xyz + xyzw`` and all + tensors use the public link ordering. DexSim physical-property + descriptors use ``wxyz`` and are converted at this boundary. + """ + masses: list[list[float]] = [] + inertias: list[list[np.ndarray]] = [] + com_poses: list[list[np.ndarray]] = [] + for entity in self.entities: + mass_row: list[float] = [] + inertia_row: list[np.ndarray] = [] + com_row: list[np.ndarray] = [] + for link_name in self.link_names: + local_name = self._entity_link_name(entity, link_name) + attr = self._entity_link_properties(entity, local_name) + mass_row.append(float(attr.mass)) + inertia_row.append(np.asarray(attr.inertia, dtype=np.float32)) + com_row.append( + np.concatenate( + ( + np.asarray(attr.com_position, dtype=np.float32), + convert_quat( + np.asarray(attr.com_quaternion, dtype=np.float32), + to="xyzw", + ), + ) + ) ) - 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, + masses.append(mass_row) + inertias.append(inertia_row) + com_poses.append(com_row) + + self._mass.copy_( + torch.as_tensor( + np.asarray(masses, dtype=np.float32), + dtype=torch.float32, + device=self.device, + ) + ) + self._inertia.copy_( + torch.as_tensor( + np.asarray(inertias, dtype=np.float32), + dtype=torch.float32, + device=self.device, ) - self.ps.gpu_fetch_link_data( - data=self._body_link_ang_vel, - gpu_indices=self.gpu_indices, - data_type=ArticulationGPUAPIReadType.LINK_ANGULAR_VELOCITY, + ) + self._com_pose.copy_( + torch.as_tensor( + np.asarray(com_poses, dtype=np.float32), + dtype=torch.float32, + device=self.device, ) - 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._mass, self._inertia, self._com_pose + + @property + def mass(self) -> torch.Tensor: + """Current link masses with shape ``(N, num_links)``.""" + return self.read_physical_properties()[0] + + @property + def inertia(self) -> torch.Tensor: + """Current link inertia diagonals with shape ``(N, num_links, 3)``.""" + return self.read_physical_properties()[1] + + @property + def com_pose(self) -> torch.Tensor: + """Current local link COM poses as ``xyz + xyzw`` tensors.""" + return self.read_physical_properties()[2] + + @property + def default_physical_properties_initialized(self) -> bool: + """Whether initialization-time link mass properties are available.""" + return ( + self._default_mass is not None + and self._default_inertia is not None + and self._default_com_pose is not None + ) + + @property + def default_mass(self) -> torch.Tensor: + """Initialization-time link masses with shape ``(N, num_links)``.""" + if self._default_mass is None: + raise RuntimeError("Default articulation link masses are unavailable.") + return self._default_mass + + @property + def default_inertia(self) -> torch.Tensor: + """Initialization-time link inertia diagonals.""" + if self._default_inertia is None: + raise RuntimeError("Default articulation link inertias are unavailable.") + return self._default_inertia + + @property + def default_com_pose(self) -> torch.Tensor: + """Initialization-time local link COM poses in ``xyz + xyzw`` order.""" + if self._default_com_pose is None: + raise RuntimeError("Default articulation link COM poses are unavailable.") + return self._default_com_pose + + def capture_default_physical_properties( + self, + *, + mass: torch.Tensor, + inertia: torch.Tensor, + com_pose: torch.Tensor, + ) -> None: + """Capture backend-resolved link mass properties exactly once.""" + expected_shapes = { + "mass": (self.num_instances, self.num_links), + "inertia": (self.num_instances, self.num_links, 3), + "com_pose": (self.num_instances, self.num_links, 7), + } + values = {"mass": mass, "inertia": inertia, "com_pose": com_pose} + for name, value in values.items(): + if tuple(value.shape) != expected_shapes[name]: + raise ValueError( + f"Expected {name} shape {expected_shapes[name]}, " + f"got {tuple(value.shape)}." + ) + if self.default_physical_properties_initialized: + raise RuntimeError( + "Default articulation link mass properties are already captured." + ) + + self._default_mass = mass.to(self.device, dtype=torch.float32).clone() + self._default_inertia = inertia.to(self.device, dtype=torch.float32).clone() + self._default_com_pose = com_pose.to(self.device, dtype=torch.float32).clone() @property def joint_stiffness(self) -> torch.Tensor: @@ -530,7 +586,9 @@ def joint_stiffness(self) -> torch.Tensor: torch.Tensor: The joint stiffness of the articulation with shape (N, dof). """ return torch.as_tensor( - np.array([entity.get_drive()[0] for entity in self.entities]), + np.array( + [self._entity_drive_properties(entity)[0] for entity in self.entities] + ), dtype=torch.float32, device=self.device, ) @@ -543,7 +601,9 @@ def joint_damping(self) -> torch.Tensor: torch.Tensor: The joint damping of the articulation with shape (N, dof). """ return torch.as_tensor( - np.array([entity.get_drive()[1] for entity in self.entities]), + np.array( + [self._entity_drive_properties(entity)[1] for entity in self.entities] + ), dtype=torch.float32, device=self.device, ) @@ -556,7 +616,9 @@ def joint_friction(self) -> torch.Tensor: torch.Tensor: The joint friction of the articulation with shape (N, dof). """ return torch.as_tensor( - np.array([entity.get_drive()[4] for entity in self.entities]), + np.array( + [self._entity_drive_properties(entity)[4] for entity in self.entities] + ), dtype=torch.float32, device=self.device, ) @@ -569,7 +631,9 @@ def joint_armature(self) -> torch.Tensor: torch.Tensor: The joint armature of the articulation with shape (N, dof). """ return torch.as_tensor( - np.array([entity.get_drive()[5] for entity in self.entities]), + np.array( + [self._entity_drive_properties(entity)[5] for entity in self.entities] + ), dtype=torch.float32, device=self.device, ) @@ -640,12 +704,47 @@ class Articulation(BatchEntity): def __init__( self, cfg: ArticulationCfg, - entities: List[_Articulation] = None, + entities: Sequence[_Articulation | SpawnedArticulation] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: - # Initialize world and physics scene - self._world = dexsim.default_world() - self._ps = self._world.get_physics_scene() + self._newton_mimic_compliance_configured = False + self._prepared_default_root_topology_revision = -1 + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared Articulation requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._world = None + self._ps = None + self._data = None + self._all_indices = torch.arange(declared_num_instances, dtype=torch.int32) + self._visual_material = [{} for _ in range(declared_num_instances)] + self.is_shared_visual_material = False + self._has_collision_visible_node_dict = {} + return + + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result + if spawn_result is None: + # Legacy initialization remains temporarily while SimulationManager + # migration is in progress. Spawn-bound facades never reach for a + # process-global World or raw PhysicsScene. + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = None self.cfg = cfg self._entities = entities @@ -654,96 +753,55 @@ def __init__( # Store all indices for batch operations self._all_indices = torch.arange(len(entities), dtype=torch.int32) - if device.type == "cuda": + if ( + spawn_result is None + and device.type == "cuda" + and not is_newton_scene(self._ps) + ): self._world.update(0.001) - self._data = ArticulationData(entities=entities, ps=self._ps, device=device) + articulation_view = None + if spawn_result is not None: + batch = spawn_result.create_articulation_batch(entities) + articulation_view = SpawnArticulationView(spawn_result, batch, device) + self._data = ArticulationData( + entities=entities, + ps=self._ps, + device=device, + articulation_view=articulation_view, + ) self.cfg: ArticulationCfg if self.cfg.init_qpos is None: self.cfg.init_qpos = torch.zeros(self.dof, dtype=torch.float32) - # Get default masses. - self.default_link_masses = self.get_mass() - - # Determine if we should use USD properties or cfg properties. - if not self.cfg.use_usd_properties: - # Set articulation configuration in DexSim - set_dexsim_articulation_cfg(entities, self.cfg) - - num_entities = len(entities) - dof = self._data.dof - default_cfg = JointDrivePropertiesCfg() - self.default_joint_damping = torch.full( - (num_entities, dof), - default_cfg.damping, - dtype=torch.float32, - device=device, - ) - self.default_joint_stiffness = torch.full( - (num_entities, dof), - default_cfg.stiffness, - dtype=torch.float32, - device=device, - ) - self.default_joint_max_effort = torch.full( - (num_entities, dof), - default_cfg.max_effort, - dtype=torch.float32, - device=device, - ) - self.default_joint_max_velocity = torch.full( - (num_entities, dof), - default_cfg.max_velocity, - dtype=torch.float32, - device=device, - ) - self.default_joint_friction = torch.full( - (num_entities, dof), - default_cfg.friction, - dtype=torch.float32, - device=device, - ) - self.default_joint_armature = torch.full( - (num_entities, dof), - default_cfg.armature, - dtype=torch.float32, - device=device, - ) + self._capture_default_physical_properties() + + preserve_asset_physics = self.cfg.resolve_asset_physics_mode() == "preserve" + self.default_joint_stiffness = self._data.joint_stiffness.clone() + self.default_joint_damping = self._data.joint_damping.clone() + self.default_joint_friction = self._data.joint_friction.clone() + self.default_joint_armature = self._data.joint_armature.clone() + self.default_joint_max_effort = self._data.qf_limits.clone() + self.default_joint_max_velocity = self._data.qvel_limits.clone() + + # Spawn descriptors already contain build-time overlays. The retained + # legacy path applies only explicitly requested drive fields here. + if ( + spawn_result is None + and not preserve_asset_physics + and self.cfg.joint_drive_props is not None + ): self._set_default_joint_drive() - else: - # Read current properties from USD-loaded entities - self.default_joint_stiffness = self._data.joint_stiffness.clone() - self.default_joint_damping = self._data.joint_damping.clone() - self.default_joint_friction = self._data.joint_friction.clone() - self.default_joint_armature = self._data.joint_armature.clone() - self.default_joint_max_effort = self._data.qf_limits.clone() - self.default_joint_max_velocity = self._data.qvel_limits.clone() - - # Write the USD properties back to cfg - usd_drive_pros = self.cfg.drive_pros - usd_drive_pros.stiffness = ( - self.default_joint_stiffness[0].cpu().numpy().tolist() - ) - usd_drive_pros.damping = ( - self.default_joint_damping[0].cpu().numpy().tolist() - ) - usd_drive_pros.friction = ( - self.default_joint_friction[0].cpu().numpy().tolist() - ) - usd_drive_pros.armature = ( - self.default_joint_armature[0].cpu().numpy().tolist() - ) - usd_drive_pros.max_effort = ( - self.default_joint_max_effort[0].cpu().numpy().tolist() - ) - usd_drive_pros.max_velocity = ( - self.default_joint_max_velocity[0].cpu().numpy().tolist() - ) - # Apply configured qpos limits if provided. This replaces the asset - # limits as the baseline and allows expanding the allowed range. - if self.cfg.qpos_limits is not None: + # Spawn-owned articulations compile both named and flattened-DOF limits + # into the source-resolved descriptor before either backend builds its + # model. The retained raw path must still apply limits at runtime. + if ( + spawn_result is None + and self.cfg.qpos_limits is not None + and not preserve_asset_physics + ): if isinstance(self.cfg.qpos_limits, dict): indices, _, values = resolve_matching_names_values( self.cfg.qpos_limits, self.joint_names @@ -766,11 +824,18 @@ def __init__( ) self.set_qpos_limits(qpos_limits) + is_usd_source = str(self.cfg.fpath).lower().endswith((".usd", ".usda", ".usdc")) self.pk_chain = None - if self.cfg.build_pk_chain: + if self.cfg.build_pk_chain and not is_usd_source: self.pk_chain = create_pk_chain( urdf_path=self.cfg.fpath, device=self.device ) + elif self.cfg.build_pk_chain: + logger.log_warning( + f"Articulation {self.uid!r} uses USD for simulation; skipping " + "the URDF-only pk_chain. Configure a solver with its matching " + "URDF when kinematics are required." + ) # For rendering purposes, each articulation can have multiple material instances associated with its links. self._visual_material: List[Dict[str, VisualMaterialInst]] = [ @@ -778,28 +843,206 @@ def __init__( ] self.is_shared_visual_material = False - # Stores mimic information for joints. - self._mimic_info = entities[0].get_mimic_info() + # Stores mimic information in the same index domain as qpos/qvel/qf. + self._mimic_info = self._state_mimic_info() 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 ( + spawn_result is None + and device.type == "cuda" + and not is_newton_scene(self._ps) + ): self._world.update(0.001) + # Spawn-bound articulations receive post-load configuration before + # their initial reset. Legacy construction keeps its historical reset. super().__init__(cfg, entities, device) + if spawn_result is None: + self.reset() self._initialize_existing_visual_material() # set default collision filter - self._set_default_collision_filter() + if spawn_result is None: + self._set_default_collision_filter() # flag for collision visible node existence self._has_collision_visible_node_dict = dict() for link_name in self.link_names: self._has_collision_visible_node_dict[link_name] = False + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized SpawnResult.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its SpawnResult binding.""" + return self._world is None + + @property + def num_instances(self) -> int: + if self._entities: + return len(self._entities) + return self._declared_num_instances + + def attach_spawn_handles( + self, + entities: Sequence[SpawnedArticulation], + ) -> None: + """Store handles and expose metadata without initializing Batch data. + + This pre-finalize step supports eager Default loading and only reads + articulation metadata. ``bind_spawn()`` performs result-dependent + Batch/Data initialization after finalization. + """ + handles = list(entities) + if len(handles) != self._declared_num_instances: + raise ValueError( + f"Articulation {self.uid!r} expected " + f"{self._declared_num_instances} Spawn handles, got {len(handles)}." + ) + self._entities = handles + self._mimic_info = self._state_mimic_info() + self.active_joint_ids = [ + index for index in range(self.dof) if index not in self.mimic_ids + ] + + def bind_spawn( + self, + result: SpawnResult, + ) -> None: + """Initialize this declared facade from Spawn articulation handles.""" + if self.is_spawn_bound: + raise RuntimeError(f"Articulation {self.uid!r} is already Spawn-bound.") + if not self.is_declared: + raise RuntimeError( + f"Articulation {self.uid!r} was not created as a Spawn declaration." + ) + + cfg = self.cfg + device = self.device + entities = list(self._entities) + if len(entities) != self._declared_num_instances: + raise ValueError( + f"Articulation {self.uid!r} expected " + f"{self._declared_num_instances} Spawn handles, got {len(entities)}." + ) + + # Build and configure the bound state off to the side. If batch + # creation or post-load configuration fails, the public facade remains + # declared and can be retried by SimulationManager.prepare(). + bound = type(self)( + cfg, + entities, + device, + spawn_result=result, + ) + bound._prepared_default_root_topology_revision = getattr( + self, + "_prepared_default_root_topology_revision", + -1, + ) + bound._apply_spawn_config() + if is_newton_gradient_mode(result): + initial_qpos = torch.as_tensor(bound.cfg.init_qpos).reshape(-1) + if initial_qpos.numel() != bound.dof: + raise ValueError( + f"Articulation {bound.uid!r} expected {bound.dof} initial " + f"joint positions, got {initial_qpos.numel()}." + ) + if torch.any(initial_qpos != 0.0): + raise NotImplementedError( + "Newton gradient mode cannot apply non-zero init_qpos after " + "Spawn finalization. Author the initial coordinates in the " + "source asset or initialize them in a differentiable task " + "before opening a Warp tape." + ) + # Spawn already authored the root pose and zero joint/dynamics + # state during model construction. Its Batch mutation APIs are + # intentionally fenced once the model requires gradients. + else: + bound.reset() + self.__dict__.clear() + self.__dict__.update(bound.__dict__) + + def _apply_spawn_config(self) -> None: + """Apply configuration that requires finalized backend resources. + + Link physics and joint-drive regex selection is resolved by + EmbodiChain against the source descriptor before finalization. Default + articulation-root properties are normally handled by the pre-runtime + hook; calling it here keeps direct facade binding safe. Render + operations also require materialized native resources. + """ + spawn_result = getattr(self, "_spawn_result", None) + self._prepare_spawn_runtime_config(spawn_result) + + self._newton_mimic_compliance_configured = _configure_newton_mimic_compliance( + result=spawn_result, + entities=self._entities, + state_joint_names=self._state_joint_names(), + mimic_ids=self.mimic_ids, + mimic_parents=self.mimic_parents, + ) + + if not self.cfg.compute_uv: + return + + for entity in self._entities: + for link_name in self.link_names: + render_body = entity.get_render_body(link_name) + if render_body is not None: + render_body.set_projective_uv() + + def _prepare_spawn_runtime_config(self, result: SpawnResult | None) -> None: + """Apply Default root properties before Direct GPU initialization. + + PhysX snapshots articulation solver iteration counts when the Direct + GPU runtime is initialized. Applying these values only during facade + binding is too late because ``World.init_gpu_physics()`` has already + performed its warm-up steps. CPU simulation accepts the late write, + which otherwise makes identical hand mimic constraints substantially + softer on CUDA. + """ + if result is None or getattr(result, "backend", None) != "dexsim": + return + + topology_revision = int(result.topology_revision) + if self._prepared_default_root_topology_revision == topology_revision: + return + + root_props = getattr(self.cfg, "root_props", None) + default_root_values_configured = root_props is not None and ( + root_props.sleep_threshold is not None + or root_props.min_position_iters is not None + or root_props.min_velocity_iters is not None + ) + if default_root_values_configured: + for entity in self._entities: + # SpawnedArticulation deliberately fences these setters, while + # its Default-native binding exposes the articulation-root API. + native_articulation = getattr(entity, "_physics_binding", None) + if native_articulation is None: + raise RuntimeError( + "Default Spawn articulation has no native physics binding." + ) + _apply_default_articulation_root_properties( + native_articulation, + root_props, + ) + self._prepared_default_root_topology_revision = topology_revision + def __str__(self) -> str: + if self.is_declared: + parent_str = ( + f"{self.__class__}: declared {self.num_instances} Spawn " + f"articulations | uid: {self.uid} | device: {self.device}" + ) + return parent_str parent_str = super().__str__() return parent_str + f" | dof: {self.dof} | num_links: {self.num_links}" @@ -810,7 +1053,9 @@ def dof(self) -> int: Returns: int: The degree of freedom of the articulation. """ - return self._data.dof + if self._data is not None: + return self._data.dof + return self._entities[0].get_dof() @cached_property def active_dof(self) -> int: @@ -828,7 +1073,9 @@ def num_links(self) -> int: Returns: int: The number of links in the articulation. """ - return self._data.num_links + if self._data is not None: + return self._data.num_links + return len(self._entities[0].get_link_names()) @cached_property def link_names(self) -> List[str]: @@ -837,7 +1084,9 @@ def link_names(self) -> List[str]: Returns: List[str]: The names of the links in the articulation. """ - return self._data.link_names + if self._data is not None: + return self._data.link_names + return self._entities[0].get_link_names() @cached_property def user_ids(self) -> torch.Tensor: @@ -868,12 +1117,202 @@ def root_link_name(self) -> str: @cached_property def joint_names(self) -> List[str]: - """Get the names of the joints in the articulation. + """Get active joint names in public qpos-buffer order. Returns: - List[str]: The names of the actived joints in the articulation. + List[str]: Active joint names aligned with qpos, qvel, and qf. + """ + if getattr(self, "_data", None) is not None: + return list(self._data.articulation_view.joint_names) + return self._state_joint_names() + + def _state_joint_names(self) -> List[str]: + """Return active joint names in the backing qpos-buffer order. + + Spawn's Newton batch layout may differ from its source articulation + order. Joint IDs sent to the batch must therefore use the layout + order. :attr:`joint_names` exposes this same order; query the Spawn + handle directly only for source-topology resolution. + """ + if not self._entities: + return [] + entity = self._entities[0] + try: + layout = entity.joint_dof_layout + except (AttributeError, RuntimeError): + return entity.get_actived_joint_names() + return [joint.name for joint in layout] + + def _source_qpos_to_state_order(self, qpos: torch.Tensor) -> torch.Tensor: + """Map source-ordered initial qpos values to the runtime state order.""" + if not self.is_spawn_bound: + return qpos + + source_joint_names = self._entities[0].get_actived_joint_names() + state_joint_names = self._state_joint_names() + if source_joint_names == state_joint_names: + return qpos + + source_indices = {name: index for index, name in enumerate(source_joint_names)} + try: + state_order = [source_indices[name] for name in state_joint_names] + except KeyError as error: + raise RuntimeError( + "Spawn articulation state layout contains a joint absent from " + "the source articulation layout." + ) from error + return qpos[..., state_order] + + def _state_mimic_info(self) -> _MimicInfo: + """Map source-articulation mimic indices to state-buffer indices.""" + entity = self._entities[0] + source_info = entity.get_mimic_info() + source_mimic_ids = np.asarray(source_info.mimic_id, dtype=np.int32).reshape(-1) + source_parent_ids = np.asarray( + source_info.mimic_parent, dtype=np.int32 + ).reshape(-1) + multipliers = np.asarray( + source_info.mimic_multiplier, dtype=np.float32 + ).reshape(-1) + offsets = np.asarray(source_info.mimic_offset, dtype=np.float32).reshape(-1) + relation_count = len(source_mimic_ids) + if not all( + len(values) == relation_count + for values in (source_parent_ids, multipliers, offsets) + ): + raise RuntimeError("Articulation mimic metadata has inconsistent lengths.") + if relation_count == 0: + return _MimicInfo( + mimic_id=source_mimic_ids, + mimic_parent=source_parent_ids, + mimic_multiplier=multipliers, + mimic_offset=offsets, + ) + + source_joint_names = entity.get_actived_joint_names() + try: + state_joint_ids = { + joint.name: int(joint.dof_start) for joint in entity.joint_dof_layout + } + except (AttributeError, RuntimeError): + state_joint_ids = { + name: index for index, name in enumerate(source_joint_names) + } + + try: + mimic_ids = np.asarray( + [ + state_joint_ids[source_joint_names[int(source_id)]] + for source_id in source_mimic_ids + ], + dtype=np.int32, + ) + parent_ids = np.asarray( + [ + state_joint_ids[source_joint_names[int(source_id)]] + for source_id in source_parent_ids + ], + dtype=np.int32, + ) + except (IndexError, KeyError) as error: + raise RuntimeError( + "Articulation mimic metadata references a joint absent from " + "the backing state layout." + ) from error + + return _MimicInfo( + mimic_id=mimic_ids, + mimic_parent=parent_ids, + mimic_multiplier=multipliers, + mimic_offset=offsets, + ) + + def _project_mimic_qpos(self, qpos: torch.Tensor) -> torch.Tensor: + """Return qpos with every mimic child projected from its parent.""" + if not self.mimic_ids: + return qpos + + projected = qpos.clone() + mimic_ids = torch.as_tensor( + self.mimic_ids, dtype=torch.long, device=qpos.device + ) + parent_ids = torch.as_tensor( + self.mimic_parents, dtype=torch.long, device=qpos.device + ) + multipliers = torch.as_tensor( + self.mimic_multipliers, dtype=qpos.dtype, device=qpos.device + ) + offsets = torch.as_tensor( + self.mimic_offsets, dtype=qpos.dtype, device=qpos.device + ) + projected[..., mimic_ids] = projected[..., parent_ids] * multipliers + offsets + return projected + + def _stabilize_newton_mimic_target_write( + self, + values: torch.Tensor, + env_ids: torch.Tensor, + joint_ids: torch.Tensor, + *, + velocity: bool, + ) -> None: + """Update weak follower-drive targets for written mimic leaders. + + The native Newton equality remains the physical coupling. This only + keeps its low-gain follower stabilizer pointed at the same commanded + relation; it never copies measured qpos or qvel into follower state. """ - return self._entities[0].get_actived_joint_names() + if not self._newton_mimic_compliance_configured: + return + + selected_columns = { + int(joint_id): column + for column, joint_id in enumerate(joint_ids.detach().cpu().tolist()) + } + follower_ids: list[int] = [] + follower_targets: list[torch.Tensor] = [] + for child_id, parent_id, multiplier, offset in zip( + self.mimic_ids, + self.mimic_parents, + self.mimic_multipliers, + self.mimic_offsets, + strict=True, + ): + parent_column = selected_columns.get(int(parent_id)) + if parent_column is None: + continue + target = values[:, parent_column] * float(multiplier) + if not velocity: + target = target + float(offset) + follower_ids.append(int(child_id)) + follower_targets.append(target) + + if not follower_ids: + return + + targets = torch.stack(follower_targets, dim=1) + follower_ids_tensor = torch.as_tensor( + follower_ids, dtype=torch.int32, device=self.device + ) + if velocity: + limits = self.body_data.qvel_limits[env_ids][:, follower_ids_tensor] + targets = targets.clamp(-limits, limits) + self._data.articulation_view.apply_qvel( + targets, + env_ids, + follower_ids_tensor, + target=True, + ) + return + + limits = self.body_data.qpos_limits[env_ids][:, follower_ids_tensor, :] + targets = targets.clamp(limits[..., 0], limits[..., 1]) + self._data.articulation_view.apply_qpos( + targets, + env_ids, + follower_ids_tensor, + target=True, + ) @cached_property def active_joint_names(self) -> List[str]: @@ -882,7 +1321,8 @@ def active_joint_names(self) -> List[str]: Returns: List[str]: The names of the active joints in the articulation. """ - return [self.joint_names[i] for i in self.active_joint_ids] + state_joint_names = self._state_joint_names() + return [state_joint_names[i] for i in self.active_joint_ids] @cached_property def all_joint_names(self) -> List[str]: @@ -984,6 +1424,87 @@ def body_data(self) -> ArticulationData: """ return self._data + @property + def default_link_masses(self) -> torch.Tensor: + """Initialization-time link masses retained for compatibility.""" + return self.body_data.default_mass + + def _capture_default_physical_properties(self) -> None: + """Capture materialized link mass properties as reset defaults.""" + if self._data.default_physical_properties_initialized: + return + mass, inertia, com_pose = self._data.read_physical_properties() + self._data.capture_default_physical_properties( + mass=mass, + inertia=inertia, + com_pose=com_pose, + ) + + def _resolve_link_names( + self, link_names: str | Sequence[str] | None + ) -> tuple[list[str], torch.Tensor]: + """Validate link names and return their public data-column indices.""" + names = ( + list(self.link_names) + if link_names is None + else [link_names] if isinstance(link_names, str) else list(link_names) + ) + unknown = [name for name in names if name not in self.link_names] + if unknown: + raise ValueError( + f"Unknown articulation links {unknown}; available links: " + f"{self.link_names}." + ) + indices = torch.as_tensor( + [self.link_names.index(name) for name in names], + dtype=torch.long, + device=self.device, + ) + return names, indices + + def _restore_default_physical_properties( + self, env_ids: Sequence[int] | torch.Tensor + ) -> None: + """Restore initialization-time link mass properties for selected rows.""" + if not self._data.default_physical_properties_initialized or len(env_ids) == 0: + return + + env_index = self._resolve_env_ids(env_ids) + env_list = env_index.detach().cpu().tolist() + default_mass = self._data.default_mass[env_index] + default_inertia = self._data.default_inertia[env_index] + default_com_pose = self._data.default_com_pose[env_index] + current_mass, current_inertia, current_com_pose = ( + value[env_index] for value in self._data.read_physical_properties() + ) + + mass_changed = not torch.allclose(current_mass, default_mass) + inertia_changed = not torch.allclose(current_inertia, default_inertia) + if mass_changed: + self.set_mass(default_mass, link_names=self.link_names, env_ids=env_list) + if mass_changed or inertia_changed: + self.set_inertia( + default_inertia, + link_names=self.link_names, + env_ids=env_list, + ) + if not torch.allclose(current_com_pose, default_com_pose): + self.set_com_pose( + default_com_pose, + link_names=self.link_names, + env_ids=env_list, + ) + + 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. @@ -1114,53 +1635,29 @@ 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. 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 articulation with shape (N, 7) or (N, 4, 4) depending on `to_matrix`. @@ -1205,7 +1702,7 @@ def get_link_pose( Args: link_name (str): The name of the link. env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. - 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 pose of the specified link with shape (N, 7) or (N, 4, 4) depending on `to_matrix`. @@ -1459,34 +1956,18 @@ def set_qpos( :, local_joint_ids, : ] qpos = qpos.clamp(selected_limits[..., 0], selected_limits[..., 1]) - - if self.device.type == "cpu": - local_joint_ids_np = ( - local_joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) - ) - for i, env_idx in enumerate(local_env_ids.detach().cpu().tolist()): - setter = ( - self._entities[env_idx].set_target_qpos - if target - else self._entities[env_idx].set_current_qpos - ) - setter(qpos[i].detach().cpu().numpy(), local_joint_ids_np) - else: - data_type = ( - ArticulationGPUAPIWriteType.JOINT_TARGET_POSITION - if target - else ArticulationGPUAPIWriteType.JOINT_POSITION - ) - - # Always fetch the latest data to avoid stale values - qpos_set = self.body_data._target_qpos if target else self.body_data._qpos - - 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, + self._data.articulation_view.apply_qpos( + qpos, + local_env_ids, + local_joint_ids, + target=target, + ) + if target: + self._stabilize_newton_mimic_target_write( + qpos, + local_env_ids, + local_joint_ids, + velocity=False, ) def get_qvel(self, target: bool = False) -> torch.Tensor: @@ -1537,55 +2018,35 @@ def set_qvel( Raises: ValueError: If the length of `env_ids` does not match the length of `qvel`. """ - local_env_ids = self._all_indices if env_ids is None else env_ids + local_env_ids = self._resolve_env_ids(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)}." ) - 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 - - 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 = self._resolve_joint_ids(joint_ids) - 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, + ) + if target: + self._stabilize_newton_mimic_target_write( + qvel, + local_env_ids, + local_joint_ids, + velocity=True, ) def set_qf( @@ -1603,30 +2064,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 get_qf(self) -> torch.Tensor: """Get the current generalized efforts (qf) of the articulation. @@ -1659,76 +2121,186 @@ def get_qf_limits( def set_mass( self, mass: torch.Tensor, - link_names: Sequence[str], - env_ids: Sequence[int] | None = None, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, ) -> None: """Set the mass of specific links in the articulation. Args: - mass (torch.Tensor): The mass values to set with shape (N, len(link_names)). - link_names (Sequence[str]): The names of the links to set the mass for. - env_ids (Sequence[int] | None, optional): Environment indices to apply the mass change. If None, applies to all environments. Defaults to None. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - - if len(local_env_ids) != len(mass): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match mass length {len(mass)}." - ) - - for link_name in link_names: - if link_name not in self.link_names: - logger.log_error( - f"Link name {link_name} not found in {self.__class__.__name__}. Available links: {self.link_names}" - ) - - 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()) + mass: Mass values with shape ``(num_envs, num_links)``. + link_names: Link names to update. If None, all links are updated. + env_ids: Environment indices. If None, all rows are updated. + """ + env_index = self._resolve_env_ids(env_ids) + env_list = env_index.detach().cpu().tolist() + names, _ = self._resolve_link_names(link_names) + mass = torch.as_tensor(mass, dtype=torch.float32, device=self.device) + expected_shape = (len(env_list), len(names)) + if tuple(mass.shape) != expected_shape: + raise ValueError( + f"Expected mass shape {expected_shape}, got {tuple(mass.shape)}." + ) + + for i, env_idx in enumerate(env_list): + entity = self._entities[env_idx] + for j, name in enumerate(names): + if self.is_spawn_bound: + local_name = self._entity_link_name(env_idx, name) + entity.set_link_mass(local_name, mass[i, j].item()) + elif self._data.is_newton_backend: + local_name = self._entity_link_name(env_idx, name) + entity.set_link_mass(local_name, mass[i, j].item()) + else: + entity.set_mass(name, mass[i, j].item()) def get_mass( self, - link_names: Sequence[str] | None = None, - env_ids: Sequence[int] | None = None, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, ) -> torch.Tensor: """Get the mass of specific links in the articulation. Args: - link_names (Sequence[str] | None, optional): The names of the links to get the mass for. If None, gets mass for all links. Defaults to None. - env_ids (Sequence[int] | None, optional): Environment indices to get the mass from. If None, gets from all environments. Defaults to None. + link_names: Link names to query. If None, all links are returned. + env_ids: Environment indices. If None, all rows are returned. Returns: - torch.Tensor: The mass of the specified links with shape (N, len(link_names)). + Selected link masses with shape ``(num_envs, num_links)``. """ - local_env_ids = self._all_indices if env_ids is None else env_ids + env_index = self._resolve_env_ids(env_ids) + _, link_index = self._resolve_link_names(link_names) + return self.body_data.mass[ + env_index[:, None], + link_index[None, :], + ] - if link_names is None: - link_names = self.link_names - else: - for link_name in link_names: - if link_name not in self.link_names: - logger.log_error( - f"Link name {link_name} not found in {self.__class__.__name__}. Available links: {self.link_names}" + def set_inertia( + self, + inertia: torch.Tensor, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set principal moments of inertia for selected links.""" + env_index = self._resolve_env_ids(env_ids) + env_list = env_index.detach().cpu().tolist() + names, _ = self._resolve_link_names(link_names) + inertia = torch.as_tensor(inertia, dtype=torch.float32, device=self.device) + expected_shape = (len(env_list), len(names), 3) + if tuple(inertia.shape) != expected_shape: + raise ValueError( + f"Expected inertia shape {expected_shape}, " + f"got {tuple(inertia.shape)}." + ) + + values = inertia.detach().cpu().numpy() + for i, env_idx in enumerate(env_list): + entity = self._entities[env_idx] + for j, name in enumerate(names): + local_name = self._entity_link_name(env_idx, name) + value = np.asarray(values[i, j], dtype=np.float32) + if self.is_spawn_bound and self._data.is_newton_backend: + entity.set_newton_link_properties( + local_name, + rigid_body=dexsim.spawn.RigidBodyPhysicsDesc.dynamic( + inertia=value + ), + ) + elif not self._data.is_newton_backend: + entity.get_physical_body(local_name).set_mass_space_inertia_tensor( + value + ) + else: + attr = entity.get_physical_attr(local_name) + attr.inertia = value + entity.set_physical_attr( + attr, + local_name, + is_replace_inertial=False, ) - mass_tensor = torch.zeros( - (len(local_env_ids), len(link_names)), - dtype=torch.float32, - device=self.device, - ) - 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() + def get_inertia( + self, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Get principal moments of inertia for selected links.""" + env_index = self._resolve_env_ids(env_ids) + _, link_index = self._resolve_link_names(link_names) + return self.body_data.inertia[ + env_index[:, None], + link_index[None, :], + ] + + def set_com_pose( + self, + com_pose: torch.Tensor, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set local COM poses in EmbodiChain ``xyz + xyzw`` convention.""" + env_index = self._resolve_env_ids(env_ids) + env_list = env_index.detach().cpu().tolist() + names, _ = self._resolve_link_names(link_names) + com_pose = torch.as_tensor(com_pose, dtype=torch.float32, device=self.device) + expected_shape = (len(env_list), len(names), 7) + if tuple(com_pose.shape) != expected_shape: + raise ValueError( + f"Expected COM pose shape {expected_shape}, " + f"got {tuple(com_pose.shape)}." + ) + + values = com_pose.detach().cpu().numpy() + for i, env_idx in enumerate(env_list): + entity = self._entities[env_idx] + for j, name in enumerate(names): + local_name = self._entity_link_name(env_idx, name) + position = np.asarray(values[i, j, :3], dtype=np.float32) + quaternion = np.asarray( + convert_quat(values[i, j, 3:7], to="wxyz"), + dtype=np.float32, ) - return mass_tensor + if self.is_spawn_bound and self._data.is_newton_backend: + entity.set_newton_link_properties( + local_name, + rigid_body=dexsim.spawn.RigidBodyPhysicsDesc.dynamic( + com_position=position, + com_quaternion=quaternion, + ), + ) + elif not self._data.is_newton_backend: + entity.get_physical_body(local_name).set_cmass_local_pose( + position, + quaternion, + ) + else: + attr = entity.get_physical_attr(local_name) + attr.com_position = position + attr.com_quaternion = quaternion + entity.set_physical_attr( + attr, + local_name, + is_replace_inertial=False, + ) + + def get_com_pose( + self, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Get local COM poses in EmbodiChain ``xyz + xyzw`` convention.""" + env_index = self._resolve_env_ids(env_ids) + _, link_index = self._resolve_link_names(link_names) + return self.body_data.com_pose[ + env_index[:, None], + link_index[None, :], + ] def get_link_physical_attr( self, link_names: str | Sequence[str] | None = None, env_ids: Sequence[int] | None = None, ) -> list[PhysicalAttr]: - """Get physical attributes for articulation links. + """Get DexSim-native physical attributes for articulation links. Args: link_names: Link names or regex patterns. If None, all links are returned. @@ -1738,6 +2310,11 @@ def get_link_physical_attr( List of :class:`~dexsim.types.PhysicalAttr`, one per (env, link) pair in row-major order (env-major). """ + if self._data is not None and self._data.is_newton_backend: + raise RuntimeError( + "get_link_physical_attr() exposes DexSim PhysicalAttr semantics; " + "use get_newton_link_properties() for Newton." + ) if link_names is None: matched_link_names = self.link_names elif isinstance(link_names, str): @@ -1752,28 +2329,88 @@ def get_link_physical_attr( local_env_ids = [0] if env_ids is None else list(env_ids) attrs: list[PhysicalAttr] = [] for env_idx in local_env_ids: + entity = self._entities[env_idx] for name in matched_link_names: - attrs.append(self._entities[env_idx].get_physical_attr(name)) + attrs.append( + entity.get_physical_attr(self._entity_link_name(env_idx, name)) + ) return attrs + def get_newton_link_properties( + self, + link_names: str | Sequence[str] | None = None, + env_ids: Sequence[int] | None = None, + ) -> list[dexsim.spawn.RigidBodyPhysicsDesc]: + """Get Newton model mass properties as typed Spawn descriptors. + + Args: + link_names: Link names or regex patterns. If None, all links are + returned. + env_ids: Environment indices. If None, only environment 0 is + queried. + + Returns: + One typed descriptor per selected ``(environment, link)`` pair in + environment-major order. + """ + if not ( + self.is_spawn_bound + and self._data is not None + and self._data.is_newton_backend + ): + raise RuntimeError( + "get_newton_link_properties() requires a Spawn-bound Newton " + "articulation." + ) + if link_names is None: + matched_link_names = self.link_names + else: + _, matched_link_names = resolve_matching_names( + keys=link_names, + list_of_strings=self.link_names, + ) + + local_env_ids = [0] if env_ids is None else list(env_ids) + properties = [] + for env_idx in local_env_ids: + entity = self._entities[env_idx] + for name in matched_link_names: + properties.append( + entity.get_newton_link_properties( + self._entity_link_name(env_idx, name) + ) + ) + return properties + def set_link_physical_attr( self, - attrs: RigidBodyAttributesCfg | RigidBodyAttributesOverrideCfg | PhysicalAttr, + attrs: RigidBodyPhysicsCfg | PhysicalAttr, link_names: str | Sequence[str] | None = None, env_ids: Sequence[int] | None = None, *, - base_attrs: RigidBodyAttributesCfg | None = None, + base_attrs: RigidBodyPhysicsCfg | None = None, replace_inertial: bool = False, ) -> None: """Set physical attributes for selected articulation links. Args: - attrs: Full, partial, or DexSim physical attributes to apply. + attrs: Grouped or DexSim physical attributes to apply. link_names: Link names or regex patterns. If None, all links are updated. 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:: + This compatibility API exposes DexSim ``PhysicalAttr`` semantics. + Newton properties must use typed Spawn descriptors. """ + is_newton = self._data is not None and self._data.is_newton_backend + if is_newton: + raise TypeError( + "set_link_physical_attr() is DexSim-only; use typed Newton " + "link properties or set_mass()/set_inertia()/set_com_pose()." + ) + if link_names is None: matched_link_names = self.link_names elif isinstance(link_names, str): @@ -1785,22 +2422,27 @@ def set_link_physical_attr( keys=link_names, list_of_strings=self.link_names ) - if isinstance(attrs, RigidBodyAttributesOverrideCfg): + if isinstance(attrs, RigidBodyPhysicsCfg): if base_attrs is None: base_attrs = self.cfg.attrs - physical_attr = attrs.merge_with(base_attrs) - if attrs.mass is not None: - replace_inertial = True - elif isinstance(attrs, RigidBodyAttributesCfg): - physical_attr = attrs.attr() + physical_attr = attrs.to_dexsim_physical_attr( + base=base_attrs.to_dexsim_physical_attr() + ) + mass_props = attrs.mass_props + if mass_props is not None and mass_props.recompute_inertia is not None: + replace_inertial = bool(mass_props.recompute_inertia) else: physical_attr = attrs local_env_ids = self._all_indices if env_ids is None else env_ids for env_idx in local_env_ids: + entity = self._entities[env_idx] for name in matched_link_names: - self._entities[env_idx].set_physical_attr( - physical_attr, name, is_replace_inertial=replace_inertial + local_name = self._entity_link_name(env_idx, name) + entity.set_physical_attr( + physical_attr, + local_name, + is_replace_inertial=replace_inertial, ) def set_joint_drive( @@ -1811,9 +2453,11 @@ def set_joint_drive( max_velocity: torch.Tensor | None = None, friction: torch.Tensor | None = None, armature: torch.Tensor | None = None, - drive_type: str = "none", + drive_type: str | None = None, joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None, + *, + target_mode: str | int | None = None, ) -> None: """Set the drive properties for the articulation. @@ -1824,24 +2468,88 @@ def set_joint_drive( max_velocity (torch.Tensor): The maximum velocity of the joint drive with shape (len(env_ids), len(joint_ids)). friction (torch.Tensor): The joint friction coefficient with shape (len(env_ids), len(joint_ids)). armature (torch.Tensor): The joint armature with shape (len(env_ids), len(joint_ids)). - drive_type (str, optional): The type of drive to apply. Defaults to "none". + drive_type: ``force``, ``acceleration``, or ``none``. ``None`` + preserves the current mode unless a target mode activates a + force drive. joint_ids (Sequence[int] | None, optional): The joint indices to apply the drive to. If None, applies to all joints. Defaults to None. env_ids (Sequence[int] | None, optional): The environment indices to apply the drive to. If None, applies to all environments. Defaults to None. + target_mode: Portable target mode: ``none``, ``position``, + ``velocity``, ``position_velocity``, ``effort``, or integer + value 0 through 4. """ local_env_ids = self._all_indices if env_ids is None else env_ids local_joint_ids = np.arange(self.dof) if joint_ids is None else joint_ids cache_env_ids = self._resolve_env_ids(env_ids) cache_joint_ids = self._resolve_joint_ids(joint_ids) + mode_cfg = JointDrivePropertiesCfg( + target_mode=target_mode, + drive_type=drive_type, + ) + resolved_target_mode, resolved_drive_type = mode_cfg._resolve_modes() + if isinstance(resolved_target_mode, dict): + raise TypeError( + "set_joint_drive() accepts one scalar target_mode; configure " + "per-joint mappings through JointDrivePropertiesCfg." + ) + target_mode_value = ( + None + if resolved_target_mode is None + else _normalize_joint_target_mode(resolved_target_mode) + ) + if target_mode_value in {1, 2, 3} and resolved_drive_type == "none": + raise ValueError( + "drive_type='none' conflicts with an active target_mode; use " + "target_mode='none' or 'effort'." + ) + def _drive_arg(value: torch.Tensor, index: int) -> float | np.ndarray: result = value[index].detach().cpu().numpy() return result.item() if result.size == 1 else result for i, env_idx in enumerate(local_env_ids): - drive_args = { - "drive_type": get_dexsim_drive_type(drive_type), - "joint_ids": local_joint_ids, - } + if self.is_spawn_bound and self.body_data.is_newton_backend: + if resolved_drive_type == "acceleration" and target_mode_value in { + 1, + 2, + 3, + }: + raise NotImplementedError( + "Newton Spawn does not have an exact equivalent of " + "the Default acceleration drive. Use " + "drive_type='force' or disable the drive." + ) + drive_args = {"joint_ids": local_joint_ids} + if target_mode_value is not None: + drive_args["target_mode"] = target_mode_value + if stiffness is not None: + drive_args["target_ke"] = _drive_arg(stiffness, i) + if damping is not None: + drive_args["target_kd"] = _drive_arg(damping, i) + if max_effort is not None: + drive_args["effort_limit"] = _drive_arg(max_effort, i) + if max_velocity is not None: + drive_args["velocity_limit"] = _drive_arg(max_velocity, i) + if friction is not None: + drive_args["friction"] = _drive_arg(friction, i) + if armature is not None: + drive_args["armature"] = _drive_arg(armature, i) + if target_mode_value in {0, 4}: + drive_args["target_ke"] = 0.0 + drive_args["target_kd"] = 0.0 + elif target_mode_value == 2: + drive_args["target_ke"] = 0.0 + self._entities[env_idx].set_newton_drive(**drive_args) + continue + + drive_args = {"joint_ids": local_joint_ids} + default_drive_type = resolved_drive_type + if target_mode_value in {0, 4}: + default_drive_type = "none" + elif target_mode_value in {1, 2, 3} and default_drive_type is None: + default_drive_type = "force" + if default_drive_type is not None: + drive_args["drive_type"] = get_dexsim_drive_type(default_drive_type) if stiffness is not None: drive_args["stiffness"] = _drive_arg(stiffness, i) if damping is not None: @@ -1854,6 +2562,11 @@ def _drive_arg(value: torch.Tensor, index: int) -> float | np.ndarray: drive_args["joint_friction"] = _drive_arg(friction, i) if armature is not None: drive_args["armature"] = _drive_arg(armature, i) + if target_mode_value in {0, 4}: + drive_args["stiffness"] = 0.0 + drive_args["damping"] = 0.0 + elif target_mode_value == 2: + drive_args["stiffness"] = 0.0 self._entities[env_idx].set_drive(**drive_args) if max_velocity is not None: @@ -1946,7 +2659,7 @@ def get_joint_drive( friction_i, armature_i, *_, - ) = self._entities[env_idx].get_drive() + ) = self._data._entity_drive_properties(self._entities[env_idx]) stiffness[i] = torch.as_tensor( stiffness_i, dtype=torch.float32, device=self.device )[local_joint_ids_tensor] @@ -1972,15 +2685,20 @@ def get_joint_drive_type( joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None, ) -> list[list[DriveType]]: - """Get the backend drive type for the selected joints. + """Get the portable drive type for the selected joints. Args: joint_ids: Joint indices to query. If None, queries all joints. env_ids: Environment indices to query. If None, queries all environments. Returns: - Backend drive types grouped by environment, with one + Drive types grouped by environment, with one :class:`~dexsim.types.DriveType` per selected joint. + + Newton has no acceleration-drive equivalent. Its passive and + direct-effort target modes map to :attr:`DriveType.NONE` because + neither installs a PD drive; position and velocity target modes + map to :attr:`DriveType.FORCE`. """ local_env_ids = self._all_indices if env_ids is None else env_ids if joint_ids is None: @@ -1994,12 +2712,63 @@ def get_joint_drive_type( drive_types: list[list[DriveType]] = [] for env_idx in local_env_ids: - entity_drive_types = self._entities[int(env_idx)].get_drive( - local_joint_ids - )[-1] - drive_types.append(list(entity_drive_types)) + entity = self._entities[int(env_idx)] + if self._data is not None and self._data.is_newton_backend: + target_modes = np.asarray(entity.get_newton_drive()[-1])[ + local_joint_ids + ] + drive_types.append( + [ + (DriveType.NONE if int(mode) in {0, 4} else DriveType.FORCE) + for mode in target_modes + ] + ) + else: + entity_drive_types = np.asarray(entity.get_drive()[-1])[local_joint_ids] + drive_types.append(list(entity_drive_types)) return drive_types + def get_joint_target_mode( + self, + joint_ids: Sequence[int] | None = None, + env_ids: Sequence[int] | None = None, + ) -> list[list[int]]: + """Get Newton ``JointTargetMode`` integer values by environment. + + Args: + joint_ids: Flattened DOF indices. If None, all DOFs are queried. + env_ids: Environment indices. If None, all environments are + queried. + + Returns: + Integer target modes grouped by selected environment. + """ + if not ( + self.is_spawn_bound + and self._data is not None + and self._data.is_newton_backend + ): + raise RuntimeError( + "get_joint_target_mode() requires a Spawn-bound Newton " "articulation." + ) + local_env_ids = self._all_indices if env_ids is None else env_ids + if joint_ids is None: + local_joint_ids = np.arange(self.dof, dtype=np.int32) + elif isinstance(joint_ids, torch.Tensor): + local_joint_ids = ( + joint_ids.detach().cpu().numpy().astype(np.int32, copy=False) + ) + else: + local_joint_ids = np.asarray(joint_ids, dtype=np.int32) + + target_modes = [] + for env_idx in local_env_ids: + modes = self._entities[int(env_idx)].get_newton_drive()[-1] + target_modes.append( + [int(value) for value in np.asarray(modes)[local_joint_ids]] + ) + return target_modes + def get_user_ids( self, link_name: str | None = None, env_ids: Sequence[int] | None = None ) -> torch.Tensor: @@ -2032,10 +2801,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.""" @@ -2092,49 +2858,76 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None: self.cfg: ArticulationCfg self.restore_visual_material(env_ids=local_env_ids) + self._restore_default_physical_properties(local_env_ids) - pos = torch.as_tensor( - self.cfg.init_pos, dtype=torch.float32, device=self.device - ) - rot = ( - torch.as_tensor(self.cfg.init_rot, dtype=torch.float32, device=self.device) - * torch.pi - / 180.0 - ) - pos = pos.unsqueeze(0).repeat(num_instances, 1) - rot = rot.unsqueeze(0).repeat(num_instances, 1) - mat = matrix_from_euler(rot, "XYZ") - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze(0) - .repeat(num_instances, 1, 1) - ) - pose[:, :3, 3] = pos - pose[:, :3, :3] = mat + if self.cfg.init_local_pose is not None: + pose = ( + torch.as_tensor( + self.cfg.init_local_pose, + dtype=torch.float32, + device=self.device, + ) + .reshape(1, 4, 4) + .repeat(num_instances, 1, 1) + ) + else: + pos = torch.as_tensor( + self.cfg.init_pos, dtype=torch.float32, device=self.device + ) + rot = ( + torch.as_tensor( + self.cfg.init_rot, dtype=torch.float32, device=self.device + ) + * torch.pi + / 180.0 + ) + pos = pos.unsqueeze(0).repeat(num_instances, 1) + rot = rot.unsqueeze(0).repeat(num_instances, 1) + pose = ( + torch.eye(4, dtype=torch.float32, device=self.device) + .unsqueeze(0) + .repeat(num_instances, 1, 1) + ) + pose[:, :3, 3] = pos + pose[:, :3, :3] = matrix_from_euler(rot, "XYZ") self.set_local_pose(pose, env_ids=local_env_ids) qpos = torch.as_tensor( self.cfg.init_qpos, dtype=torch.float32, device=self.device ) qpos = qpos.unsqueeze(0).repeat(num_instances, 1) + qpos = self._source_qpos_to_state_order(qpos) + if ( + self.body_data.is_newton_backend + and not self._newton_mimic_compliance_configured + ): + # Native Newton mimic constraints can generate a large corrective + # impulse when initialized away from their equality manifold. + qpos = self._project_mimic_qpos(qpos) self.set_qpos(qpos, target=False, env_ids=local_env_ids) # Set drive target to hold position. self.set_qpos(qpos, target=True, env_ids=local_env_ids) 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: + def _set_default_joint_drive( + self, + joint_drive_props: JointDrivePropertiesCfg | dict | None = None, + ) -> None: """Set default joint drive parameters based on the configuration.""" import numbers from embodichain.utils.string import resolve_matching_names_values - drive_props = [ + if joint_drive_props is None: + joint_drive_props = self.cfg.joint_drive_props + if joint_drive_props is None: + return + + joint_property_targets = [ ("damping", self.default_joint_damping), ("stiffness", self.default_joint_stiffness), ("max_effort", self.default_joint_max_effort), @@ -2143,8 +2936,12 @@ def _set_default_joint_drive(self) -> None: ("armature", self.default_joint_armature), ] - for prop_name, default_array in drive_props: - value = getattr(self.cfg.drive_pros, prop_name, None) + for prop_name, default_array in joint_property_targets: + value = ( + joint_drive_props.get(prop_name) + if isinstance(joint_drive_props, dict) + else getattr(joint_drive_props, prop_name, None) + ) if value is None: continue if isinstance(value, numbers.Number): @@ -2160,11 +2957,19 @@ def _set_default_joint_drive(self) -> None: except Exception as e: logger.log_error(f"Failed to set {prop_name}: {e}") - drive_pros = self.cfg.drive_pros - if isinstance(drive_pros, dict): - drive_type = drive_pros.get("drive_type", "none") + if isinstance(joint_drive_props, dict): + drive_type = joint_drive_props.get("drive_type") + target_mode = joint_drive_props.get("target_mode") else: - drive_type = getattr(drive_pros, "drive_type", "none") + drive_type = getattr(joint_drive_props, "drive_type", None) + target_mode = getattr(joint_drive_props, "target_mode", None) + if isinstance(target_mode, dict): + logger.log_warning( + "Per-joint target_mode mappings require a Spawn-bound " + "articulation; the retained raw-articulation path preserves " + "its current target modes." + ) + target_mode = None # Apply drive parameters to all articulations in the batch self.set_joint_drive( @@ -2175,6 +2980,7 @@ def _set_default_joint_drive(self) -> None: friction=self.default_joint_friction, armature=self.default_joint_armature, drive_type=drive_type, + target_mode=target_mode, ) def compute_fk( @@ -2384,7 +3190,12 @@ def set_visual_material( for link_name in link_names: mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{link_name}") for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_material(link_name, mat_inst.mat) + if self.is_spawn_bound: + self._entities[env_idx].set_material_inst( + link_name, mat_inst.mat + ) + else: + self._entities[env_idx].set_material(link_name, mat_inst.mat) self._visual_material[env_idx][link_name] = mat_inst if update_default: self._original_visual_material[env_idx][link_name] = ( @@ -2402,7 +3213,12 @@ def set_visual_material( mat_inst = mat.create_instance( f"{mat.uid}_{self.uid}_{link_name}_{env_idx}" ) - self._entities[env_idx].set_material(link_name, mat_inst.mat) + if self.is_spawn_bound: + self._entities[env_idx].set_material_inst( + link_name, mat_inst.mat + ) + else: + self._entities[env_idx].set_material(link_name, mat_inst.mat) self._visual_material[env_idx][link_name] = mat_inst if update_default: self._original_visual_material[env_idx][link_name] = ( @@ -2627,6 +3443,17 @@ def set_physical_visible( ) link_names = self.link_names if link_names is None else link_names + if self.is_spawn_bound: + for env_idx in self._all_indices: + entity = self._entities[env_idx] + for link_name in link_names: + self._spawn_result.set_physical_visible( + (entity, link_name), rgba, visible + ) + for link_name in link_names: + self._has_collision_visible_node_dict[link_name] = True + return + # create collision visible node if not exist if visible: for i, env_idx in enumerate(self._all_indices): @@ -2677,12 +3504,18 @@ def set_self_collision( ) def destroy(self) -> None: + if self.is_declared or self.is_spawn_bound: + # SpawnResult is the sole owner of native lifetime. + return env = self._world.get_env() arenas = env.get_all_arenas() if len(arenas) == 0: 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) __all__ = ["ArticulationData", "Articulation", "ArticulationJointKinematics"] diff --git a/embodichain/lab/sim/objects/backends/__init__.py b/embodichain/lab/sim/objects/backends/__init__.py new file mode 100644 index 000000000..3d039017d --- /dev/null +++ b/embodichain/lab/sim/objects/backends/__init__.py @@ -0,0 +1,40 @@ +# ---------------------------------------------------------------------------- +# 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, +) +from .spawn import SpawnArticulationView, SpawnRigidBodyView + +__all__ = [ + "ArticulationViewBase", + "RigidBodyViewBase", + "DefaultArticulationView", + "DefaultRigidBodyView", + "NewtonArticulationView", + "NewtonRigidBodyView", + "apply_collision_filter_for_entities", + "apply_collision_filter_for_envs", + "is_newton_scene", + "SpawnArticulationView", + "SpawnRigidBodyView", +] diff --git a/embodichain/lab/sim/objects/backends/base.py b/embodichain/lab/sim/objects/backends/base.py new file mode 100644 index 000000000..dfb480776 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/base.py @@ -0,0 +1,386 @@ +# ---------------------------------------------------------------------------- +# 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 COM-local poses as ``(x, y, z, qx, qy, qz, qw)``.""" + ... + + @abstractmethod + def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply COM-local poses from ``(x, y, z, qx, qy, qz, qw)``.""" + ... + + # -- 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.""" + ... + + def fetch_damping( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch linear/angular damping into ``data`` as ``(N, 2)``.""" + raise NotImplementedError("This backend view does not expose damping.") + + def apply_damping(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + """Apply linear/angular damping from an ``(N, 2)`` tensor.""" + raise NotImplementedError("This backend view does not expose damping.") + + def fetch_collision_filter( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + """Fetch collision-filter rows into ``data`` as ``(N, 4)``.""" + raise NotImplementedError( + "This backend view does not expose collision filters." + ) + + def apply_collision_filter( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + """Apply collision-filter rows from an ``(N, 4)`` tensor.""" + raise NotImplementedError( + "This backend view does not expose collision filters." + ) + + +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 EmbodiChain ``xyz + xyzw`` tensors.""" + ... + + @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 000000000..0af741ff8 --- /dev/null +++ b/embodichain/lab/sim/objects/backends/default.py @@ -0,0 +1,736 @@ +# ---------------------------------------------------------------------------- +# 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-backend rigid body data adapter over DexSim entities. + + Encapsulates both GPU (DexSim) 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(pose_cpu[:, 3:7]) + 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 at build time with CollisionPropertiesCfg 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-backend articulation data adapter over DexSim entities.""" + + 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, + ) + 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, + ) + return torch.cat((data[..., 4:], data[..., :4]), 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 = pose[:, 3:7] + 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 000000000..85bb833be --- /dev/null +++ b/embodichain/lab/sim/objects/backends/newton.py @@ -0,0 +1,1078 @@ +# ---------------------------------------------------------------------------- +# 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 TYPE_CHECKING, Any, Sequence +import numpy as np +import torch + +from dexsim.models import MeshObject +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 + +if TYPE_CHECKING: + from dexsim.engine.newton_physics.newton_physics_scene import NewtonPhysicsScene + from dexsim.spawn import SpawnResult +else: + NewtonPhysicsScene = Any + +__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") + ) + + +_DEFAULT_MIMIC_NATURAL_FREQUENCY = 1.0e3 +_DEFAULT_MIMIC_DAMPING_RATIO = 1.0e1 +_MIMIC_FOLLOWER_TARGET_GAIN_RATIO = 1.0e-2 + + +def _default_mujoco_mimic_solref(physics_dt: float, num_substeps: int) -> np.ndarray: + """Approximate Default's mimic compliance with MuJoCo ``solref``. + + Positive MuJoCo ``solref`` uses ``(timeconst, dampratio)`` and therefore + retains the effective-mass scaling of PhysX articulation mimic joints. + MuJoCo's reference-safety rule clamps ``timeconst`` to twice the solver + timestep, so apply the same bound explicitly. + """ + if not np.isfinite(physics_dt) or physics_dt <= 0.0: + raise ValueError("Newton physics_dt must be finite and positive.") + if num_substeps <= 0: + raise ValueError("Newton num_substeps must be positive.") + + solver_dt = physics_dt / num_substeps + natural_time_constant = 1.0 / ( + _DEFAULT_MIMIC_NATURAL_FREQUENCY * _DEFAULT_MIMIC_DAMPING_RATIO + ) + return np.asarray( + ( + max(natural_time_constant, 2.0 * solver_dt), + _DEFAULT_MIMIC_DAMPING_RATIO, + ), + dtype=np.float32, + ) + + +def _configure_newton_mimic_compliance( + *, + result: SpawnResult | None, + entities: Sequence[object], + state_joint_names: Sequence[str], + mimic_ids: Sequence[int], + mimic_parents: Sequence[int], +) -> bool: + """Tune native MuJoCo-Warp mimic constraints toward Default behavior. + + MuJoCo's default equality ``solref`` is underdamped relative to Default's + articulation mimic. Map Default's natural-frequency and damping-ratio + parameters to MuJoCo's mass-scaled positive convention as a stable + approximation. A follower drive with one percent of its leader's gains also + tracks the leader's *target* relation between solver updates. Keeping the + native equality rows enabled preserves mechanical force coupling; the drive + is only a stabilizer and never mirrors measured follower state. + """ + if result is None or result.backend != "newton" or not mimic_ids: + return False + + from dexsim.engine.newton_physics.backend_registry import get_newton_backend + + backend = get_newton_backend(result.world) + if ( + backend is None + or backend.solver_type != "mujoco_warp" + or backend.cfg.requires_grad + or (backend.model is not None and backend.model.requires_grad) + ): + return False + + relation_names = [ + (state_joint_names[child_id], state_joint_names[parent_id]) + for child_id, parent_id in zip(mimic_ids, mimic_parents, strict=True) + ] + first_binding = getattr(entities[0], "_physics_binding", None) + runtime = getattr(first_binding, "_runtime", None) + if runtime is None: + raise RuntimeError("Newton Spawn articulation has no finalized runtime.") + + model = runtime.model + target_ke = np.asarray(model.joint_target_ke.numpy()).reshape(-1) + target_kd = np.asarray(model.joint_target_kd.numpy()).reshape(-1) + target_mode = np.asarray(model.joint_target_mode.numpy()).reshape(-1) + expected_pairs: set[tuple[int, int]] = set() + for entity in entities: + binding = getattr(entity, "_physics_binding", None) + if binding is None or getattr(binding, "_runtime", None) is not runtime: + raise RuntimeError( + "Newton mimic configuration requires one shared finalized runtime." + ) + runtime_joints = {joint.name: joint for joint in binding.joints} + follower_ke: list[float] = [] + follower_kd: list[float] = [] + follower_mode: list[int] = [] + for child_name, parent_name in relation_names: + try: + child = runtime_joints[child_name] + parent = runtime_joints[parent_name] + except KeyError as error: + raise RuntimeError( + "Newton mimic metadata references a missing runtime joint." + ) from error + if int(child.qd_size) != 1 or int(parent.qd_size) != 1: + raise NotImplementedError( + "MuJoCo-Warp mimic compliance requires scalar joints." + ) + expected_pairs.add((int(child.joint_id), int(parent.joint_id))) + parent_dof = int(parent.qd_start) + follower_ke.append( + float(target_ke[parent_dof]) * _MIMIC_FOLLOWER_TARGET_GAIN_RATIO + ) + follower_kd.append( + float(target_kd[parent_dof]) * _MIMIC_FOLLOWER_TARGET_GAIN_RATIO + ) + follower_mode.append(int(target_mode[parent_dof])) + + configured = entity.set_newton_drive( + joint_ids=np.asarray(mimic_ids, dtype=np.int32), + target_ke=np.asarray(follower_ke, dtype=np.float32), + target_kd=np.asarray(follower_kd, dtype=np.float32), + target_mode=np.asarray(follower_mode, dtype=np.int32), + ) + if configured != len(mimic_ids): + raise RuntimeError( + "Newton failed to configure every mimic follower stabilizer." + ) + + mimic_joint0 = np.asarray(model.constraint_mimic_joint0.numpy()).reshape(-1) + mimic_joint1 = np.asarray(model.constraint_mimic_joint1.numpy()).reshape(-1) + row_by_pair = { + (int(child), int(parent)): row + for row, (child, parent) in enumerate( + zip(mimic_joint0, mimic_joint1, strict=True) + ) + } + try: + constraint_rows = np.asarray( + [row_by_pair[pair] for pair in expected_pairs], dtype=np.int32 + ) + except KeyError as error: + raise RuntimeError( + f"Newton model has no mimic constraint for joint pair {error.args[0]}." + ) from error + + solver = runtime.solver + mapping = getattr(solver, "mjc_eq_to_newton_mimic", None) + mjw_model = getattr(solver, "mjw_model", None) + if mapping is None or mjw_model is None: + raise RuntimeError("MuJoCo-Warp did not materialize Newton mimic rows.") + + mapping_values = np.asarray(mapping.numpy()) + selected = np.isin(mapping_values, constraint_rows) + if int(selected.sum()) != len(constraint_rows): + raise RuntimeError( + "MuJoCo-Warp mimic row mapping does not match the articulation." + ) + eq_solref = np.asarray(mjw_model.eq_solref.numpy()).copy() + mimic_solref = _default_mujoco_mimic_solref( + float(backend.cfg.dt), + int(backend.cfg.num_substeps), + ) + eq_solref[selected] = mimic_solref + mjw_model.eq_solref.assign(eq_solref) + + # Keep the optional CPU mirror coherent for debugging and CPU execution. + mj_model = getattr(solver, "mj_model", None) + if mj_model is not None and len(eq_solref) > 0: + mj_model.eq_solref[:] = eq_solref[0] + return True + + +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)`` and is **arena-local**: callers of + ``set_local_pose`` / ``get_local_pose`` pass and receive poses relative to + the per-env arena root node. DexSim Newton's ``POSE`` batch API, however, is + a **world-frame** body pose (see ``NewtonRigidDataType.POSE``). This adapter + bridges the two by adding (apply) / subtracting (fetch) the arena root + xy offset around the batch call, so the backend conforms to the same + local-pose contract as the default backend. Arenas are planar-translated + with identity rotation, so only the xy translation differs. + """ + + _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 + # Cached per-entity arena root xy offsets (world frame), entity-ordered. + # Arena root positions are static after build, so we cache to keep the + # per-step ``fetch_pose`` path off a Python loop over arenas. + self._arena_xy_cache: tuple[torch.Tensor, int] | None = None + # Cached (sorted body ids, argsort indices) for body_id -> entity-index + # lookup via ``torch.searchsorted``. Invalidated whenever body IDs are + # re-resolved (pre- -> post-finalization interleaved layout). + self._body_id_sort_cache: tuple[torch.Tensor, torch.Tensor] | None = None + + # -- 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 + ) + # Body IDs changed -> any sorted-lookup cache is stale. + self._body_id_sort_cache = None + 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, + ) + # DexSim POSE is world-frame; convert to arena-local for the public API. + offsets = self._arena_xy_offsets_for(body_ids) + if offsets is not None: + data[:, :2] -= offsets + + def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> None: + # Public pose is arena-local; convert to world-frame for DexSim POSE. + offsets = self._arena_xy_offsets_for(body_ids) + if offsets is not None: + pose = pose.clone() + pose[:, :2] += offsets + self._apply_data(body_ids, self._get_data_type().POSE, pose) + + def _arena_xy_offsets_for( + self, body_ids: torch.Tensor | None + ) -> torch.Tensor | None: + """Per-row arena root xy offsets (world frame) aligned with ``body_ids``. + + Returns ``None`` when no conversion is needed (no arenas, or arena count + does not match entity count), which makes the caller a safe no-op. + Otherwise returns an ``(N, 2)`` float32 tensor to add (local->world, + :meth:`apply_pose`) or subtract (world->local, :meth:`fetch_pose`). + """ + all_offsets = self._all_arena_xy_offsets() + if all_offsets is None: + return None + if body_ids is None: + # ``fetch_pose`` default: rows are entity-ordered, matching the cache. + return all_offsets + # Map each requested body id back to its entity index, then index the + # entity-ordered offset cache. ``body_ids`` passed to ``apply_pose`` + # always originate from ``select_body_ids`` (a subset of + # ``_body_ids_tensor``), so every value is guaranteed to match. Use a + # cached sorted-id lookup + binary search (O((N+M) log N)) instead of an + # O(N*M) equality matrix. + sorted_ids, sort_idx = self._body_id_sort_lookup() + bids = body_ids.to(device=sorted_ids.device, dtype=sorted_ids.dtype) + entity_idx = sort_idx[torch.searchsorted(sorted_ids, bids)] + return all_offsets[entity_idx] + + def _body_id_sort_lookup(self) -> tuple[torch.Tensor, torch.Tensor]: + """Cached ``(sorted body ids, argsort indices)`` for body_id lookup. + + ``sort_idx`` maps a position in the sorted body-id array back to the + entity index, so ``sort_idx[searchsorted(sorted_ids, ids)]`` resolves + any body id to its entity index in O(log N). + """ + if self._body_id_sort_cache is None: + self._ensure_body_ids() + sorted_ids, sort_idx = torch.sort(self._body_ids_tensor) # type: ignore[arg-type] + self._body_id_sort_cache = (sorted_ids, sort_idx) + return self._body_id_sort_cache + + def _all_arena_xy_offsets(self) -> torch.Tensor | None: + """Cached entity-ordered ``(num_instances, 2)`` arena root xy offsets.""" + from embodichain.lab.sim.utility import get_dexsim_arenas + + arenas = get_dexsim_arenas() + n = len(self.entities) + if len(arenas) == 0 or len(arenas) != n: + # Cannot map entities to arenas safely; no-op (preserve behavior). + return None + if self._arena_xy_cache is not None and self._arena_xy_cache[1] == n: + return self._arena_xy_cache[0] + offsets = torch.as_tensor( + np.stack( + [arena.get_root_node().get_local_pose()[:2, 3] for arena in arenas] + ).astype(np.float32), + device=self.device, + ) + self._arena_xy_cache = (offsets, n) + return offsets + + # -- 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/backends/spawn.py b/embodichain/lab/sim/objects/backends/spawn.py new file mode 100644 index 000000000..739d450de --- /dev/null +++ b/embodichain/lab/sim/objects/backends/spawn.py @@ -0,0 +1,694 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""EmbodiChain tensor-layout adapters for :mod:`dexsim.spawn` batches. + +The classes in this module deliberately know nothing about Default backend scenes or +Newton runtime objects. Backend selection, handle rebinding, and topology +revision tracking remain owned by DexSim's ``SpawnResult`` and batch classes. +EmbodiChain only adapts logical row selections and its public pose convention +``(x, y, z, qx, qy, qz, qw)``. + +Row and DOF selection is delegated to DexSim's public batches. This adapter is +therefore limited to EmbodiChain naming and tensor-layout conversion. +""" + +from __future__ import annotations + +from numbers import Integral +from typing import TYPE_CHECKING, Any, Sequence + +import torch + +from .base import ArticulationViewBase, RigidBodyViewBase + +if TYPE_CHECKING: + from dexsim.spawn import ArticulationBatch, RigidBodyBatch, SpawnResult + +__all__ = ["SpawnArticulationView", "SpawnRigidBodyView"] + +_NEWTON_ROOT_POSE_ATOL = 1.0e-6 + + +def _create_newton_standalone_state_sync( + model: Any, + body_ids: Sequence[int], +) -> Any: + """Create DexSim's reusable FREE-joint synchronization selection.""" + from dexsim.engine.newton_physics.rigid_body.state_sync import ( + StandaloneRigidStateSync, + ) + + return StandaloneRigidStateSync.from_body_ids(model, body_ids) + + +def _checked_batch_call( + batch: Any, + method_name: str, + /, + *args: Any, + **kwargs: Any, +) -> Any: + """Call one Spawn batch operation and reject native failure statuses.""" + status = getattr(batch, method_name)(*args, **kwargs) + if isinstance(status, Integral) and status < 0: + raise RuntimeError( + f"DexSim Spawn batch operation {method_name!r} failed with " + f"status {status}." + ) + return status + + +def _rows( + selection: Sequence[int] | torch.Tensor | None, + count: int, + device: torch.device, +) -> torch.Tensor: + if selection is None: + return torch.arange(count, dtype=torch.long, device=device) + result = torch.as_tensor(selection, dtype=torch.long, device=device).reshape(-1) + if torch.any(result < 0) or torch.any(result >= count): + raise IndexError(f"Batch row selection is outside [0, {count}).") + return result + + +def _spawn_pose(data: torch.Tensor) -> torch.Tensor: + """Convert rigid-body ``xyz+xyzw`` poses to Spawn ``xyzw+xyz``.""" + result = torch.empty_like(data, dtype=torch.float32) + result[..., 0:4] = data[..., 3:7] + result[..., 4:7] = data[..., 0:3] + return result + + +def _embodichain_pose(data: torch.Tensor) -> torch.Tensor: + """Convert Spawn ``xyzw+xyz`` poses to rigid-body ``xyz+xyzw``.""" + result = torch.empty_like(data, dtype=torch.float32) + result[..., 0:3] = data[..., 4:7] + result[..., 3:7] = data[..., 0:4] + return result + + +def _spawn_articulation_pose(data: torch.Tensor) -> torch.Tensor: + """Convert articulation ``xyz+xyzw`` poses to Spawn ``xyzw+xyz``.""" + return _spawn_pose(data) + + +def _embodichain_articulation_pose(data: torch.Tensor) -> torch.Tensor: + """Convert Spawn ``xyzw+xyz`` poses to articulation ``xyz+xyzw``.""" + return _embodichain_pose(data) + + +class _SpawnSelectionAdapter: + """Shared row-selection support for fixed-size Spawn batches.""" + + def __init__(self, batch: Any, device: torch.device, row_count: int) -> None: + self._batch = batch + self.device = device + self._row_count = row_count + + def _fetch_rows( + self, + method_name: str, + out: torch.Tensor, + selection: Sequence[int] | torch.Tensor | None, + tail_shape: tuple[int, ...], + ) -> torch.Tensor: + rows = _rows(selection, self._row_count, self.device) + selected = torch.empty( + (len(rows), *tail_shape), + dtype=torch.float32, + device=self.device, + ) + if len(rows): + _checked_batch_call(self._batch.select(rows), method_name, selected) + out.copy_(selected.to(device=out.device, dtype=out.dtype)) + return out + + def _apply_rows( + self, + method_name: str, + values: torch.Tensor, + selection: Sequence[int] | torch.Tensor, + tail_shape: tuple[int, ...], + ) -> None: + rows = _rows(selection, self._row_count, self.device) + values = values.to(device=self.device, dtype=torch.float32) + expected_shape = (len(rows), *tail_shape) + if tuple(values.shape) != expected_shape: + raise ValueError( + f"Expected selected data shape {expected_shape}, got " + f"{tuple(values.shape)}." + ) + if len(rows): + _checked_batch_call(self._batch.select(rows), method_name, values) + + +class SpawnRigidBodyView(_SpawnSelectionAdapter, RigidBodyViewBase): + """Backend-neutral rigid-body view backed by ``RigidBodyBatch``.""" + + def __init__( + self, + result: SpawnResult, + batch: RigidBodyBatch, + device: torch.device, + ) -> None: + super().__init__(batch, device, len(batch)) + self.result = result + self.batch = batch + self._body_ids_tensor = torch.arange( + len(batch), dtype=torch.int32, device=device + ) + self._newton_state_sync: tuple[int, Any, Any] | None = None + + @property + def is_ready(self) -> bool: + return True + + @property + def is_newton_backend(self) -> bool: + return self.result.backend == "newton" + + @property + def body_ids(self) -> list[int]: + return list(range(self._row_count)) + + @property + def body_ids_tensor(self) -> torch.Tensor: + return self._body_ids_tensor + + def select_body_ids(self, indices: Sequence[int] | torch.Tensor) -> torch.Tensor: + return self._body_ids_tensor[indices] + + def fetch_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + spawn = torch.empty((len(data), 7), dtype=torch.float32, device=self.device) + self._fetch_rows("fetch_pose", spawn, body_ids, (7,)) + data.copy_(_embodichain_pose(spawn).to(data.device, data.dtype)) + + def apply_pose(self, pose: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_pose", + _spawn_pose(pose.to(self.device, torch.float32)), + body_ids, + (7,), + ) + if self.is_newton_backend and len(body_ids): + self._synchronize_newton_standalone_state() + + def _synchronize_newton_standalone_state(self) -> None: + """Keep Newton standalone-body FREE joints coherent after state writes. + + DexSim 0.4.3's device ``RigidBodyBatch`` state writes update maximal + ``body_q`` or ``body_qd`` state, while MuJoCo-Warp advances standalone + rigid bodies from their reduced FREE-joint state. Cache one selection + for this stable batch and project both state buffers after each write. + """ + topology_revision = int(self.result.topology_revision) + cached = self._newton_state_sync + if cached is None or cached[0] != topology_revision: + # Accessing ``_binding`` refreshes a stale stable batch. DexSim + # currently exposes neither the Newton runtime nor this required + # synchronization through the public Batch API. + binding = self.batch._binding + runtime = getattr(binding, "_runtime", None) + indices = getattr(binding, "_indices", None) + if runtime is None or indices is None: + raise RuntimeError( + "Newton rigid-body batch has no finalized runtime selection." + ) + selected_body_ids = indices.detach().cpu().tolist() + state_sync = _create_newton_standalone_state_sync( + runtime.model, + selected_body_ids, + ) + cached = (topology_revision, runtime, state_sync) + self._newton_state_sync = cached + + _, runtime, state_sync = cached + state_sync.synchronize((runtime.current_state, runtime.other_state)) + + def fetch_com_local_pose( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + spawn = torch.empty((len(data), 7), dtype=torch.float32, device=self.device) + self._fetch_rows("fetch_com_local_pose", spawn, body_ids, (7,)) + data.copy_(_embodichain_pose(spawn).to(data.device, data.dtype)) + + def apply_com_local_pose(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_com_local_pose", + _spawn_pose(data.to(self.device, torch.float32)), + body_ids, + (7,), + ) + + def fetch_linear_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_linear_velocity", data, body_ids, (3,)) + + def fetch_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_angular_velocity", data, body_ids, (3,)) + + def apply_linear_velocity(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows( + "apply_linear_velocity", + data, + body_ids, + (3,), + ) + if self.is_newton_backend and len(body_ids): + self._synchronize_newton_standalone_state() + + def apply_angular_velocity( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_rows( + "apply_angular_velocity", + data, + body_ids, + (3,), + ) + if self.is_newton_backend and len(body_ids): + self._synchronize_newton_standalone_state() + + def fetch_linear_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_linear_acceleration", data, body_ids, (3,)) + + def fetch_angular_acceleration( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_angular_acceleration", data, body_ids, (3,)) + + def apply_force(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_force", data, body_ids, (3,)) + + def apply_torque(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_torque", data, body_ids, (3,)) + + def fetch_mass( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_mass", data, body_ids, (1,)) + + def apply_mass(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_mass", data, body_ids, (1,)) + + def fetch_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_inertia_diagonal", data, body_ids, (3,)) + + def apply_inertia_diagonal( + self, data: torch.Tensor, body_ids: torch.Tensor + ) -> None: + self._apply_rows( + "apply_inertia_diagonal", + data, + body_ids, + (3,), + ) + + def fetch_friction( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_friction", data, body_ids, (1,)) + + def apply_friction(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_friction", data, body_ids, (1,)) + + def fetch_restitution( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_restitution", data, body_ids, (1,)) + + def apply_restitution(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_restitution", data, body_ids, (1,)) + + def fetch_contact_offset( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_contact_offset", data, body_ids, (1,)) + + def apply_contact_offset(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_contact_offset", data, body_ids, (1,)) + + def fetch_damping( + self, data: torch.Tensor, body_ids: torch.Tensor | None = None + ) -> None: + self._fetch_rows("fetch_damping", data, body_ids, (2,)) + + def apply_damping(self, data: torch.Tensor, body_ids: torch.Tensor) -> None: + self._apply_rows("apply_damping", data, body_ids, (2,)) + + def fetch_collision_filter( + self, + data: torch.Tensor, + body_ids: torch.Tensor | None = None, + ) -> None: + rows = _rows(body_ids, self._row_count, self.device) + selected = torch.empty( + (len(rows), 4), + dtype=data.dtype, + device=self.device, + ) + if len(rows): + _checked_batch_call( + self.batch.select(rows), + "fetch_collision_filter", + selected, + ) + data.copy_(selected.to(device=data.device, dtype=data.dtype)) + + def apply_collision_filter( + self, + data: torch.Tensor, + body_ids: torch.Tensor, + ) -> None: + rows = _rows(body_ids, self._row_count, self.device) + expected_shape = (len(rows), 4) + if tuple(data.shape) != expected_shape: + raise ValueError( + f"Expected selected data shape {expected_shape}, got " + f"{tuple(data.shape)}." + ) + if len(rows): + _checked_batch_call( + self.batch.select(rows), + "apply_collision_filter", + data, + ) + + +class SpawnArticulationView(_SpawnSelectionAdapter, ArticulationViewBase): + """Backend-neutral articulation state view backed by ``ArticulationBatch``. + + Joint selections currently require one scalar DOF per selected joint. The + public DexSim layout already describes multi-DOF joints; supporting them + without ambiguity requires a DOF-selection API in DexSim and is therefore + kept as an explicit boundary rather than guessed here. + """ + + def __init__( + self, + result: SpawnResult, + batch: ArticulationBatch, + device: torch.device, + ) -> None: + super().__init__(batch, device, len(batch)) + self.result = result + self.batch = batch + self._validate_homogeneous_layout() + self._articulation_ids = torch.arange( + len(batch), dtype=torch.int32, device=device + ) + + def _validate_homogeneous_layout(self) -> None: + """Require the uniform topology promised by one EC Articulation.""" + dof_counts = tuple(self.batch.dof_counts) + link_counts = tuple(self.batch.link_counts) + joint_names = tuple(self.batch.joint_names_per_articulation) + link_names = tuple(self.batch.link_names_per_articulation) + if dof_counts and len(set(dof_counts)) != 1: + raise ValueError( + "One EmbodiChain Articulation cannot bind heterogeneous Spawn " + f"DOF counts: {dof_counts}." + ) + if link_counts and len(set(link_counts)) != 1: + raise ValueError( + "One EmbodiChain Articulation cannot bind heterogeneous Spawn " + f"link counts: {link_counts}." + ) + if joint_names and any(names != joint_names[0] for names in joint_names[1:]): + raise ValueError( + "One EmbodiChain Articulation requires identical active-joint " + "ordering in every Spawn row." + ) + if link_names and any(names != link_names[0] for names in link_names[1:]): + raise ValueError( + "One EmbodiChain Articulation requires identical link ordering " + "in every Spawn row." + ) + layouts = tuple(self.batch.joint_layouts_per_articulation) + if layouts and any(layout.dof_count != 1 for layout in layouts[0]): + raise NotImplementedError( + "EmbodiChain's Articulation API currently indexes joints and " + "scalar DOFs interchangeably. Spawn multi-DOF joints require " + "an explicit DOF-selection API before they can be bound safely." + ) + + @property + def dof(self) -> int: + """Scalar DOF width shared by every articulation row.""" + return self.batch.dof_width + + @property + def num_links(self) -> int: + """Link count shared by every articulation row.""" + return self.batch.link_width + + @property + def joint_names(self) -> list[str]: + """Active joints in public flattened-DOF order.""" + rows = self.batch.joint_names_per_articulation + return [] if not rows else list(rows[0]) + + @property + def link_names(self) -> list[str]: + """Links in public link-buffer order.""" + rows = self.batch.link_names_per_articulation + return [] if not rows else list(rows[0]) + + @property + def is_ready(self) -> bool: + return True + + @property + def is_newton_backend(self) -> bool: + return self.result.backend == "newton" + + @property + def articulation_ids_tensor(self) -> torch.Tensor: + return self._articulation_ids + + def select_articulation_ids( + self, env_ids: Sequence[int] | torch.Tensor + ) -> torch.Tensor: + return self._articulation_ids[env_ids] + + def fetch_root_pose(self, data: torch.Tensor) -> torch.Tensor: + spawn = torch.empty_like(data, dtype=torch.float32, device=self.device) + _checked_batch_call(self.batch, "fetch_root_pose", spawn) + data.copy_(_embodichain_articulation_pose(spawn).to(data.device, data.dtype)) + return data + + def fetch_root_linear_velocity(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_root_linear_velocity", data) + return data + + def fetch_root_angular_velocity(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_root_angular_velocity", data) + return data + + def fetch_qpos(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_joint_position", data) + return data + + def fetch_target_qpos(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_joint_target_position", data) + return data + + def fetch_qvel(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_joint_velocity", data) + return data + + def fetch_target_qvel(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_joint_target_velocity", data) + return data + + def fetch_qacc(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_joint_acceleration", data) + return data + + def fetch_qf(self, data: torch.Tensor) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_joint_force", data) + return data + + def fetch_link_pose(self, data: torch.Tensor) -> torch.Tensor: + spawn = torch.empty_like(data, dtype=torch.float32, device=self.device) + _checked_batch_call(self.batch, "fetch_link_pose", spawn) + data.copy_(_embodichain_articulation_pose(spawn).to(data.device, data.dtype)) + return data + + def fetch_link_velocity( + self, + data: torch.Tensor, + linear_data: torch.Tensor, + angular_data: torch.Tensor, + ) -> torch.Tensor: + _checked_batch_call(self.batch, "fetch_link_linear_velocity", linear_data) + _checked_batch_call(self.batch, "fetch_link_angular_velocity", angular_data) + data[..., 0:3] = linear_data + data[..., 3:6] = angular_data + return data + + def apply_root_pose( + self, pose: torch.Tensor, env_ids: Sequence[int] | torch.Tensor + ) -> None: + rows = _rows(env_ids, self._row_count, self.device) + spawn_pose = _spawn_articulation_pose(pose.to(self.device, torch.float32)) + if self.is_newton_backend and len(rows): + current_pose = torch.empty_like(spawn_pose) + self._fetch_rows( + "fetch_root_pose", + current_pose, + rows, + (7,), + ) + translation_matches = torch.all( + torch.abs(current_pose[:, 4:7] - spawn_pose[:, 4:7]) + <= _NEWTON_ROOT_POSE_ATOL, + dim=1, + ) + quaternion_delta = torch.minimum( + torch.amax(torch.abs(current_pose[:, 0:4] - spawn_pose[:, 0:4]), dim=1), + torch.amax(torch.abs(current_pose[:, 0:4] + spawn_pose[:, 0:4]), dim=1), + ) + changed = ~( + translation_matches & (quaternion_delta <= _NEWTON_ROOT_POSE_ATOL) + ) + rows = rows[changed] + spawn_pose = spawn_pose[changed] + + self._apply_rows( + "apply_root_pose", + spawn_pose, + rows, + (7,), + ) + + def _joint_columns(self, joint_ids: Sequence[int] | torch.Tensor) -> torch.Tensor: + ids = torch.as_tensor(joint_ids, dtype=torch.long, device=self.device) + layouts = self.batch.joint_layouts_per_articulation + if not layouts: + return ids + reference = layouts[0] + columns: list[int] = [] + for joint_id in ids.detach().cpu().tolist(): + layout = reference[joint_id] + if layout.dof_count != 1: + raise NotImplementedError( + "SpawnArticulationView needs DexSim DOF selection for " + f"multi-DOF joint {layout.name!r}." + ) + columns.append(layout.dof_start) + return torch.as_tensor(columns, dtype=torch.long, device=self.device) + + def _apply_joint_selection( + self, + values: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + apply_method: str, + ) -> None: + rows = _rows(env_ids, self._row_count, self.device) + columns = self._joint_columns(joint_ids) + values = values.to(device=self.device, dtype=torch.float32) + expected = (len(rows), len(columns)) + if tuple(values.shape) != expected: + raise ValueError( + f"Expected selected joint data shape {expected}, got " + f"{tuple(values.shape)}." + ) + if len(rows): + _checked_batch_call( + self.batch.select(rows), + apply_method, + values, + dof_ids=columns, + ) + + def apply_qpos( + self, + qpos: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + self._apply_joint_selection( + qpos, + env_ids, + joint_ids, + apply_method=( + "apply_joint_target_position" if target else "apply_joint_position" + ), + ) + + def apply_qvel( + self, + qvel: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + *, + target: bool, + ) -> None: + self._apply_joint_selection( + qvel, + env_ids, + joint_ids, + apply_method=( + "apply_joint_target_velocity" if target else "apply_joint_velocity" + ), + ) + + def apply_qf( + self, + qf: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor, + joint_ids: Sequence[int] | torch.Tensor, + ) -> None: + self._apply_joint_selection( + qf, + env_ids, + joint_ids, + apply_method="apply_joint_force", + ) + + def clear_dynamics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + rows = _rows(env_ids, self._row_count, self.device) + if not len(rows): + return + zeros = torch.zeros( + (len(rows), self.batch.dof_width), + dtype=torch.float32, + device=self.device, + ) + selected = self.batch.select(rows) + _checked_batch_call(selected, "apply_joint_velocity", zeros) + _checked_batch_call(selected, "apply_joint_target_velocity", zeros) + _checked_batch_call(selected, "apply_joint_force", zeros) + + def compute_kinematics(self, env_ids: Sequence[int] | torch.Tensor) -> None: + rows = _rows(env_ids, self._row_count, self.device) + if not len(rows): + return + _checked_batch_call(self.batch.select(rows), "compute_kinematics") diff --git a/embodichain/lab/sim/objects/cloth_object.py b/embodichain/lab/sim/objects/cloth_object.py index ef8419080..4fd1df7a5 100644 --- a/embodichain/lab/sim/objects/cloth_object.py +++ b/embodichain/lab/sim/objects/cloth_object.py @@ -14,441 +14,24 @@ # limitations under the License. # ---------------------------------------------------------------------------- -from __future__ import annotations +"""Compatibility exports for the surface-deformable object API.""" -import torch -import dexsim -import numpy as np -from functools import cached_property +from __future__ import annotations -from dataclasses import dataclass -from typing import List, Sequence, Union +from embodichain.lab.sim.cfg import ClothObjectCfg, SurfaceDeformableObjectCfg -from dexsim.models import MeshObject -from dexsim.engine import ClothBody, PhysicsScene -from dexsim.types import ClothBodyGPUAPIReadWriteType -from scipy.spatial import cKDTree -from embodichain.lab.sim.common import ( - BatchEntity, -) -from embodichain.lab.sim.material import ( - VisualMaterial, - VisualMaterialInst, - _capture_render_materials, - _restore_render_materials, - _wrap_first_render_material, -) -from embodichain.utils.math import ( - matrix_from_euler, -) -from embodichain.utils import logger -from embodichain.lab.sim.cfg import ( - ClothObjectCfg, +from .deformable.surface import ( + ClothBodyData, + ClothObject, + SurfaceDeformableData, + SurfaceDeformableObject, ) -from embodichain.utils.math import xyz_quat_to_4x4_matrix - -__all__ = ["ClothBodyData", "ClothObject", "ClothObjectCfg"] - - -@dataclass -class ClothBodyData: - """Data manager for cloth. - - Note: - 1. The pose data managed by dexsim is in the format of (qx, qy, qz, qw, x, y, z), but in EmbodiChain, we use (x, y, z, qw, qx, qy, qz) format. - """ - - def __init__( - self, entities: List[MeshObject], ps: PhysicsScene, device: torch.device - ) -> None: - """Initialize the ClothBodyData. - - Args: - entities (List[MeshObject]): List of MeshObjects representing the cloth bodies. - ps (PhysicsScene): The physics scene. - device (torch.device): The device to use for the cloth body data. - """ - self.entities = entities - # TODO: cloth body data can only be stored in cuda device for now. - self.device = device - # TODO: inorder to retrieve arena position, we need to access the node of each entity. - self.ps = ps - self.num_instances = len(entities) - - self.cloth_bodies: Sequence[ClothBody] = [ - self.entities[i].get_physical_body() for i in range(self.num_instances) - ] - self.n_vertices = self.cloth_bodies[0].get_num_vertices() - - self._rest_position_buffer = torch.empty( - (self.num_instances, self.n_vertices, 4), - device=self.device, - dtype=torch.float32, - ) - for i, cloth_body in enumerate(self.cloth_bodies): - self._rest_position_buffer[i] = cloth_body.get_position_inv_mass_buffer() - - self._vertex_position = torch.zeros( - (self.num_instances, self.n_vertices, 3), - device=self.device, - dtype=torch.float32, - ) - - self._vertex_velocity = torch.zeros( - (self.num_instances, self.n_vertices, 3), - device=self.device, - dtype=torch.float32, - ) - - @property - def rest_vertices(self): - """Get the rest position buffer of the cloth bodies.""" - return self._rest_position_buffer[:, :, :3].clone() - - @property - def vertex_position(self): - """Get the current vertex position buffer of the cloth bodies.""" - for i, clothbody in enumerate(self.cloth_bodies): - self._vertex_position[i] = clothbody.get_position_inv_mass_buffer()[:, :3] - return self._vertex_position.clone() - - @property - def vertex_velocity(self): - """Get the current vertex velocity buffer of the cloth bodies.""" - for i, clothbody in enumerate(self.cloth_bodies): - self._vertex_velocity[i] = clothbody.get_velocity_buffer()[:, 3:] - return self._vertex_velocity.clone() - - -class ClothObject(BatchEntity): - """ClothObject represents a batch of cloth body in the simulation.""" - - def __init__( - self, - cfg: ClothObjectCfg, - entities: List[MeshObject] = None, - device: torch.device = torch.device("cpu"), - ) -> None: - self._world: dexsim.World = dexsim.default_world() - self._ps = self._world.get_physics_scene() - self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() - - self._data = ClothBodyData(entities=entities, ps=self._ps, device=device) - - self._world.update(0.001) - self._surface_triangles = self._build_surface_triangles( - entities[0], - self._data.rest_vertices[0].detach().cpu().numpy(), - ) - - self._visual_material: List[VisualMaterialInst | None] = [None] * len(entities) - self.is_shared_visual_material = False - - super().__init__(cfg=cfg, entities=entities, device=device) - - self._initialize_existing_visual_material() - - self._set_default_collision_filter() - - @staticmethod - def _build_surface_triangles( - entity: MeshObject, - rest_vertices: np.ndarray, - ) -> np.ndarray: - """Map render triangles onto DexSim's welded cloth vertex buffer.""" - render_body = entity.get_render_body() - render_vertices: list[np.ndarray] = [] - render_triangles: list[np.ndarray] = [] - vertex_offset = 0 - for mesh_id in range(render_body.get_mesh_count()): - vertices = np.asarray( - render_body.get_vertices(mesh_id), - dtype=np.float32, - ) - triangles = np.asarray( - render_body.get_triangles(mesh_id), - dtype=np.int64, - ) - render_vertices.append(vertices) - render_triangles.append(triangles + vertex_offset) - vertex_offset += len(vertices) - - vertices = np.concatenate(render_vertices, axis=0) - triangles = np.concatenate(render_triangles, axis=0) - distances, cloth_vertex_ids = cKDTree(rest_vertices).query(vertices) - scale = max(float(np.ptp(rest_vertices, axis=0).max()), 1.0) - if float(distances.max(initial=0.0)) > scale * 1.0e-5: - raise RuntimeError( - "Could not map cloth render vertices onto the physical vertex buffer." - ) - return np.asarray(cloth_vertex_ids[triangles], dtype=np.int32) - - def _initialize_existing_visual_material(self) -> None: - """Wrap asset-parsed materials during cloth-object construction. - - For a multi-segment render body, the first segment with a valid - material is registered as the environment's representative material. - """ - self._original_visual_material = [[] for _ in self._entities] - self._original_visual_material_inst = [None] * len(self._entities) - for env_idx, entity in enumerate(self._entities): - render_body = entity.get_render_body() - if render_body is None: - continue - original_materials = _capture_render_materials(render_body) - self._original_visual_material[env_idx] = original_materials - wrapped = _wrap_first_render_material(original_materials) - if wrapped is not None: - self._visual_material[env_idx] = wrapped - self._original_visual_material_inst[env_idx] = wrapped - - def set_visual_material( - self, - mat: VisualMaterial, - env_ids: Sequence[int] | None = None, - shared: bool = False, - ) -> None: - """Set visual material for the cloth object. - - Args: - mat: The material template to assign. - env_ids: Environment indices. If None, all instances are used. - shared: Whether selected environments share one material instance. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - if shared: - if len(local_env_ids) != self.num_instances: - logger.log_error("Cannot share material instance for partial env_ids.") - mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}") - for env_idx in local_env_ids: - self._entities[env_idx].set_material(mat_inst.mat) - self._visual_material[env_idx] = mat_inst - self.is_shared_visual_material = True - else: - for env_idx in local_env_ids: - mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}") - self._entities[env_idx].set_material(mat_inst.mat) - self._visual_material[env_idx] = mat_inst - self.is_shared_visual_material = False - - def restore_visual_material(self, env_ids: Sequence[int] | None = None) -> None: - """Restore visual materials captured when the cloth object was created. - - Args: - env_ids: Environment indices. If None, all instances are restored. - """ - if not hasattr(self, "_original_visual_material"): - return - local_env_ids = self._all_indices if env_ids is None else env_ids - for env_idx in local_env_ids: - render_body = self._entities[env_idx].get_render_body() - if render_body is None: - continue - _restore_render_materials( - render_body, self._original_visual_material[env_idx] - ) - self._visual_material[env_idx] = self._original_visual_material_inst[ - env_idx - ] - self.is_shared_visual_material = False - - def get_visual_material_inst( - self, env_ids: Sequence[int] | None = None - ) -> List[VisualMaterialInst | None]: - """Get the material instance registered for each selected environment. - - Args: - env_ids: Environment indices. If None, all instances are returned. - - Returns: - The existing material wrappers, or None where an asset has no material. - """ - ids = env_ids if env_ids is not None else range(self.num_instances) - return [self._visual_material[i] for i in ids] - - def _set_default_collision_filter(self) -> None: - collision_filter_data = torch.zeros( - size=(self.num_instances, 4), dtype=torch.int32 - ) - for i in range(self.num_instances): - collision_filter_data[i, 0] = i - collision_filter_data[i, 1] = 1 - self.set_collision_filter(collision_filter_data) - - def set_collision_filter( - self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None - ) -> None: - """Set collision filter data for the cloth object. - - Args: - filter_data (torch.Tensor): [N, 4] of int. - First element of each object is arena id. - If 2nd element is 0, the object will collision with all other objects in world. - 3rd and 4th elements are not used currently. - - 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 - - if len(local_env_ids) != len(filter_data): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." - ) - - filter_data_np = filter_data.cpu().numpy().astype(np.uint32) - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].get_physical_body().set_collision_filter_data( - filter_data_np[i] - ) - - @property - def body_data(self) -> ClothBodyData | None: - """Get the cloth body data manager for this cloth object. - - Returns: - ClothBodyData | None: The cloth body data manager. - """ - return self._data - - def get_rest_vertex_position(self) -> torch.Tensor: - """Get the rest vertex position of the cloth bodies. - - Returns: - torch.Tensor: The rest vertex position of the cloth bodies, shape (num_instances, n_vertices, 3). - """ - return self._data.rest_vertices - - def get_current_vertex_position(self) -> torch.Tensor: - """Get the current vertex position of the cloth bodies. - - Returns: - torch.Tensor: The current vertex position of the cloth bodies, shape (num_instances, n_vertices, 3). - """ - return self._data.vertex_position - - def get_current_vertex_velocity(self) -> torch.Tensor: - """Get the current vertex velocity of the cloth bodies. - - Returns: - torch.Tensor: The current vertex velocity of the cloth bodies, shape (num_instances, n_vertices, 3). - """ - return self._data.vertex_velocity - - def get_triangles(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: - """Get surface triangle indices for selected cloth instances. - - Args: - env_ids: Environment indices. If ``None``, returns all instances. - - Returns: - Triangle indices with shape ``(N, num_triangles, 3)``. - """ - ids = self._all_indices if env_ids is None else env_ids - triangles = torch.as_tensor( - self._surface_triangles, - dtype=torch.int32, - device=self.device, - ) - return triangles.unsqueeze(0).expand(len(ids), -1, -1).clone() - - def set_local_pose( - self, pose: torch.Tensor, env_ids: Sequence[int] | None = None - ) -> None: - """Set local pose of the cloth object. - - Args: - pose (torch.Tensor): The local pose of the cloth object with shape (N, 7) or (N, 4, 4). - env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used. - """ - from embodichain.lab.sim import SimulationManager - - sim = SimulationManager.get_instance() - - local_env_ids = self._all_indices if env_ids is None else env_ids - - if len(local_env_ids) != len(pose): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." - ) - - if pose.dim() == 2 and pose.shape[1] == 7: - pose4x4 = xyz_quat_to_4x4_matrix(pose) - elif pose.dim() == 3 and pose.shape[1:3] == (4, 4): - pose4x4 = pose - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - arena_offsets = sim.arena_offsets - for i, env_idx in enumerate(local_env_ids): - # TODO: cloth body cannot directly set by `set_local_pose` currently. - rest_vertices = self.body_data.rest_vertices[i] - rotation = pose4x4[i][:3, :3] - translation = pose4x4[i][:3, 3] - - # apply transformation to local rest vertices and back - rest_vertices_local = rest_vertices - arena_offsets[i] - transformed_vertices = rest_vertices_local @ rotation.T + translation - transformed_vertices = transformed_vertices + arena_offsets[i] - - cloth_body: ClothBody = self._entities[env_idx].get_physical_body() - position_buffer = cloth_body.get_position_inv_mass_buffer() - velocity_buffer = cloth_body.get_velocity_buffer() - position_buffer[:, :3] = transformed_vertices - velocity_buffer[:, 3:] = 0.0 - - cloth_body.mark_dirty(ClothBodyGPUAPIReadWriteType.ALL) - # TODO: currently cloth body has no wake up interface, use set_wake_counter and pass in a positive value to wake it up - cloth_body.set_wake_counter(0.4) - - def get_local_pose(self, to_matrix=False): - """Get local pose of the cloth 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. - - Returns: - torch.Tensor: The local pose of the cloth object with shape (N, 7) or (N, 4, 4) depending on `to_matrix`. - """ - raise NotImplementedError( - "Getting local pose for ClothObject is not supported." - ) - - 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.restore_visual_material(env_ids=local_env_ids) - - # TODO: set attr for cloth body after loading in physics scene. - - # rest cloth body to init_pos - pos = torch.as_tensor( - self.cfg.init_pos, dtype=torch.float32, device=self.device - ) - rot = ( - torch.as_tensor(self.cfg.init_rot, dtype=torch.float32, device=self.device) - * torch.pi - / 180.0 - ) - pos = pos.unsqueeze(0).repeat(num_instances, 1) - rot = rot.unsqueeze(0).repeat(num_instances, 1) - mat = matrix_from_euler(rot, "XYZ") - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze(0) - .repeat(num_instances, 1, 1) - ) - pose[:, :3, 3] = pos - pose[:, :3, :3] = mat - self.set_local_pose(pose, env_ids=local_env_ids) - def destroy(self) -> None: - # TODO: not tested yet - 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) +__all__ = [ + "ClothBodyData", + "ClothObject", + "ClothObjectCfg", + "SurfaceDeformableData", + "SurfaceDeformableObject", + "SurfaceDeformableObjectCfg", +] diff --git a/embodichain/lab/sim/objects/deformable/__init__.py b/embodichain/lab/sim/objects/deformable/__init__.py new file mode 100644 index 000000000..91bf6b72c --- /dev/null +++ b/embodichain/lab/sim/objects/deformable/__init__.py @@ -0,0 +1,47 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Unified deformable-object API with DexSim volume/surface specializations.""" + +from __future__ import annotations + +from .base import DeformableObject +from .data import DeformableObjectData +from .surface import ( + ClothBodyData, + ClothObject, + SurfaceDeformableData, + SurfaceDeformableObject, +) +from .volume import ( + SoftBodyData, + SoftObject, + VolumeDeformableData, + VolumeDeformableObject, +) + +__all__ = [ + "ClothBodyData", + "ClothObject", + "DeformableObject", + "DeformableObjectData", + "SoftBodyData", + "SoftObject", + "SurfaceDeformableData", + "SurfaceDeformableObject", + "VolumeDeformableData", + "VolumeDeformableObject", +] diff --git a/embodichain/lab/sim/objects/deformable/base.py b/embodichain/lab/sim/objects/deformable/base.py new file mode 100644 index 000000000..86740a5c2 --- /dev/null +++ b/embodichain/lab/sim/objects/deformable/base.py @@ -0,0 +1,413 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Common facade for volume and surface deformable objects.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from copy import deepcopy +from typing import TYPE_CHECKING, Any, ClassVar, Literal, Sequence + +import dexsim +import numpy as np +import torch + +from embodichain.lab.sim.cfg import DeformableObjectCfg +from embodichain.lab.sim.common import BatchEntity +from embodichain.lab.sim.material import ( + VisualMaterial, + VisualMaterialInst, + _capture_render_materials, + _restore_render_materials, + _wrap_first_render_material, +) +from embodichain.utils import logger +from embodichain.utils.math import matrix_from_euler, xyz_quat_to_4x4_matrix + +from .data import DeformableObjectData + +if TYPE_CHECKING: + from dexsim.engine import PhysicsScene + from dexsim.spawn import SpawnResult + +__all__ = ["DeformableObject"] + + +class DeformableObject(BatchEntity, ABC): + """Common facade for a batch of deformable assets. + + The public nodal and surface contracts are backend-neutral. The concrete + implementations in this package currently bind them to DexSim soft-body + and cloth buffers. Newton support can be added as a separate implementation + without changing manager or visualization consumers. + """ + + deformable_type: ClassVar[Literal["volume", "surface"]] + spawn_kind: ClassVar[str] + display_name: ClassVar[str] + + def __init__( + self, + cfg: DeformableObjectCfg, + entities: Sequence[Any] | None = None, + device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, + ) -> None: + if cfg.deformable_type != self.deformable_type: + raise ValueError( + f"{type(self).__name__} requires deformable_type=" + f"{self.deformable_type!r}, got {cfg.deformable_type!r}." + ) + + if entities is None: + self._initialize_declared(cfg, device, declared_num_instances) + return + + entities = list(entities) + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result + if spawn_result is None: + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps: PhysicsScene | None = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = self._world.get_physics_scene() + self._all_indices = list(range(len(entities))) + + self._data = self._create_data(entities, self._ps, device) + if spawn_result is None: + self._world.update(0.001) + self._initialize_topology(entities) + + self._visual_material: list[VisualMaterialInst | None] = [None] * len(entities) + self.is_shared_visual_material = False + + super().__init__(cfg=cfg, entities=entities, device=device) + self._initialize_existing_visual_material() + self.reset() + self._set_default_collision_filter() + + def _initialize_declared( + self, + cfg: DeformableObjectCfg, + device: torch.device, + declared_num_instances: int | None, + ) -> None: + """Initialize a facade before Spawn materializes native handles.""" + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + f"A declared {type(self).__name__} requires " + "declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities: list[Any] = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._world = None + self._ps = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._visual_material = [None] * declared_num_instances + self.is_shared_visual_material = False + + @abstractmethod + def _create_data( + self, + entities: Sequence[Any], + physics_scene: PhysicsScene, + device: torch.device, + ) -> DeformableObjectData: + """Create the concrete backend data view.""" + + def _initialize_topology(self, entities: Sequence[Any]) -> None: + """Initialize implementation-specific surface topology.""" + del entities + + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized Spawn result.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its Spawn result binding.""" + return self._world is None + + @property + def num_instances(self) -> int: + """Return the materialized or declared instance count.""" + return len(self._entities) if self._entities else self._declared_num_instances + + @property + def data(self) -> DeformableObjectData | None: + """Return the common deformable data view after Spawn binding.""" + return self._data + + def attach_spawn_handles(self, entities: Sequence[Any]) -> None: + """Store materialized handles before final Spawn binding.""" + self._entities = list(entities) + + def bind_spawn(self, result: SpawnResult) -> None: + """Bind a declared facade to finalized native handles in place.""" + entities = list(self._entities) + if self.cfg.shape.compute_uv: + for entity in entities: + entity.compute_uv_mapping() + type(self).__init__( + self, + self.cfg, + entities, + self.device, + spawn_result=result, + ) + + def __str__(self) -> str: + if self.is_declared: + return ( + f"{self.__class__}: declared {self.num_instances} Spawn " + f"{self.display_name} objects | uid: {self.uid} | " + f"device: {self.device}" + ) + return super().__str__() + + def _initialize_existing_visual_material(self) -> None: + """Capture and wrap materials parsed from the source asset.""" + self._original_visual_material = [[] for _ in self._entities] + self._original_visual_material_inst = [None] * len(self._entities) + for env_idx, entity in enumerate(self._entities): + render_body = entity.get_render_body() + if render_body is None: + continue + original_materials = _capture_render_materials(render_body) + self._original_visual_material[env_idx] = original_materials + wrapped = _wrap_first_render_material(original_materials) + if wrapped is not None: + self._visual_material[env_idx] = wrapped + self._original_visual_material_inst[env_idx] = wrapped + + def set_visual_material( + self, + mat: VisualMaterial, + env_ids: Sequence[int] | None = None, + shared: bool = False, + ) -> None: + """Assign visual material instances to selected environments.""" + local_env_ids = self._resolve_env_ids(env_ids) + if shared: + if len(local_env_ids) != self.num_instances: + logger.log_error("Cannot share material instance for partial env_ids.") + mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}") + for env_idx in local_env_ids: + self._entities[env_idx].set_material(mat_inst.mat) + self._visual_material[env_idx] = mat_inst + self.is_shared_visual_material = True + return + + for env_idx in local_env_ids: + mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}") + self._entities[env_idx].set_material(mat_inst.mat) + self._visual_material[env_idx] = mat_inst + self.is_shared_visual_material = False + + def restore_visual_material(self, env_ids: Sequence[int] | None = None) -> None: + """Restore materials captured when the deformable was created.""" + if not hasattr(self, "_original_visual_material"): + return + for env_idx in self._resolve_env_ids(env_ids): + render_body = self._entities[env_idx].get_render_body() + if render_body is None: + continue + _restore_render_materials( + render_body, self._original_visual_material[env_idx] + ) + self._visual_material[env_idx] = self._original_visual_material_inst[ + env_idx + ] + self.is_shared_visual_material = False + + def get_visual_material_inst( + self, env_ids: Sequence[int] | None = None + ) -> list[VisualMaterialInst | None]: + """Return registered material wrappers for selected environments.""" + return [self._visual_material[i] for i in self._resolve_env_ids(env_ids)] + + def _set_default_collision_filter(self) -> None: + collision_filter_data = torch.zeros( + size=(self.num_instances, 4), dtype=torch.int32 + ) + collision_filter_data[:, 0] = torch.arange( + self.num_instances, dtype=torch.int32 + ) + collision_filter_data[:, 1] = 1 + self.set_collision_filter(collision_filter_data) + + def set_collision_filter( + self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None + ) -> None: + """Set native collision-filter data for selected environments.""" + local_env_ids = self._resolve_env_ids(env_ids) + if len(local_env_ids) != len(filter_data): + logger.log_error( + f"Length of env_ids {len(local_env_ids)} does not match filter " + f"data length {len(filter_data)}." + ) + filter_data_np = filter_data.detach().cpu().numpy().astype(np.uint32) + for i, env_idx in enumerate(local_env_ids): + self._entities[env_idx].get_physical_body().set_collision_filter_data( + filter_data_np[i] + ) + + def _resolve_env_ids(self, env_ids: Sequence[int] | None) -> list[int]: + if env_ids is None: + return list(self._all_indices) + if isinstance(env_ids, torch.Tensor): + ids = env_ids.detach().cpu().reshape(-1).tolist() + else: + ids = list(env_ids) + resolved = [int(env_id) for env_id in ids] + if any(env_id < 0 or env_id >= self.num_instances for env_id in resolved): + raise IndexError( + f"Environment IDs {resolved!r} are outside [0, {self.num_instances})." + ) + return resolved + + def set_local_pose( + self, pose: torch.Tensor, env_ids: Sequence[int] | None = None + ) -> None: + """Set deformable pose by transforming its rest-node buffers.""" + from embodichain.lab.sim import SimulationManager + + local_env_ids = self._resolve_env_ids(env_ids) + if len(local_env_ids) != len(pose): + logger.log_error( + f"Length of env_ids {len(local_env_ids)} does not match pose " + f"length {len(pose)}." + ) + if pose.dim() == 2 and pose.shape[1] == 7: + pose4x4 = xyz_quat_to_4x4_matrix(pose) + elif pose.dim() == 3 and pose.shape[1:] == (4, 4): + pose4x4 = pose + else: + logger.log_error( + f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." + ) + + sim = SimulationManager.get_instance() + self._apply_local_pose( + pose4x4.to(device=self.device, dtype=torch.float32), + local_env_ids, + sim.arena_offsets, + ) + + @abstractmethod + def _apply_local_pose( + self, + pose: torch.Tensor, + env_ids: Sequence[int], + arena_offsets: torch.Tensor, + ) -> None: + """Apply rest-node transforms to native backend buffers.""" + + def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: + """Reject root-pose reads because deformables have no rigid root pose.""" + del to_matrix + raise NotImplementedError( + f"Getting local pose for {type(self).__name__} is not supported." + ) + + def get_current_nodal_position(self) -> torch.Tensor: + """Return current simulation-node positions in world frame.""" + self._require_data() + return self.data.nodal_pos_w + + def get_current_nodal_velocity(self) -> torch.Tensor: + """Return current simulation-node velocities in world frame.""" + self._require_data() + return self.data.nodal_vel_w + + def get_current_nodal_state(self) -> torch.Tensor: + """Return current simulation-node state ``[position, velocity]``.""" + self._require_data() + return self.data.nodal_state_w + + def get_default_nodal_state(self) -> torch.Tensor: + """Return default simulation-node state ``[position, velocity]``.""" + self._require_data() + return self.data.default_nodal_state_w + + def _require_data(self) -> None: + if self.data is None: + raise RuntimeError( + f"{type(self).__name__} data is unavailable before Spawn finalization." + ) + + @abstractmethod + def get_surface_vertices(self) -> torch.Tensor: + """Return visualization/collision surface vertices in world frame.""" + + @abstractmethod + def get_surface_triangles( + self, env_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Return surface triangle indices for selected environments.""" + + def get_triangles(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: + """Compatibility alias for :meth:`get_surface_triangles`.""" + return self.get_surface_triangles(env_ids=env_ids) + + def reset(self, env_ids: Sequence[int] | None = None) -> None: + """Restore initial pose, zero nodal velocity, and source materials.""" + local_env_ids = self._resolve_env_ids(env_ids) + self.restore_visual_material(env_ids=local_env_ids) + num_instances = len(local_env_ids) + + pos = torch.as_tensor( + self.cfg.init_pos, dtype=torch.float32, device=self.device + ).repeat(num_instances, 1) + rot = ( + torch.as_tensor(self.cfg.init_rot, dtype=torch.float32, device=self.device) + * torch.pi + / 180.0 + ).repeat(num_instances, 1) + pose = ( + torch.eye(4, dtype=torch.float32, device=self.device) + .unsqueeze(0) + .repeat(num_instances, 1, 1) + ) + pose[:, :3, 3] = pos + pose[:, :3, :3] = matrix_from_euler(rot, "XYZ") + self.set_local_pose(pose, env_ids=local_env_ids) + + def destroy(self) -> None: + """Destroy legacy directly-created native entities. + + Spawn-bound entities are owned and released by ``SpawnResult``. + """ + if self.is_spawn_bound or self.is_declared: + return + 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) diff --git a/embodichain/lab/sim/objects/deformable/data.py b/embodichain/lab/sim/objects/deformable/data.py new file mode 100644 index 000000000..f9210e415 --- /dev/null +++ b/embodichain/lab/sim/objects/deformable/data.py @@ -0,0 +1,64 @@ +# ---------------------------------------------------------------------------- +# 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-neutral data contract for deformable simulation objects.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch + +__all__ = ["DeformableObjectData"] + + +class DeformableObjectData(ABC): + """Common nodal-state view for volume and surface deformables. + + Positions and velocities use the simulation world frame. Concrete + backends own how the buffers are fetched; consumers can rely on a stable + ``(num_instances, num_nodes, 3)`` contract. + """ + + @property + @abstractmethod + def nodal_pos_w(self) -> torch.Tensor: + """Return current simulation-node positions in world frame.""" + + @property + @abstractmethod + def nodal_vel_w(self) -> torch.Tensor: + """Return current simulation-node velocities in world frame.""" + + @property + @abstractmethod + def default_nodal_state_w(self) -> torch.Tensor: + """Return default nodal state ``[position, velocity]`` in world frame.""" + + @property + def nodal_state_w(self) -> torch.Tensor: + """Return current nodal state ``[position, velocity]`` in world frame.""" + return torch.cat((self.nodal_pos_w, self.nodal_vel_w), dim=-1) + + @property + def root_pos_w(self) -> torch.Tensor: + """Return the mean nodal position for each deformable instance.""" + return self.nodal_pos_w.mean(dim=1) + + @property + def root_vel_w(self) -> torch.Tensor: + """Return the mean nodal velocity for each deformable instance.""" + return self.nodal_vel_w.mean(dim=1) diff --git a/embodichain/lab/sim/objects/deformable/surface.py b/embodichain/lab/sim/objects/deformable/surface.py new file mode 100644 index 000000000..bd3df53fd --- /dev/null +++ b/embodichain/lab/sim/objects/deformable/surface.py @@ -0,0 +1,237 @@ +# ---------------------------------------------------------------------------- +# 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 surface-deformable object implementation.""" + +from __future__ import annotations + +from typing import Any, Sequence + +import numpy as np +import torch +from dexsim.engine import ClothBody, PhysicsScene +from dexsim.models import MeshObject +from dexsim.types import ClothBodyGPUAPIReadWriteType +from scipy.spatial import cKDTree + +from .base import DeformableObject +from .data import DeformableObjectData + +__all__ = [ + "ClothBodyData", + "ClothObject", + "SurfaceDeformableData", + "SurfaceDeformableObject", +] + + +class SurfaceDeformableData(DeformableObjectData): + """DexSim cloth buffers exposed through the common nodal contract.""" + + def __init__( + self, + entities: Sequence[MeshObject], + ps: PhysicsScene, + device: torch.device, + ) -> None: + self.entities = list(entities) + self.device = device + self.ps = ps + self.num_instances = len(self.entities) + self.cloth_bodies: Sequence[ClothBody] = [ + entity.get_physical_body() for entity in self.entities + ] + self.n_vertices = self.cloth_bodies[0].get_num_vertices() + + self._rest_position_buffer = torch.empty( + (self.num_instances, self.n_vertices, 4), + device=self.device, + dtype=torch.float32, + ) + for i, cloth_body in enumerate(self.cloth_bodies): + self._rest_position_buffer[i] = cloth_body.get_rest_position_buffer() + + self._vertex_position = torch.zeros( + (self.num_instances, self.n_vertices, 3), + device=self.device, + dtype=torch.float32, + ) + self._vertex_velocity = torch.zeros_like(self._vertex_position) + self._default_nodal_state_w = torch.cat( + ( + self._rest_position_buffer[..., :3], + torch.zeros_like(self._rest_position_buffer[..., :3]), + ), + dim=-1, + ) + + @property + def rest_vertices(self) -> torch.Tensor: + """Return rest surface vertices in simulation world frame.""" + return self._rest_position_buffer[..., :3].clone() + + @property + def vertex_position(self) -> torch.Tensor: + """Return current surface vertices in simulation world frame.""" + for i, cloth_body in enumerate(self.cloth_bodies): + self._vertex_position[i] = cloth_body.get_position_inv_mass_buffer()[:, :3] + return self._vertex_position.clone() + + @property + def vertex_velocity(self) -> torch.Tensor: + """Return current surface-vertex velocities.""" + for i, cloth_body in enumerate(self.cloth_bodies): + # DexSim stores velocity in the first xyz channels. The fourth + # channel is padding/metadata and must not be exposed as velocity. + self._vertex_velocity[i] = cloth_body.get_velocity_buffer()[:, :3] + return self._vertex_velocity.clone() + + @property + def nodal_pos_w(self) -> torch.Tensor: + return self.vertex_position + + @property + def nodal_vel_w(self) -> torch.Tensor: + return self.vertex_velocity + + @property + def default_nodal_state_w(self) -> torch.Tensor: + return self._default_nodal_state_w.clone() + + +class SurfaceDeformableObject(DeformableObject): + """A batch of DexSim surface deformables backed by ``ClothBody``.""" + + deformable_type = "surface" + spawn_kind = "cloth_object" + display_name = "surface deformable" + + def _create_data( + self, + entities: Sequence[Any], + physics_scene: PhysicsScene, + device: torch.device, + ) -> SurfaceDeformableData: + return SurfaceDeformableData(entities, physics_scene, device) + + def _initialize_topology(self, entities: Sequence[Any]) -> None: + self._surface_triangles = self._build_surface_triangles( + entities[0], + self.body_data.rest_vertices[0].detach().cpu().numpy(), + self.body_data.cloth_bodies[0].get_initial_transform(), + ) + + @property + def body_data(self) -> SurfaceDeformableData | None: + """Compatibility view of the DexSim cloth data.""" + return self._data + + @staticmethod + def _build_surface_triangles( + entity: MeshObject, + rest_vertices: np.ndarray, + initial_transform: np.ndarray, + ) -> np.ndarray: + """Map render triangles onto DexSim's welded cloth vertex buffer.""" + render_body = entity.get_render_body() + render_vertices: list[np.ndarray] = [] + render_triangles: list[np.ndarray] = [] + vertex_offset = 0 + for mesh_id in range(render_body.get_mesh_count()): + vertices = np.asarray(render_body.get_vertices(mesh_id), dtype=np.float32) + triangles = np.asarray(render_body.get_triangles(mesh_id), dtype=np.int64) + render_vertices.append(vertices) + render_triangles.append(triangles + vertex_offset) + vertex_offset += len(vertices) + + vertices = np.concatenate(render_vertices, axis=0) + triangles = np.concatenate(render_triangles, axis=0) + initial_transform = np.asarray(initial_transform, dtype=np.float32).reshape( + 4, 4 + ) + vertices = vertices @ initial_transform[:3, :3].T + initial_transform[:3, 3] + distances, cloth_vertex_ids = cKDTree(rest_vertices).query(vertices) + scale = max(float(np.ptp(rest_vertices, axis=0).max()), 1.0) + if float(distances.max(initial=0.0)) > scale * 1.0e-5: + raise RuntimeError( + "Could not map surface-deformable render vertices onto the " + "physical vertex buffer." + ) + return np.asarray(cloth_vertex_ids[triangles], dtype=np.int32) + + def _apply_local_pose( + self, + pose: torch.Tensor, + env_ids: Sequence[int], + arena_offsets: torch.Tensor, + ) -> None: + self._require_data() + rest_vertices = self.body_data.rest_vertices + for i, env_idx in enumerate(env_ids): + cloth_body: ClothBody = self._entities[env_idx].get_physical_body() + initial_transform = torch.as_tensor( + cloth_body.get_initial_transform(), + dtype=torch.float32, + device=self.device, + ) + rest_vertices_local = ( + rest_vertices[env_idx] - initial_transform[:3, 3] + ) @ initial_transform[:3, :3] + rotation = pose[i, :3, :3] + translation = pose[i, :3, 3] + arena_offset = torch.as_tensor( + arena_offsets[env_idx], dtype=torch.float32, device=self.device + ) + transformed_vertices = ( + rest_vertices_local @ rotation.T + translation + arena_offset + ) + + cloth_body.get_position_inv_mass_buffer()[:, :3] = transformed_vertices + cloth_body.get_velocity_buffer()[:, :3] = 0.0 + cloth_body.mark_dirty(ClothBodyGPUAPIReadWriteType.ALL) + cloth_body.set_wake_counter(0.4) + + def get_rest_vertex_position(self) -> torch.Tensor: + """Return rest surface-vertex positions.""" + self._require_data() + return self.body_data.rest_vertices + + def get_current_vertex_position(self) -> torch.Tensor: + """Return current surface-vertex positions.""" + return self.get_current_nodal_position() + + def get_current_vertex_velocity(self) -> torch.Tensor: + """Return current surface-vertex velocities.""" + return self.get_current_nodal_velocity() + + def get_surface_vertices(self) -> torch.Tensor: + """Return the live cloth surface used for visualization.""" + return self.get_current_vertex_position() + + def get_surface_triangles( + self, env_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Return surface triangle indices for selected instances.""" + ids = self._resolve_env_ids(env_ids) + triangles = torch.as_tensor( + self._surface_triangles, dtype=torch.int32, device=self.device + ) + return triangles.unsqueeze(0).expand(len(ids), -1, -1).clone() + + +# Compatibility names retained for existing environments and tutorials. +ClothBodyData = SurfaceDeformableData +ClothObject = SurfaceDeformableObject diff --git a/embodichain/lab/sim/objects/deformable/volume.py b/embodichain/lab/sim/objects/deformable/volume.py new file mode 100644 index 000000000..b3ecf11ea --- /dev/null +++ b/embodichain/lab/sim/objects/deformable/volume.py @@ -0,0 +1,282 @@ +# ---------------------------------------------------------------------------- +# 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 volume-deformable object implementation.""" + +from __future__ import annotations + +from functools import cached_property +from typing import Any, Sequence + +import numpy as np +import torch +from dexsim.engine import PhysicsScene, SoftBody +from dexsim.models import MeshObject +from dexsim.types import SoftBodyGPUAPIReadWriteType +from scipy.spatial import ConvexHull, QhullError + +from embodichain.utils import logger + +from .base import DeformableObject +from .data import DeformableObjectData + +__all__ = [ + "SoftBodyData", + "SoftObject", + "VolumeDeformableData", + "VolumeDeformableObject", +] + + +class VolumeDeformableData(DeformableObjectData): + """DexSim soft-body buffers exposed through the common nodal contract.""" + + def __init__( + self, + entities: Sequence[MeshObject], + ps: PhysicsScene, + device: torch.device, + ) -> None: + self.entities = list(entities) + self.device = device + self.ps = ps + self.num_instances = len(self.entities) + self.soft_bodies: Sequence[SoftBody] = [ + entity.get_physical_body() for entity in self.entities + ] + self.n_collision_vertices = self.soft_bodies[0].get_num_vertices() + self.n_sim_vertices = self.soft_bodies[0].get_num_sim_vertices() + + self._rest_position_buffer = torch.empty( + (self.num_instances, self.n_collision_vertices, 4), + device=self.device, + dtype=torch.float32, + ) + self._rest_sim_position_buffer = torch.empty( + (self.num_instances, self.n_sim_vertices, 4), + device=self.device, + dtype=torch.float32, + ) + for i, soft_body in enumerate(self.soft_bodies): + self._rest_position_buffer[i] = soft_body.get_position_inv_mass_buffer() + self._rest_sim_position_buffer[i] = ( + soft_body.get_sim_position_inv_mass_buffer() + ) + + self._collision_position = torch.zeros( + (self.num_instances, self.n_collision_vertices, 3), + device=self.device, + dtype=torch.float32, + ) + self._sim_vertex_position = torch.zeros( + (self.num_instances, self.n_sim_vertices, 3), + device=self.device, + dtype=torch.float32, + ) + self._sim_vertex_velocity = torch.zeros_like(self._sim_vertex_position) + self._default_nodal_state_w = torch.cat( + ( + self._rest_sim_position_buffer[..., :3], + torch.zeros_like(self._rest_sim_position_buffer[..., :3]), + ), + dim=-1, + ) + + @property + def rest_collision_vertices(self) -> torch.Tensor: + """Return rest collision vertices in simulation world frame.""" + return self._rest_position_buffer[..., :3].clone() + + @property + def rest_sim_vertices(self) -> torch.Tensor: + """Return rest simulation vertices in simulation world frame.""" + return self._rest_sim_position_buffer[..., :3].clone() + + @property + def collision_position(self) -> torch.Tensor: + """Return current collision vertices in simulation world frame.""" + for i, soft_body in enumerate(self.soft_bodies): + self._collision_position[i] = soft_body.get_position_inv_mass_buffer()[ + :, :3 + ] + return self._collision_position.clone() + + @property + def sim_vertex_position(self) -> torch.Tensor: + """Return current simulation vertices in simulation world frame.""" + for i, soft_body in enumerate(self.soft_bodies): + self._sim_vertex_position[i] = soft_body.get_sim_position_inv_mass_buffer()[ + :, :3 + ] + return self._sim_vertex_position.clone() + + @property + def sim_vertex_velocity(self) -> torch.Tensor: + """Return current simulation-vertex velocities.""" + for i, soft_body in enumerate(self.soft_bodies): + self._sim_vertex_velocity[i] = soft_body.get_sim_velocity_buffer()[:, :3] + return self._sim_vertex_velocity.clone() + + @property + def nodal_pos_w(self) -> torch.Tensor: + return self.sim_vertex_position + + @property + def nodal_vel_w(self) -> torch.Tensor: + return self.sim_vertex_velocity + + @property + def default_nodal_state_w(self) -> torch.Tensor: + return self._default_nodal_state_w.clone() + + @cached_property + def collision_surface_triangles(self) -> torch.Tensor: + """Return a stable convex-hull topology over collision vertices.""" + vertices = self.rest_collision_vertices[0].detach().cpu().numpy() + if vertices.shape[0] < 4: + logger.log_warning( + "Volume-deformable collision geometry has fewer than four " + "vertices; its visualization surface will be empty." + ) + triangles = np.empty((0, 3), dtype=np.int32) + else: + try: + triangles = np.asarray(ConvexHull(vertices).simplices, dtype=np.int32) + except QhullError as error: + try: + triangles = np.asarray( + ConvexHull(vertices, qhull_options="QJ").simplices, + dtype=np.int32, + ) + except QhullError: + logger.log_warning( + "Unable to build a volume-deformable visualization " + f"surface from collision vertices: {error!r}" + ) + triangles = np.empty((0, 3), dtype=np.int32) + return torch.as_tensor(triangles, dtype=torch.int32, device=self.device) + + +class VolumeDeformableObject(DeformableObject): + """A batch of DexSim volume deformables backed by ``SoftBody``.""" + + deformable_type = "volume" + spawn_kind = "soft_object" + display_name = "volume deformable" + + def _create_data( + self, + entities: Sequence[Any], + physics_scene: PhysicsScene, + device: torch.device, + ) -> VolumeDeformableData: + return VolumeDeformableData(entities, physics_scene, device) + + @property + def body_data(self) -> VolumeDeformableData | None: + """Compatibility view of the DexSim soft-body data.""" + return self._data + + def _apply_local_pose( + self, + pose: torch.Tensor, + env_ids: Sequence[int], + arena_offsets: torch.Tensor, + ) -> None: + self._require_data() + rest_collision_vertices = self.body_data.rest_collision_vertices + rest_sim_vertices = self.body_data.rest_sim_vertices + for i, env_idx in enumerate(env_ids): + soft_body: SoftBody = self._entities[env_idx].get_physical_body() + initial_transform = torch.as_tensor( + soft_body.get_initial_transform(), + dtype=torch.float32, + device=self.device, + ) + initial_rotation = initial_transform[:3, :3] + initial_translation = initial_transform[:3, 3] + rest_collision_local = ( + rest_collision_vertices[env_idx] - initial_translation + ) @ initial_rotation + rest_sim_local = ( + rest_sim_vertices[env_idx] - initial_translation + ) @ initial_rotation + rotation = pose[i, :3, :3] + translation = pose[i, :3, 3] + arena_offset = torch.as_tensor( + arena_offsets[env_idx], dtype=torch.float32, device=self.device + ) + + collision_positions = ( + rest_collision_local @ rotation.T + translation + arena_offset + ) + sim_positions = rest_sim_local @ rotation.T + translation + arena_offset + + soft_body.get_position_inv_mass_buffer()[:, :3] = collision_positions + soft_body.get_sim_position_inv_mass_buffer()[:, :3] = sim_positions + soft_body.get_sim_velocity_buffer()[:, :3] = 0.0 + soft_body.mark_dirty(SoftBodyGPUAPIReadWriteType.ALL) + soft_body.set_wake_counter(0.4) + + def get_rest_collision_vertices(self) -> torch.Tensor: + """Return rest collision vertices.""" + self._require_data() + return self.body_data.rest_collision_vertices + + def get_rest_sim_vertices(self) -> torch.Tensor: + """Return rest simulation vertices.""" + self._require_data() + return self.body_data.rest_sim_vertices + + def get_current_collision_vertices(self) -> torch.Tensor: + """Return current collision vertices.""" + self._require_data() + return self.body_data.collision_position + + def get_current_sim_vertices(self) -> torch.Tensor: + """Return current simulation vertices.""" + return self.get_current_nodal_position() + + def get_current_sim_vertex_velocities(self) -> torch.Tensor: + """Return current simulation-vertex velocities.""" + return self.get_current_nodal_velocity() + + def get_surface_vertices(self) -> torch.Tensor: + """Return the live collision surface used for visualization.""" + return self.get_current_collision_vertices() + + def get_collision_surface_triangles( + self, env_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Return convex-hull triangles over collision vertices.""" + self._require_data() + ids = self._resolve_env_ids(env_ids) + return ( + self.body_data.collision_surface_triangles.unsqueeze(0) + .expand(len(ids), -1, -1) + .clone() + ) + + def get_surface_triangles( + self, env_ids: Sequence[int] | None = None + ) -> torch.Tensor: + """Return the volume deformable's collision-surface topology.""" + return self.get_collision_surface_triangles(env_ids=env_ids) + + +# Compatibility names retained for existing environments and tutorials. +SoftBodyData = VolumeDeformableData +SoftObject = VolumeDeformableObject diff --git a/embodichain/lab/sim/objects/light.py b/embodichain/lab/sim/objects/light.py index 065267333..f497a96ae 100644 --- a/embodichain/lab/sim/objects/light.py +++ b/embodichain/lab/sim/objects/light.py @@ -46,6 +46,7 @@ def __init__( ) -> None: super().__init__(cfg, entities, device) + self.reset() def set_color( self, colors: torch.Tensor, env_ids: Sequence[int] | None = None diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index 39185a8a0..f9b3e3bbe 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -20,14 +20,23 @@ import dexsim import numpy as np -from dataclasses import dataclass, MISSING -from typing import List, Sequence, Union +from copy import deepcopy +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Sequence from functools import cached_property from dexsim.models import MeshObject from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType from dexsim.engine import CudaArray, MaterialInst, PhysicsScene -from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg +from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg, RigidObjectCfg +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.physics.newton import is_newton_gradient_mode from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim import ( VisualMaterial, @@ -45,10 +54,13 @@ get_combined_triangles, get_combined_vertices, ) -from embodichain.utils.math import convert_quat from embodichain.utils.math import matrix_from_quat, quat_from_matrix, matrix_from_euler from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedObject + +_UINT64_MAX = (1 << 64) - 1 __all__ = ["RigidBodyData", "RigidObject", "RigidObjectCfg"] @@ -56,12 +68,16 @@ 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__( - self, entities: List[MeshObject], ps: PhysicsScene, device: torch.device + self, + entities: List[MeshObject], + ps: PhysicsScene | None, + device: torch.device, + body_view: RigidBodyViewBase | None = None, ) -> None: """Initialize the RigidBodyData. @@ -75,16 +91,21 @@ 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 body_view is not None: + self.body_view = body_view + elif is_newton_scene(ps): + self.body_view: RigidBodyViewBase = NewtonRigidBodyView( + entities=entities, scene=ps, device=device ) - 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( @@ -102,77 +123,127 @@ 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) - self.default_com_pose = torch.zeros( - (self.num_instances, 7), dtype=torch.float32, device=self.device - ) + # Initialization-time physical-property snapshots. These are captured + # after backend materialization and remain unchanged by runtime writes. + self._default_mass: torch.Tensor | None = None + self._default_inertia: torch.Tensor | None = None + self._default_com_pose: torch.Tensor | None = None + + # center of mass pose in format (x, y, z, qx, qy, qz, qw) 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 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, + def default_physical_properties_initialized(self) -> bool: + """Whether the backend-resolved physical-property defaults are available.""" + return ( + self._default_mass is not None + and self._default_inertia is not None + and self._default_com_pose is not None + ) + + @property + def default_mass(self) -> torch.Tensor: + """Initialization-time mass with shape ``(N,)``.""" + if self._default_mass is None: + raise RuntimeError("Default rigid-body mass has not been captured yet.") + return self._default_mass + + @property + def default_inertia(self) -> torch.Tensor: + """Initialization-time inertia diagonal with shape ``(N, 3)``.""" + if self._default_inertia is None: + raise RuntimeError("Default rigid-body inertia has not been captured yet.") + return self._default_inertia + + @property + def default_com_pose(self) -> torch.Tensor: + """Initialization-time local COM pose as an ``xyz + xyzw`` tensor.""" + if self._default_com_pose is None: + raise RuntimeError("Default rigid-body COM pose has not been captured yet.") + return self._default_com_pose + + def capture_default_physical_properties( + self, + *, + mass: torch.Tensor, + inertia: torch.Tensor, + com_pose: torch.Tensor, + ) -> None: + """Capture backend-resolved physical properties exactly once.""" + expected_shapes = { + "mass": (self.num_instances,), + "inertia": (self.num_instances, 3), + "com_pose": (self.num_instances, 7), + } + values = {"mass": mass, "inertia": inertia, "com_pose": com_pose} + for name, value in values.items(): + if tuple(value.shape) != expected_shapes[name]: + raise ValueError( + f"Expected {name} shape {expected_shapes[name]}, got {tuple(value.shape)}." + ) + + if self.default_physical_properties_initialized: + raise RuntimeError( + "Default rigid-body physical properties are already captured." ) - 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._default_mass = mass.to(self.device, dtype=torch.float32).clone() + self._default_inertia = inertia.to(self.device, dtype=torch.float32).clone() + self._default_com_pose = com_pose.to(self.device, dtype=torch.float32).clone() + + @property + def is_newton_backend(self) -> bool: + return bool( + getattr( + self.body_view, + "is_newton_backend", + isinstance(self.body_view, NewtonRigidBodyView), ) - self._pose[:, :4] = convert_quat(self._pose[:, :4], to="wxyz") - self._pose = self._pose[:, [4, 5, 6, 0, 1, 2, 3]] - return self._pose + ) + + @property + def 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.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: @@ -185,39 +256,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: @@ -228,21 +279,33 @@ def acc(self) -> torch.Tensor: """ return torch.cat((self.lin_acc, self.ang_acc), dim=-1) + @property + def mass(self) -> torch.Tensor: + """Get current masses with shape ``(N,)``.""" + if not self.body_view.is_ready: + logger.log_error("RigidBodyData mass requested but body view is not ready.") + self.body_view.fetch_mass(self._mass) + return self._mass.squeeze(-1) + + @property + def inertia(self) -> torch.Tensor: + """Get current inertia diagonals with shape ``(N, 3)``.""" + if not self.body_view.is_ready: + logger.log_error( + "RigidBodyData inertia requested but body view is not ready." + ) + self.body_view.fetch_inertia_diagonal(self._inertia) + return self._inertia + @property def com_pose(self) -> torch.Tensor: """Get the center of mass pose of the rigid bodies. Returns: - torch.Tensor: The center of mass pose with shape (N, 7). + torch.Tensor: The center-of-mass pose with shape ``(N, 7)`` in + ``(x, y, z, qx, qy, qz, qw)`` order. """ - 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 @@ -261,36 +324,90 @@ def __init__( cfg: RigidObjectCfg, entities: List[MeshObject] = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared RigidObject requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self.body_type = cfg.body_type + self._entities = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._ps = None + self._world = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._visual_material = [None] * declared_num_instances + self.is_shared_visual_material = False + self._has_collision_visible_node = False + return + + self._declared_num_instances = len(entities) + self._spawn_result = spawn_result self.body_type = cfg.body_type - self._world = dexsim.default_world() - self._ps = self._world.get_physics_scene() + if spawn_result is None: + self._world = dexsim.default_world() + from embodichain.lab.sim.sim_manager import get_physics_scene + + self._ps = get_physics_scene() + else: + self._world = spawn_result.world + self._ps = None self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() # data for managing body data (only for dynamic and kinematic bodies) on GPU. self._data: RigidBodyData | None = None if self.is_static is False: - self._data = RigidBodyData(entities=entities, ps=self._ps, device=device) + body_view = None + if spawn_result is not None: + from embodichain.lab.sim.objects.backends import SpawnRigidBodyView + + batch = spawn_result.create_rigid_body_batch(entities) + body_view = SpawnRigidBodyView(spawn_result, batch, device) + self._data = RigidBodyData( + entities=entities, + ps=self._ps, + device=device, + body_view=body_view, + ) # For rendering purposes, each instance can have its own material. self._visual_material: List[VisualMaterialInst] = [None] * len(entities) self.is_shared_visual_material = False - # Determine if we should use USD properties or cfg properties. - if not cfg.use_usd_properties: + source_path = getattr(cfg.shape, "fpath", None) + is_usd_source = str(source_path).lower().endswith((".usd", ".usda", ".usdc")) + preserve_asset_physics = ( + is_usd_source and cfg.resolve_asset_physics_mode() == "preserve" + ) + + # Procedural/non-USD sources have no authored physics to preserve. + if spawn_result is None and not preserve_asset_physics: for entity in entities: entity.set_body_scale(*cfg.body_scale) - entity.set_physical_attr(cfg.attrs.attr()) - else: + 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.to_dexsim_physical_attr()) + elif spawn_result is None: # Read current properties from USD-loaded entities and write back to cfg # Use first entity as reference first_entity: MeshObject = entities[0] cfg.body_scale = tuple(first_entity.get_body_scale()) - cfg.attrs = RigidBodyAttributesCfg().from_dict( - first_entity.get_physical_attr().as_dict() + cfg.attrs = RigidBodyPhysicsCfg.from_dexsim_physical_attr( + first_entity.get_physical_attr() ) super().__init__(cfg, entities, device) @@ -298,36 +415,110 @@ def __init__( self._initialize_existing_visual_material() # set default collision filter - self._set_default_collision_filter() + if spawn_result is None: + self._set_default_collision_filter() - 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() + # Cache reset-relative physical properties after backend materialization. + if self._data is not None: + self._capture_default_physical_properties() # TODO: Must be called after setting all attributes. # May be improved in the future. - if cfg.attrs.enable_collision is False: + if ( + spawn_result is None + and cfg.attrs.collision_props is not None + and cfg.attrs.collision_props.collision_enabled is False + ): flag = torch.zeros(len(entities), dtype=torch.bool) self.enable_collision(flag) # reserve flag for collision visible node existence self._has_collision_visible_node = False + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to one finalized SpawnResult.""" + return self._spawn_result is not None + + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for its SpawnResult binding.""" + return self._world is None + + @property + def num_instances(self) -> int: + if self._entities: + return len(self._entities) + return self._declared_num_instances + + def attach_spawn_handles( + self, + entities: Sequence[SpawnedObject], + ) -> None: + """Store materialized handles without initializing runtime Batch data. + + Default may call this before Spawn finalization so native metadata is + available early. ``bind_spawn()`` remains responsible for creating + result-dependent Batch/Data state after finalization. + """ + handles = list(entities) + if len(handles) != self._declared_num_instances: + raise ValueError( + f"RigidObject {self.uid!r} expected " + f"{self._declared_num_instances} Spawn handles, got {len(handles)}." + ) + self._entities = handles + + def bind_spawn( + self, + result: SpawnResult, + ) -> None: + """Atomically bind a declared facade to stable Spawn handles.""" + if self.is_spawn_bound: + raise RuntimeError(f"RigidObject {self.uid!r} is already Spawn-bound.") + if not self.is_declared: + raise RuntimeError( + f"RigidObject {self.uid!r} was not created as a Spawn declaration." + ) + + cfg = self.cfg + device = self.device + entities = list(self._entities) + if len(entities) != self._declared_num_instances: + raise ValueError( + f"RigidObject {self.uid!r} expected " + f"{self._declared_num_instances} Spawn handles, got {len(entities)}." + ) + + bound = type(self)( + cfg, + entities, + device, + spawn_result=result, + ) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) + def __str__(self) -> str: - parent_str = super().__str__() - max_hull = self.cfg.max_convex_hull_num - if max_hull is MISSING: - if isinstance(self.cfg.shape, MeshCfg): - max_hull = self.cfg.shape.max_convex_hull_num - else: - max_hull = 1 + if self.is_declared: + parent_str = ( + f"{self.__class__}: declared {self.num_instances} Spawn objects " + f"| uid: {self.uid} | device: {self.device}" + ) + else: + parent_str = super().__str__() + max_hull = ( + self.cfg.shape.collision.max_hulls + if isinstance(self.cfg.shape, MeshCfg) + and self.cfg.shape.collision is not None + and self.cfg.shape.collision.max_hulls is not None + else 1 + ) return ( parent_str - + f" | body type: {self.body_type} | max_convex_hull_num: {max_hull}" + + f" | body type: {self.body_type} | collision max_hulls: {max_hull}" ) @cached_property @@ -356,12 +547,135 @@ def body_data(self) -> RigidBodyData | None: return self._data + @property + def default_mass(self) -> torch.Tensor: + """Initialization-time mass retained for backward compatibility.""" + if self._data is None: + raise RuntimeError( + "Static rigid objects do not have a default mass buffer." + ) + return self._data.default_mass + + def _capture_default_physical_properties(self) -> None: + """Capture materialized mass properties as immutable reset defaults.""" + if self._data is None or self._data.default_physical_properties_initialized: + return + if not self._data.body_view.is_ready: + logger.log_error( + "Cannot capture default rigid-body physical properties before " + "the backend view is ready." + ) + self._data.capture_default_physical_properties( + mass=self.get_mass(), + inertia=self.get_inertia(), + com_pose=self._data.com_pose, + ) + + def _restore_default_physical_properties(self, env_ids: Sequence[int]) -> None: + """Restore initialization-time mass properties for selected rows.""" + if ( + self._data is None + or not self._data.default_physical_properties_initialized + or self.is_non_dynamic + or len(env_ids) == 0 + ): + return + + index = torch.as_tensor(env_ids, dtype=torch.long, device=self.device) + self.set_mass(self._data.default_mass[index], env_ids=env_ids) + self.set_inertia(self._data.default_inertia[index], env_ids=env_ids) + self.set_com_pose(self._data.default_com_pose[index], env_ids=env_ids) + + 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 created + from grouped Spawn descriptors may not carry a legacy ``attr`` mirror. + Used by not-ready setter paths to tolerate that representation. + """ + 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. @@ -425,6 +739,26 @@ def set_collision_filter( f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." ) + if self.is_spawn_bound: + if self._data is None: + raise NotImplementedError( + "Runtime collision-filter updates are unavailable for static " + "Spawn rigid objects." + ) + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_collision_filter(filter_data, body_ids) + return + + 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( @@ -447,50 +781,45 @@ 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 = 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] - 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(target_pose[:, 3:7]) + 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`. @@ -514,7 +843,6 @@ def get_local_pose_cpu( np.array([entity.get_rotation_quat() for entity in entities]), dtype=torch.float32, ) - quats = convert_quat(quats, to="wxyz") pose = torch.cat((xyzs, quats), dim=-1) return pose @@ -522,7 +850,7 @@ 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]) @@ -580,28 +908,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 " + "unprepared; call SimulationManager.prepare() first." + ) + else: + logger.log_error("Cannot apply force or torque before body view is ready.") def set_velocity( self, @@ -638,57 +976,142 @@ 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 unprepared; " + "call SimulationManager.prepare() first." + ) + else: + logger.log_error("Cannot set velocity before body view is ready.") def set_attrs( self, - attrs: Union[RigidBodyAttributesCfg, List[RigidBodyAttributesCfg]], + attrs: RigidBodyPhysicsCfg | list[RigidBodyPhysicsCfg], env_ids: Sequence[int] | None = None, ) -> None: """Set physical attributes for the rigid object. Args: - attrs (Union[RigidBodyAttributesCfg, List[RigidBodyAttributesCfg]]): The physical attributes to set. + attrs: Grouped physical attributes, shared or one per environment. env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. """ local_env_ids = self._all_indices if env_ids is None else env_ids - if isinstance(attrs, List) and len(local_env_ids) != len(attrs): + if isinstance(attrs, list) and len(local_env_ids) != len(attrs): logger.log_error( f"Length of env_ids {len(local_env_ids)} does not match attrs length {len(attrs)}." ) - # TODO: maybe need to improve the physical attributes setter efficiency. - if isinstance(attrs, RigidBodyAttributesCfg): - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].set_physical_attr(attrs.attr()) + # Resolve per-env physical attrs into a flat list aligned with local_env_ids. + if isinstance(attrs, RigidBodyPhysicsCfg): + physical_attrs = [attrs.to_dexsim_physical_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.to_dexsim_physical_attr() for a in attrs] + + if self.is_spawn_bound: + if self._data is None: + raise NotImplementedError( + "Runtime physical attributes are unavailable for static " + "Spawn rigid objects." + ) + body_ids = self._data.body_ids_for(local_env_ids) + view = self._data.body_view + + def _stack(field: str) -> torch.Tensor: + return torch.as_tensor( + [getattr(attr, field) for attr in physical_attrs], + dtype=torch.float32, + device=self.device, + ).unsqueeze(-1) + + if any( + attr.static_friction != attr.dynamic_friction for attr in physical_attrs + ): + logger.log_warning( + "DexSim Spawn exposes one backend-neutral friction value; " + "set_attrs() uses dynamic_friction for both coefficients." + ) + 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) + view.apply_damping( + torch.cat( + (_stack("linear_damping"), _stack("angular_damping")), + dim=1, + ), + body_ids, + ) + return + + 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 prepared; physical attributes are mirrored " + "to metadata and applied at the next prepare()." + ) + 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 @@ -706,9 +1129,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 + # Default-backend 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. @@ -721,9 +1159,38 @@ 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.is_spawn_bound and self.is_static: + # Static actors have no finite runtime mass (and Newton therefore + # gives them no body id), but the legacy API exposed their authored + # configuration. Preserve that readable metadata contract without + # manufacturing a dynamic-body batch solely for property queries. + configured_mass = self.cfg.attrs.to_dexsim_physical_attr().mass + value = 0.0 if configured_mass is None else float(configured_mass) + return torch.full( + (len(local_env_ids),), + value, + dtype=torch.float32, + device=self.device, + ) + + if self._data is not None and self._data.body_view.is_ready: + if env_ids is None: + return self._data.mass + body_ids = self._data.body_ids_for(local_env_ids) + buf = torch.empty( + (len(local_env_ids), 1), + dtype=torch.float32, + device=self.device, + ) + 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) @@ -744,12 +1211,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 Default-backend 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. @@ -762,11 +1247,28 @@ 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.is_spawn_bound and self.is_static: + return torch.full( + (len(local_env_ids),), + float(self.cfg.attrs.to_dexsim_physical_attr().dynamic_friction), + dtype=torch.float32, + device=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) + 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) @@ -779,6 +1281,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 @@ -787,13 +1295,31 @@ 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 self.is_spawn_bound: + if self._data is None: + raise NotImplementedError( + "Runtime damping is unavailable for static Spawn rigid objects." + ) + body_ids = self._data.body_ids_for(local_env_ids) + self._data.body_view.apply_damping(damping, body_ids) + return + + 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: @@ -807,14 +1333,38 @@ def get_damping(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: """ local_env_ids = self._all_indices if env_ids is None else env_ids + if self.is_spawn_bound: + if self._data is None: + return torch.tensor( + [ + self.cfg.attrs.to_dexsim_physical_attr().linear_damping, + self.cfg.attrs.to_dexsim_physical_attr().angular_damping, + ], + dtype=torch.float32, + device=self.device, + ).repeat(len(local_env_ids), 1) + body_ids = self._data.body_ids_for(local_env_ids) + damping = torch.empty( + (len(local_env_ids), 2), + dtype=torch.float32, + device=self.device, + ) + self._data.body_view.fetch_damping(damping, body_ids) + return damping + 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) @@ -835,11 +1385,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 + # Default-backend 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. @@ -852,13 +1417,37 @@ 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.is_spawn_bound and self.is_static: + # Static actors have infinite mass, so no finite inertia tensor is + # represented by either Spawn backend. + return torch.zeros( + (len(local_env_ids), 3), + dtype=torch.float32, + device=self.device, + ) + + if self._data is not None and self._data.body_view.is_ready: + if env_ids is None: + return self._data.inertia + body_ids = self._data.body_ids_for(local_env_ids) + buf = torch.empty( + (len(local_env_ids), 3), + dtype=torch.float32, + device=self.device, + ) + 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( @@ -1109,7 +1698,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). @@ -1128,11 +1717,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. @@ -1142,9 +1733,28 @@ 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 self.is_spawn_bound: + raise NotImplementedError( + "Changing actor topology after Spawn binding requires a public " + "descriptor mutation transaction and is not implemented yet." + ) + + if is_newton_scene(self._ps): + logger.log_warning( + "Newton backend does not support changing RigidObject body type at " + "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'." @@ -1254,36 +1864,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 unprepared; " + "call SimulationManager.prepare() first." ) + else: + logger.log_error("Cannot clear dynamics before body view is ready.") def set_physical_visible( self, @@ -1300,6 +1903,13 @@ def set_physical_visible( if len(rgba) != 4: logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.") + if self.is_spawn_bound: + color = np.asarray(rgba, dtype=np.float32) + for entity in self._entities: + self._spawn_result.set_physical_visible(entity, color, visible) + self._has_collision_visible_node = True + return + # create collision visible node if not exist if visible: if not self._has_collision_visible_node: @@ -1329,14 +1939,19 @@ 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.restore_visual_material(env_ids=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) + if self.cfg.init_local_pose is not None: + return ( + torch.as_tensor( + self.cfg.init_local_pose, + dtype=torch.float32, + device=self.device, + ) + .reshape(1, 4, 4) + .repeat(num_instances, 1, 1) + ) pos = torch.as_tensor( self.cfg.init_pos, dtype=torch.float32, device=self.device ) @@ -1355,14 +1970,70 @@ 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. + + The Default backend runs a full reset. Newton applies init pose in + ``BUILDER`` via the scene batch API; velocities are cleared after + preparation through :meth:`SimulationManager.prepare`. + """ + if self.is_spawn_bound: + if self._spawn_result.backend == "dexsim": + # DexSim Direct GPU readiness performs native warm-up updates. + # Re-apply the authored state after the batch becomes usable + # so prepare() itself is not an observable simulation step. + self.reset() + else: + # Newton finalization materializes the descriptor pose without + # advancing simulation; only one-step dynamics buffers need + # clearing after batch binding. + if not is_newton_gradient_mode(self._spawn_result): + self.clear_dynamics() + return + + 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 + + self.restore_visual_material(env_ids=local_env_ids) + + # Preserve the legacy Default-backend attribute reset before restoring + # the backend-resolved mass-property snapshot below. + if not self.is_spawn_bound and not is_newton_scene(self._ps): + self.set_attrs(self.cfg.attrs, env_ids=local_env_ids) + + self._restore_default_physical_properties(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: + if self.is_declared or self.is_spawn_bound: + # SimulationManager owns topology removal and SpawnResult lifetime. + # Direct facade destruction must never bypass that owner. + return env = self._world.get_env() arenas = env.get_all_arenas() if len(arenas) == 0: arenas = [env] for i, entity in enumerate(self._entities): - arenas[i].remove_actor(entity) + if is_newton_scene(self._ps): + arenas[i].remove_actor(entity.get_name()) + else: + arenas[i].remove_actor(entity) diff --git a/embodichain/lab/sim/objects/rigid_object_group.py b/embodichain/lab/sim/objects/rigid_object_group.py index 0f6192d28..2d035f2e3 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/embodichain/lab/sim/objects/rigid_object_group.py @@ -16,249 +16,302 @@ from __future__ import annotations -import torch -import dexsim -import numpy as np +from copy import deepcopy +from typing import TYPE_CHECKING, Sequence -from dataclasses import dataclass -from typing import List, Sequence, Union +import numpy as np +import torch -from dexsim.models import MeshObject -from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType -from dexsim.engine import CudaArray, PhysicsScene -from embodichain.lab.sim.cfg import ( - RigidObjectGroupCfg, - RigidBodyAttributesCfg, -) -from embodichain.lab.sim import ( - BatchEntity, -) -from embodichain.lab.sim.material import VisualMaterial, VisualMaterialInst -from ._mesh_utils import ( - get_combined_triangles, - get_combined_vertices, +from embodichain.lab.sim import BatchEntity +from embodichain.lab.sim.cfg import RigidObjectGroupCfg +from embodichain.lab.sim.material import VisualMaterial +from embodichain.lab.sim.objects.backends.spawn import SpawnRigidBodyView +from embodichain.utils.math import ( + matrix_from_euler, + matrix_from_quat, + quat_from_matrix, ) -from embodichain.utils.math import convert_quat -from embodichain.utils.math import matrix_from_quat, quat_from_matrix, matrix_from_euler -from embodichain.utils import logger + +from ._mesh_utils import get_combined_triangles, get_combined_vertices + +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedObject __all__ = ["RigidBodyGroupData", "RigidObjectGroup", "RigidObjectGroupCfg"] -@dataclass class RigidBodyGroupData: - """Data manager for rigid body group with body type of dynamic or kinematic.""" + """Expose one flat Spawn rigid-body batch as ``[env, object, ...]`` tensors.""" def __init__( - self, entities: List[List[MeshObject]], ps: PhysicsScene, device: torch.device + self, + body_view: SpawnRigidBodyView, + *, + num_instances: int, + num_objects: int, + device: torch.device, ) -> None: - """Initialize the RigidBodyGroupData. - - Args: - entities (List[List[MeshObject]]): List of List MeshObjects representing the rigid body group. - ps (PhysicsScene): The physics scene. - device (torch.device): The device to use for the rigid body group data. - """ - self.entities = entities - self.ps = ps - self.num_instances = len(entities) - self.num_objects = len(entities[0]) + self.body_view = body_view + self.num_instances = num_instances + self.num_objects = num_objects self.device = device - - # get gpu indices for the rigid bodies with shape of (num_instances, num_objects) - self.gpu_indices = ( - torch.as_tensor( - [ - [entity.get_gpu_index() for entity in instance] - for instance in entities - ], - dtype=torch.int32, - device=self.device, - ) - if self.device.type == "cuda" - else None + self._pose = torch.empty( + (num_instances, num_objects, 7), dtype=torch.float32, device=device ) - - # Initialize rigid body group data tensors. Shape of (num_instances, num_objects, data_dim) - self._pose = torch.zeros( - (self.num_instances, self.num_objects, 7), + self._lin_vel = torch.empty( + (num_instances, num_objects, 3), dtype=torch.float32, device=device + ) + self._ang_vel = torch.empty_like(self._lin_vel) + self._mass = torch.empty( + (num_instances, num_objects, 1), dtype=torch.float32, - device=self.device, + device=device, ) - self._lin_vel = torch.zeros( - (self.num_instances, self.num_objects, 3), + self._inertia = torch.empty( + (num_instances, num_objects, 3), dtype=torch.float32, - device=self.device, + device=device, ) - self._ang_vel = torch.zeros( - (self.num_instances, self.num_objects, 3), + self._com_pose = torch.empty( + (num_instances, num_objects, 7), dtype=torch.float32, - device=self.device, + device=device, ) + self._default_mass: torch.Tensor | None = None + self._default_inertia: torch.Tensor | None = None + self._default_com_pose: torch.Tensor | None = None @property def pose(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch pose from CPU entities - xyzs = torch.as_tensor( - [ - [entity.get_location() for entity in instance] - for instance in self.entities - ], - device=self.device, - ) - quats = torch.as_tensor( - [ - [entity.get_rotation_quat() for entity in instance] - for instance in self.entities - ], - device=self.device, - ) - quats = convert_quat(quats.reshape(-1, 4), to="wxyz").reshape( - -1, self.num_objects, 4 - ) - return torch.cat((xyzs, quats), dim=-1) - else: - pose = self._pose.reshape(-1, 7) - self.ps.gpu_fetch_rigid_body_data( - data=pose, - gpu_indices=self.gpu_indices.flatten(), - data_type=RigidBodyGPUAPIReadType.POSE, - ) - pose = convert_quat(pose[:, :4], to="wxyz") - pose = pose[:, [4, 5, 6, 0, 1, 2, 3]] - return self._pose + """Local poses in EmbodiChain ``xyz + xyzw`` order.""" + flat = self._pose.reshape(-1, 7) + self.body_view.fetch_pose(flat) + return self._pose @property def lin_vel(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch linear velocity from CPU entities - self._lin_vel = torch.as_tensor( - [ - [entity.get_linear_velocity() for entity in instance] - for instance in self.entities - ], - dtype=torch.float32, - device=self.device, - ) - else: - lin_vel = self._lin_vel.reshape(-1, 3) - self.ps.gpu_fetch_rigid_body_data( - data=lin_vel, - gpu_indices=self.gpu_indices.flatten(), - data_type=RigidBodyGPUAPIReadType.LINEAR_VELOCITY, - ) + self.body_view.fetch_linear_velocity(self._lin_vel.reshape(-1, 3)) return self._lin_vel @property def ang_vel(self) -> torch.Tensor: - if self.device.type == "cpu": - # Fetch angular velocity from CPU entities - self._ang_vel = torch.as_tensor( - [ - [entity.get_angular_velocity() for entity in instance] - for instance in self.entities - ], - dtype=torch.float32, - device=self.device, - ) - else: - ang_vel = self._ang_vel.reshape(-1, 3) - self.ps.gpu_fetch_rigid_body_data( - data=ang_vel, - gpu_indices=self.gpu_indices.flatten(), - data_type=RigidBodyGPUAPIReadType.ANGULAR_VELOCITY, - ) + self.body_view.fetch_angular_velocity(self._ang_vel.reshape(-1, 3)) return self._ang_vel @property def vel(self) -> torch.Tensor: - """Get the linear and angular velocities of the rigid bodies. - - Returns: - torch.Tensor: The linear and angular velocities concatenated, with shape (num_instances, num_objects, 6). - """ + """Linear and angular velocities with shape ``[env, object, 6]``.""" return torch.cat((self.lin_vel, self.ang_vel), dim=-1) + @property + def mass(self) -> torch.Tensor: + """Current masses with shape ``[env, object]``.""" + self.body_view.fetch_mass(self._mass.reshape(-1, 1)) + return self._mass.squeeze(-1) + + @property + def inertia(self) -> torch.Tensor: + """Current inertia diagonals with shape ``[env, object, 3]``.""" + self.body_view.fetch_inertia_diagonal(self._inertia.reshape(-1, 3)) + return self._inertia + + @property + def com_pose(self) -> torch.Tensor: + """Current local COM poses in Group ``xyz + xyzw`` convention.""" + flat = self._com_pose.reshape(-1, 7) + self.body_view.fetch_com_local_pose(flat) + return self._com_pose + + @property + def default_physical_properties_initialized(self) -> bool: + """Whether initialization-time mass properties are available.""" + return ( + self._default_mass is not None + and self._default_inertia is not None + and self._default_com_pose is not None + ) + + @property + def default_mass(self) -> torch.Tensor: + """Initialization-time masses with shape ``[env, object]``.""" + if self._default_mass is None: + raise RuntimeError("Default rigid-object Group masses are unavailable.") + return self._default_mass + + @property + def default_inertia(self) -> torch.Tensor: + """Initialization-time inertia diagonals.""" + if self._default_inertia is None: + raise RuntimeError("Default rigid-object Group inertias are unavailable.") + return self._default_inertia + + @property + def default_com_pose(self) -> torch.Tensor: + """Initialization-time local COM poses in ``xyz + xyzw`` order.""" + if self._default_com_pose is None: + raise RuntimeError("Default rigid-object Group COM poses are unavailable.") + return self._default_com_pose + + def capture_default_physical_properties( + self, + *, + mass: torch.Tensor, + inertia: torch.Tensor, + com_pose: torch.Tensor, + ) -> None: + """Capture backend-resolved Group mass properties exactly once.""" + expected_shapes = { + "mass": (self.num_instances, self.num_objects), + "inertia": (self.num_instances, self.num_objects, 3), + "com_pose": (self.num_instances, self.num_objects, 7), + } + values = {"mass": mass, "inertia": inertia, "com_pose": com_pose} + for name, value in values.items(): + if tuple(value.shape) != expected_shapes[name]: + raise ValueError( + f"Expected {name} shape {expected_shapes[name]}, " + f"got {tuple(value.shape)}." + ) + if self.default_physical_properties_initialized: + raise RuntimeError( + "Default rigid-object Group mass properties are already captured." + ) + + self._default_mass = mass.to(self.device, dtype=torch.float32).clone() + self._default_inertia = inertia.to(self.device, dtype=torch.float32).clone() + self._default_com_pose = com_pose.to(self.device, dtype=torch.float32).clone() + class RigidObjectGroup(BatchEntity): - """RigidObjectGroup represents a batch of rigid bodies in the simulation.""" + """A two-dimensional view over rigid objects owned by DexSim Spawn.""" def __init__( self, cfg: RigidObjectGroupCfg, - entities: List[List[MeshObject]] = None, + entities: Sequence[Sequence[SpawnedObject]] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: self.body_type = cfg.body_type + self._declared_num_objects = len(cfg.rigid_objects) - self._world = dexsim.default_world() - self._ps = self._world.get_physics_scene() - - self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() - self._all_obj_indices = torch.arange( - len(entities[0]), dtype=torch.int32 - ).tolist() - - # data for managing body data (only for dynamic and kinematic bodies) on GPU. - self._data = RigidBodyGroupData(entities=entities, ps=self._ps, device=device) + if entities is None: + if declared_num_instances is None or declared_num_instances <= 0: + raise ValueError( + "A declared RigidObjectGroup requires declared_num_instances > 0." + ) + self.cfg = deepcopy(cfg) + self.uid = self.cfg.uid + self.device = device + self._entities: list[list[SpawnedObject]] = [] + self._declared_num_instances = declared_num_instances + self._spawn_result = None + self._data = None + self._all_indices = list(range(declared_num_instances)) + self._all_obj_indices = list(range(self._declared_num_objects)) + return - body_cfgs = list(cfg.rigid_objects.values()) - for instance in entities: - for i, body in enumerate(instance): - body.set_body_scale(*body_cfgs[i].body_scale) - body.set_physical_attr(body_cfgs[i].attrs.attr()) + rows = [list(instance) for instance in entities] + if not rows or any( + len(instance) != self._declared_num_objects for instance in rows + ): + raise ValueError( + "RigidObjectGroup Spawn handles must have shape " + "[num_instances, num_objects]." + ) + if spawn_result is None: + raise ValueError( + "RigidObjectGroup entities must be owned by a SpawnResult." + ) - if device.type == "cuda": - self._world.update(0.001) + self._declared_num_instances = len(rows) + self._spawn_result = spawn_result + self._all_indices = list(range(len(rows))) + self._all_obj_indices = list(range(self._declared_num_objects)) + flat_entities = [entity for instance in rows for entity in instance] + batch = spawn_result.create_rigid_body_batch(flat_entities) + body_view = SpawnRigidBodyView(spawn_result, batch, device) + self._data = RigidBodyGroupData( + body_view, + num_instances=len(rows), + num_objects=self._declared_num_objects, + device=device, + ) - super().__init__(cfg, entities, device) + super().__init__(cfg, rows, device) + self._capture_default_physical_properties() + self.reset() - # set default collision filter - self._set_default_collision_filter() + @property + def is_declared(self) -> bool: + """Whether this facade is waiting for Spawn materialization.""" + return self._spawn_result is None - # reserve flag for collision visible node existence - n_instances = len(self._entities[0]) - self._has_collision_visible_node_list = [False] * n_instances + @property + def is_spawn_bound(self) -> bool: + """Whether this facade is bound to a SpawnResult.""" + return self._spawn_result is not None - def __str__(self) -> str: - parent_str = super().__str__() - return ( - parent_str - + f" | body type: {self.body_type} | num_objects: {self.num_objects}" - ) + @property + def num_instances(self) -> int: + return len(self._entities) if self._entities else self._declared_num_instances @property def num_objects(self) -> int: - """Get the number of objects in each rigid body instance. - - Returns: - int: The number of objects in each rigid body instance. - """ - return self._data.num_objects + return self._declared_num_objects @property def body_data(self) -> RigidBodyGroupData: - """Get the rigid body data manager for this rigid object. - - Returns: - RigidBodyGroupData: The rigid body data manager. - """ + if self._data is None: + raise RuntimeError( + f"RigidObjectGroup {self.uid!r} is not bound; call SimulationManager.prepare()." + ) return self._data - @property - def body_state(self) -> torch.Tensor: - """Get the body state of the rigid object. - - The body state of a rigid object is represented as a tensor with the following format: - [x, y, z, qw, qx, qy, qz, lin_x, lin_y, lin_z, ang_x, ang_y, ang_z] + def _capture_default_physical_properties(self) -> None: + """Capture materialized Group mass properties as reset defaults.""" + data = self.body_data + if data.default_physical_properties_initialized: + return + data.capture_default_physical_properties( + mass=data.mass, + inertia=data.inertia, + com_pose=data.com_pose, + ) - If the rigid object is static, linear and angular velocities will be zero. + def _restore_default_physical_properties( + self, env_ids: Sequence[int] | torch.Tensor | None + ) -> None: + """Restore initialization-time Group mass properties for selected rows.""" + data = self.body_data + if self.is_non_dynamic or not data.default_physical_properties_initialized: + return + env, objects, _ = self._selected_indices(env_ids) + if not env: + return + env_index = torch.as_tensor(env, dtype=torch.long, device=self.device) + obj_index = torch.as_tensor(objects, dtype=torch.long, device=self.device) + self.set_mass( + data.default_mass[env_index[:, None], obj_index[None, :]], + env_ids=env, + obj_ids=objects, + ) + self.set_inertia( + data.default_inertia[env_index[:, None], obj_index[None, :]], + env_ids=env, + obj_ids=objects, + ) + self.set_com_pose( + data.default_com_pose[env_index[:, None], obj_index[None, :]], + env_ids=env, + obj_ids=objects, + ) - Returns: - torch.Tensor: The body state of the rigid object with shape (num_instances, num_objects, 13), - where N is the number of instances. - """ + @property + def body_state(self) -> torch.Tensor: + """Pose and velocity with shape ``[env, object, 13]``.""" return torch.cat( (self.body_data.pose, self.body_data.lin_vel, self.body_data.ang_vel), dim=-1, @@ -266,46 +319,192 @@ def body_state(self) -> torch.Tensor: @property def is_non_dynamic(self) -> bool: - """Check if the rigid object is non-dynamic (static or kinematic). + return self.body_type in ("static", "kinematic") - Returns: - bool: True if the rigid object is non-dynamic, False otherwise. + def attach_spawn_handles(self, entities: Sequence[SpawnedObject]) -> None: + """Store env-major handles without initializing the group's Batch data. + + ``bind_spawn()`` creates the result-dependent runtime view after Spawn + finalization. """ - return self.body_type in ("static", "kinematic") + expected = self._declared_num_instances * self.num_objects + if len(entities) != expected: + raise ValueError( + f"RigidObjectGroup {self.uid!r} expected {expected} Spawn handles, " + f"got {len(entities)}." + ) + self._entities = [ + list(entities[start : start + self.num_objects]) + for start in range(0, len(entities), self.num_objects) + ] + + def bind_spawn(self, result: SpawnResult) -> None: + """Atomically bind the declaration facade to env-major Spawn handles.""" + if self.is_spawn_bound: + raise RuntimeError(f"RigidObjectGroup {self.uid!r} is already Spawn-bound.") + if not self.is_declared: + raise RuntimeError( + f"RigidObjectGroup {self.uid!r} was not created as a Spawn declaration." + ) + + cfg = self.cfg + device = self.device + rows = [list(row) for row in self._entities] + if len(rows) != self._declared_num_instances or any( + len(row) != self.num_objects for row in rows + ): + raise ValueError( + f"RigidObjectGroup {self.uid!r} expected " + f"{self._declared_num_instances}x{self.num_objects} Spawn handles." + ) - def _set_default_collision_filter(self) -> None: - collision_filter_data = torch.zeros( - size=(self.num_instances, 4), dtype=torch.int32 + bound = type(self)( + cfg, + rows, + device, + spawn_result=result, ) - for i in range(self.num_instances): - collision_filter_data[i, 0] = i - collision_filter_data[i, 1] = 1 - self.set_collision_filter(collision_filter_data) + self.__dict__.clear() + self.__dict__.update(bound.__dict__) - def set_collision_filter( - self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None + def __str__(self) -> str: + if self.is_declared: + return ( + f"{self.__class__}: declared {self.num_instances}x{self.num_objects} " + f"Spawn objects | uid: {self.uid} | device: {self.device}" + ) + return ( + super().__str__() + + f" | body type: {self.body_type} | num_objects: {self.num_objects}" + ) + + def _selected_indices( + self, + env_ids: Sequence[int] | torch.Tensor | None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> tuple[list[int], list[int], torch.Tensor]: + env = ( + self._all_indices + if env_ids is None + else torch.as_tensor(env_ids).reshape(-1).cpu().tolist() + ) + objects = ( + self._all_obj_indices + if obj_ids is None + else torch.as_tensor(obj_ids).reshape(-1).cpu().tolist() + ) + if any(index < 0 or index >= self.num_instances for index in env): + raise IndexError("RigidObjectGroup environment index is out of range.") + if any(index < 0 or index >= self.num_objects for index in objects): + raise IndexError("RigidObjectGroup object index is out of range.") + rows = torch.as_tensor( + [ + env_id * self.num_objects + obj_id + for env_id in env + for obj_id in objects + ], + dtype=torch.long, + device=self.device, + ) + return env, objects, rows + + def get_mass( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Return selected masses with shape ``[env, object]``.""" + env, objects, _ = self._selected_indices(env_ids, obj_ids) + env_index = torch.as_tensor(env, dtype=torch.long, device=self.device) + obj_index = torch.as_tensor(objects, dtype=torch.long, device=self.device) + return self.body_data.mass[env_index[:, None], obj_index[None, :]] + + def set_mass( + self, + mass: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | None = None, + obj_ids: Sequence[int] | torch.Tensor | None = None, ) -> None: - """set collision filter data for the rigid object group. + """Set selected masses from a tensor shaped ``[env, object]``.""" + env, objects, rows = self._selected_indices(env_ids, obj_ids) + mass = torch.as_tensor(mass, dtype=torch.float32, device=self.device) + expected_shape = (len(env), len(objects)) + if tuple(mass.shape) != expected_shape: + raise ValueError( + f"Expected mass shape {expected_shape}, got {tuple(mass.shape)}." + ) + self.body_data.body_view.apply_mass(mass.reshape(-1, 1), rows) + + def get_inertia( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Return selected inertia diagonals with shape ``[env, object, 3]``.""" + env, objects, _ = self._selected_indices(env_ids, obj_ids) + env_index = torch.as_tensor(env, dtype=torch.long, device=self.device) + obj_index = torch.as_tensor(objects, dtype=torch.long, device=self.device) + return self.body_data.inertia[env_index[:, None], obj_index[None, :]] - Args: - filter_data (torch.Tensor): [N, 4] of int. - First element of each object is arena id. - If 2nd element is 0, the object will collision with all other objects in world. - 3rd and 4th elements are not used currently. + def set_inertia( + self, + inertia: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | None = None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set selected inertia diagonals.""" + env, objects, rows = self._selected_indices(env_ids, obj_ids) + inertia = torch.as_tensor(inertia, dtype=torch.float32, device=self.device) + expected_shape = (len(env), len(objects), 3) + if tuple(inertia.shape) != expected_shape: + raise ValueError( + f"Expected inertia shape {expected_shape}, " + f"got {tuple(inertia.shape)}." + ) + self.body_data.body_view.apply_inertia_diagonal(inertia.reshape(-1, 3), rows) - env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. Defaults to None. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids + def get_com_pose( + self, + env_ids: Sequence[int] | torch.Tensor | None = None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> torch.Tensor: + """Return selected local COM poses in Group ``xyz + xyzw`` order.""" + env, objects, _ = self._selected_indices(env_ids, obj_ids) + env_index = torch.as_tensor(env, dtype=torch.long, device=self.device) + obj_index = torch.as_tensor(objects, dtype=torch.long, device=self.device) + return self.body_data.com_pose[env_index[:, None], obj_index[None, :]] - if len(local_env_ids) != len(filter_data): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." + def set_com_pose( + self, + com_pose: torch.Tensor, + env_ids: Sequence[int] | torch.Tensor | None = None, + obj_ids: Sequence[int] | torch.Tensor | None = None, + ) -> None: + """Set selected local COM poses in Group ``xyz + xyzw`` order.""" + env, objects, rows = self._selected_indices(env_ids, obj_ids) + com_pose = torch.as_tensor(com_pose, dtype=torch.float32, device=self.device) + expected_shape = (len(env), len(objects), 7) + if tuple(com_pose.shape) != expected_shape: + raise ValueError( + f"Expected COM pose shape {expected_shape}, " + f"got {tuple(com_pose.shape)}." ) + self.body_data.body_view.apply_com_local_pose(com_pose.reshape(-1, 7), rows) - filter_data_np = filter_data.cpu().numpy().astype(np.uint32) - for i, env_idx in enumerate(local_env_ids): - for entity in self._entities[env_idx]: - entity.get_physical_body().set_collision_filter_data(filter_data_np[i]) + def set_collision_filter( + self, + filter_data: torch.Tensor, + env_ids: Sequence[int] | None = None, + ) -> None: + """Set one collision filter value for every selected member in each env.""" + env, objects, rows = self._selected_indices(env_ids) + values = filter_data.to(device=self.device, dtype=torch.int32).reshape(-1, 4) + if len(values) != len(env): + raise ValueError( + f"Expected {len(env)} collision filters, got {len(values)}." + ) + expanded = values[:, None, :].expand(-1, len(objects), -1).reshape(-1, 4) + self.body_data.body_view.apply_collision_filter(expanded, rows) def set_local_pose( self, @@ -313,96 +512,40 @@ def set_local_pose( env_ids: Sequence[int] | None = None, obj_ids: Sequence[int] | None = None, ) -> None: - """Set local pose of the rigid object group. - - Args: - pose (torch.Tensor): The local pose of the rigid object group with shape (num_instances, num_objects, 7) or - (num_instances, num_objects, 4, 4). - env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. - obj_ids (Sequence[int] | None, optional): Object indices within the group. If None, all objects are set. Defaults to None. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - local_obj_ids = self._all_obj_indices if obj_ids is None else obj_ids - - if len(local_env_ids) != len(pose): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." + """Set Group poses in ``xyz+xyzw`` or homogeneous-matrix form.""" + env, objects, rows = self._selected_indices(env_ids, obj_ids) + expected_prefix = (len(env), len(objects)) + pose = pose.to(device=self.device, dtype=torch.float32) + if tuple(pose.shape) == (*expected_prefix, 7): + target = pose.reshape(-1, 7) + elif tuple(pose.shape) == (*expected_prefix, 4, 4): + flat = pose.reshape(-1, 4, 4) + target = torch.cat( + ( + flat[:, :3, 3], + quat_from_matrix(flat[:, :3, :3]), + ), + dim=-1, ) - - if self.device.type == "cpu": - pose = pose.cpu() - if pose.dim() == 3 and pose.shape[2] == 7: - reshape_pose = pose.reshape(-1, 7) - pose_matrix = ( - torch.eye(4).unsqueeze(0).repeat(reshape_pose.shape[0], 1, 1) - ) - pose_matrix[:, :3, 3] = reshape_pose[:, :3] - pose_matrix[:, :3, :3] = matrix_from_quat(reshape_pose[:, 3:7]) - pose = pose_matrix.reshape(-1, len(local_obj_ids), 4, 4) - elif pose.dim() == 4 and pose.shape[2:] == (4, 4): - pass - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (num_instances, num_objects, 7) or (num_instances, num_objects, 4, 4)." - ) - - for i, env_idx in enumerate(local_env_ids): - for j, obj_idx in enumerate(local_obj_ids): - self._entities[env_idx][obj_idx].set_local_pose(pose[i, j]) - else: - if pose.dim() == 3 and pose.shape[2] == 7: - xyz = pose[..., :3].reshape(-1, 3) - quat = pose[..., 3:7].reshape(-1, 4) - quat = convert_quat(quat, to="xyzw") - elif pose.dim() == 4 and pose.shape[2:] == (4, 4): - xyz = pose[..., :3, 3].reshape(-1, 3) - mat = pose[..., :3, :3].reshape(-1, 3, 3) - quat = quat_from_matrix(mat) - quat = convert_quat(quat, to="xyzw") - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - # we should keep `pose_` life cycle to the end of the function. - pose = torch.cat((quat, xyz), dim=-1) - indices = self.body_data.gpu_indices[local_env_ids][ - :, local_obj_ids - ].flatten() - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=pose.clone(), - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.POSE, - ) - self._world.sync_poses_gpu_to_cpu( - rigid_pose=CudaArray(pose), rigid_gpu_indices=CudaArray(indices) + raise ValueError( + f"Expected pose shape {(*expected_prefix, 7)} or " + f"{(*expected_prefix, 4, 4)}, got {tuple(pose.shape)}." ) + self.body_data.body_view.apply_pose(target, rows) def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: - """Get local pose of the rigid object group. - - Args: - to_matrix (bool, optional): If True, return the pose as a 4x4 matrix. If False, return as (x, y, z, qw, qx, qy, qz). Defaults to False. - - Returns: - torch.Tensor: The local pose of the rigid object with shape (num_instances, num_objects, 7) or (num_instances, num_objects, 4, 4) depending on `to_matrix`. - """ + """Return all Group poses as ``xyz+xyzw`` or homogeneous matrices.""" pose = self.body_data.pose - if to_matrix: - pose = pose.reshape(-1, 7) - xyz = pose[:, :3] - mat = matrix_from_quat(pose[:, 3:7]) - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze(0) - .repeat(self.num_instances * self.num_objects, 1, 1) - ) - pose[:, :3, 3] = xyz - pose[:, :3, :3] = mat - pose = pose.reshape(self.num_instances, self.num_objects, 4, 4) - return pose + if not to_matrix: + return pose + flat = pose.reshape(-1, 7) + result = torch.eye(4, dtype=torch.float32, device=self.device).repeat( + len(flat), 1, 1 + ) + result[:, :3, 3] = flat[:, :3] + result[:, :3, :3] = matrix_from_quat(flat[:, 3:7]) + return result.reshape(self.num_instances, self.num_objects, 4, 4) def get_object_vertices( self, @@ -410,34 +553,19 @@ def get_object_vertices( env_ids: Sequence[int] | None = None, scale: bool = False, ) -> torch.Tensor: - """Get one constituent object's vertices across selected environments. - - Args: - object_id: Constituent object index within the group. - env_ids: Environment indices. If ``None``, returns all instances. - scale: Whether to apply each object's body scale. - - Returns: - Vertices with shape ``(N, num_vertices, 3)``. - """ - if not 0 <= object_id < self.num_objects: - raise IndexError( - f"object_id {object_id} is outside [0, {self.num_objects - 1}]." - ) - ids = self._all_indices if env_ids is None else env_ids + """Return one member's render vertices across selected environments.""" + env, objects, _ = self._selected_indices(env_ids, [object_id]) + object_id = objects[0] vertices = np.asarray( - [ - get_combined_vertices(self._entities[env_id][object_id]) - for env_id in ids - ], + [get_combined_vertices(self._entities[index][object_id]) for index in env], dtype=np.float32, ) if scale: scales = np.asarray( - [self._entities[env_id][object_id].get_body_scale() for env_id in ids], + [self._entities[index][object_id].get_body_scale() for index in env], dtype=np.float32, ) - vertices = vertices * scales[:, None, :] + vertices *= scales[:, None, :] return torch.as_tensor(vertices, dtype=torch.float32, device=self.device) def get_object_triangles( @@ -445,35 +573,17 @@ def get_object_triangles( object_id: int, env_ids: Sequence[int] | None = None, ) -> torch.Tensor: - """Get one constituent object's triangle indices. - - Args: - object_id: Constituent object index within the group. - env_ids: Environment indices. If ``None``, returns all instances. - - Returns: - Triangle indices with shape ``(N, num_triangles, 3)``. - """ - if not 0 <= object_id < self.num_objects: - raise IndexError( - f"object_id {object_id} is outside [0, {self.num_objects - 1}]." - ) - ids = self._all_indices if env_ids is None else env_ids + """Return one member's render triangles across selected environments.""" + env, objects, _ = self._selected_indices(env_ids, [object_id]) + object_id = objects[0] triangles = np.asarray( - [ - get_combined_triangles(self._entities[env_id][object_id]) - for env_id in ids - ], + [get_combined_triangles(self._entities[index][object_id]) for index in env], dtype=np.int32, ) return torch.as_tensor(triangles, dtype=torch.int32, device=self.device) def get_user_ids(self) -> torch.Tensor: - """Get the user ids of the rigid body group. - - Returns: - torch.Tensor: A tensor of shape (num_envs, num_objects) representing the user ids of the rigid body group. - """ + """Return render user ids with shape ``[env, object]``.""" return torch.as_tensor( [ [entity.get_user_id() for entity in instance] @@ -484,164 +594,79 @@ def get_user_ids(self) -> torch.Tensor: ) def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: - """Clear the dynamics of the rigid bodies by resetting velocities and applying zero forces and torques. - - Args: - env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used. - """ + """Clear velocity and one-step wrench buffers for selected envs.""" if self.is_non_dynamic: return - - local_env_ids = self._all_indices if env_ids is None else env_ids - - if self.device.type == "cpu": - for env_idx in local_env_ids: - for entity in self._entities[env_idx]: - entity.clear_dynamics() - else: - # Apply zero force and torque to the rigid bodies. - zeros = torch.zeros( - (len(local_env_ids) * self.num_objects, 3), - dtype=torch.float32, - device=self.device, - ) - indices = self.body_data.gpu_indices[local_env_ids].flatten() - torch.cuda.synchronize(self.device) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.LINEAR_VELOCITY, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.ANGULAR_VELOCITY, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.FORCE, - ) - self._ps.gpu_apply_rigid_body_data( - data=zeros, - gpu_indices=indices, - data_type=RigidBodyGPUAPIWriteType.TORQUE, - ) + _, _, rows = self._selected_indices(env_ids) + zeros = torch.zeros((len(rows), 3), dtype=torch.float32, device=self.device) + view = self.body_data.body_view + view.apply_linear_velocity(zeros, rows) + view.apply_angular_velocity(zeros, rows) + view.apply_force(zeros, rows) + view.apply_torque(zeros, rows) def set_visual_material( - self, mat: VisualMaterial, env_ids: Sequence[int] | None = None + self, + mat: VisualMaterial, + env_ids: Sequence[int] | None = None, ) -> None: - """Set visual material for the rigid object group. - - Note: - For each entity in the rigid object group, a unique material instance will be created and shared - among all objects in that entity. - - Args: - mat (VisualMaterial): The material to set. - env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - - for i, env_idx in enumerate(local_env_ids): - mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}") - for j, entity in enumerate(self._entities[env_idx]): - entity.set_material(mat_inst.mat) - - # Note: The rigid object group is not supported to change the visual material once created. - # If needed, we should create a visual material dict to store the material instances, and - # implement a get_visual_material method to retrieve the material instances. + """Assign one material instance to all members in each selected env.""" + env, _, _ = self._selected_indices(env_ids) + for env_id in env: + material = mat.create_instance(f"{mat.uid}_{self.uid}_{env_id}") + for entity in self._entities[env_id]: + entity.set_material(material.mat) def reset(self, env_ids: Sequence[int] | None = None) -> None: - local_env_ids = self._all_indices if env_ids is None else env_ids - num_instances = len(local_env_ids) - - self.cfg: RigidObjectGroupCfg - body_cfgs = list(self.cfg.rigid_objects.values()) - - init_pos = [] - init_rot = [] - for cfg in body_cfgs: - init_pos.append(cfg.init_pos) - init_rot.append(cfg.init_rot) - - # (num_objects, 3) - pos = torch.as_tensor(init_pos, dtype=torch.float32, device=self.device) - rot = ( - torch.as_tensor(init_rot, dtype=torch.float32, device=self.device) - * torch.pi - / 180.0 - ) - # Convert pos and rot to shape (num_instances, num_objects, dim) - pos = pos.unsqueeze_(0).repeat(num_instances, 1, 1) - rot = rot.unsqueeze_(0).repeat(num_instances, 1, 1) - - mat = matrix_from_euler(rot.reshape(-1, 3), "XYZ") - # Init pose with shape (num_instances, num_objects, 4, 4) - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze_(0) - .repeat(num_instances * self.num_objects, 1, 1) - ) - pose[:, :3, 3] = pos.reshape(-1, 3) - pose[:, :3, :3] = mat - pose = pose.reshape(num_instances, self.num_objects, 4, 4) - self.set_local_pose(pose, env_ids=local_env_ids) - - self.clear_dynamics(env_ids=local_env_ids) + env, _, _ = self._selected_indices(env_ids) + self._restore_default_physical_properties(env) + member_poses = [] + for cfg in self.cfg.rigid_objects.values(): + if cfg.init_local_pose is not None: + member_poses.append( + torch.as_tensor( + cfg.init_local_pose, + dtype=torch.float32, + device=self.device, + ).reshape(4, 4) + ) + continue + pose = torch.eye(4, dtype=torch.float32, device=self.device) + pose[:3, 3] = torch.as_tensor( + cfg.init_pos, dtype=torch.float32, device=self.device + ) + rotation = torch.as_tensor( + cfg.init_rot, dtype=torch.float32, device=self.device + ) + pose[:3, :3] = matrix_from_euler( + (rotation * torch.pi / 180.0).reshape(1, 3), "XYZ" + )[0] + member_poses.append(pose) + pose = torch.stack(member_poses).repeat(len(env), 1, 1) + self.set_local_pose(pose.reshape(len(env), self.num_objects, 4, 4), env_ids=env) + self.clear_dynamics(env_ids=env) def set_physical_visible( self, visible: bool = True, rgba: Sequence[float] | None = None, - ): - """set collion render visibility - - Args: - visible (bool, optional): is collision body visible. Defaults to True. - rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None. - """ - rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7) - if len(rgba) != 4: - logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.") - - # create collision visible node if not exist - if visible: - for i, env_idx in enumerate(self._all_indices): - for intance_id, entity in enumerate(self._entities[env_idx]): - if not self._has_collision_visible_node_list[intance_id]: - entity.create_physical_visible_node( - np.array( - [ - rgba[0], - rgba[1], - rgba[2], - rgba[3], - ] - ) - ) - self._has_collision_visible_node_list[intance_id] = True - - # create collision visible node if not exist - for i, env_idx in enumerate(self._all_indices): - for entity in self._entities[env_idx]: - entity.set_physical_visible(visible) + ) -> None: + """Set collision-geometry visibility for every Group member.""" + color = np.asarray( + (0.8, 0.2, 0.2, 0.7) if rgba is None else rgba, + dtype=np.float32, + ) + if color.shape != (4,): + raise ValueError("Collision visualization color must contain four values.") + for instance in self._entities: + for entity in instance: + self._spawn_result.set_physical_visible(entity, color, visible) def set_visible(self, visible: bool = True) -> None: - """Set the visibility of the rigid object group. - - Args: - visible (bool, optional): Whether the rigid object group is visible. Defaults to True. - """ - for i, env_idx in enumerate(self._all_indices): - for entity in self._entities[env_idx]: + """Set render visibility for every Group member.""" + for instance in self._entities: + for entity in instance: entity.set_visible(visible) def destroy(self) -> None: - env = self._world.get_env() - arenas = env.get_all_arenas() - if len(arenas) == 0: - arenas = [env] - for i, instance in enumerate(self._entities): - for entity in instance: - arenas[i].remove_actor(entity) + """Leave topology destruction to SimulationManager and SpawnResult.""" diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index 7b8a1340e..b7025170f 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -19,7 +19,7 @@ import torch import numpy as np -from typing import Dict, List, Literal, Sequence, Tuple +from typing import TYPE_CHECKING, Dict, List, Literal, Sequence, Tuple from dataclasses import dataclass, field from tensordict import TensorDict @@ -39,6 +39,9 @@ ) from embodichain.utils import logger +if TYPE_CHECKING: + from dexsim.spawn import SpawnResult, SpawnedArticulation + @dataclass class ControlGroup: @@ -71,11 +74,14 @@ class Robot(Articulation): def __init__( self, cfg: RobotCfg, - entities: List[_Articulation], + entities: List[_Articulation | SpawnedArticulation] | None = None, device: torch.device = torch.device("cpu"), + *, + spawn_result: SpawnResult | None = None, + declared_num_instances: int | None = None, ) -> None: - self._entities = entities + self._entities = [] if entities is None else entities self.cfg = cfg # Initialize joint ids for control parts. @@ -91,12 +97,18 @@ def __init__( # cache I/O unless a task actually requests workspace sampling. self._workspaces: Dict[str, RobotWorkspace] = {} - if self.cfg.control_parts: + if entities is not None and self.cfg.control_parts: self._init_control_parts(self.cfg.control_parts) - super().__init__(cfg, entities, device) + super().__init__( + cfg, + entities, + device, + spawn_result=spawn_result, + declared_num_instances=declared_num_instances, + ) - if self.cfg.solver_cfg: + if entities is not None and self.cfg.solver_cfg: self.init_solver(self.cfg.solver_cfg) def __str__(self) -> str: @@ -106,6 +118,19 @@ def __str__(self) -> str: + f" | control_parts: {self.control_parts}, solvers: {self._solvers}" ) + def attach_spawn_handles( + self, + entities: Sequence[SpawnedArticulation], + ) -> None: + """Store handles and expose robot metadata without creating Batch data. + + Runtime Batch/Data initialization remains the responsibility of + ``bind_spawn()`` after Spawn finalization. + """ + super().attach_spawn_handles(entities) + if self.cfg.control_parts: + self._init_control_parts(self.cfg.control_parts) + @property def control_parts(self) -> Dict[str, List[str]] | None: """Get the control parts of the robot.""" @@ -784,7 +809,9 @@ def compute_fk( to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix. Returns: - torch.Tensor: The forward kinematics result with shape (num_envs, 7) or (num_envs, 4, 4) if `to_matrix` is True. + torch.Tensor: The forward-kinematics result with shape + ``(num_envs, 7)`` in ``(x, y, z, qx, qy, qz, qw)`` order, or + ``(num_envs, 4, 4)`` if ``to_matrix`` is True. """ local_env_ids = self._all_indices if env_ids is None else env_ids @@ -848,7 +875,8 @@ def compute_ik( The input pose should be in the local arena frame. Args: - pose (torch.Tensor): The end effector pose of the robot, (num_envs, 7) or (num_envs, 4, 4). + pose (torch.Tensor): The end-effector pose as ``(num_envs, 7)`` in + ``(x, y, z, qx, qy, qz, qw)`` order or ``(num_envs, 4, 4)``. joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (num_envs, dof). If None, the zero joint positions will be used as the seed. name (str | None): The name of the control part to compute the IK for. If None, the default part is used. @@ -932,7 +960,9 @@ def compute_batch_fk( to_matrix (bool): If True, returns the transformation in the form of a 4x4 matrix. Returns: - torch.Tensor: The forward kinematics result with shape (num_envs, batch, 7) or (num_envs, batch, 4, 4) if `to_matrix` is True. + torch.Tensor: The forward-kinematics result with shape + ``(num_envs, batch, 7)`` in ``xyz + xyzw`` order, or + ``(num_envs, batch, 4, 4)`` if ``to_matrix`` is True. """ local_env_ids = self._all_indices if env_ids is None else env_ids if not self._solvers: @@ -992,7 +1022,8 @@ def compute_batch_ik( The input pose should be in the local arena frame. Args: - pose (torch.Tensor): The end effector pose of the robot, (num_envs, n_batch, 7) or (num_envs, n_batch, 4, 4). + pose (torch.Tensor): End-effector poses as ``(num_envs, n_batch, 7)`` + in ``xyz + xyzw`` order or ``(num_envs, n_batch, 4, 4)``. joint_seed (torch.Tensor | None): The joint positions to use as a seed for the IK computation, (num_envs, n_batch, dof). If None, the zero joint positions will be used as the seed. name (str | None): The name of the control part to compute the IK for. If None, the default part is used. env_ids (Sequence[int] | None): Environment indices to apply the positions. Defaults to all environments. @@ -1088,8 +1119,7 @@ def _init_control_parts(self, control_parts: Dict[str, List[str]]) -> None: joint names or regular expressions that match joint names. """ joint_name_to_ids = { - name: i - for i, name in enumerate(self._entities[0].get_actived_joint_names()) + name: i for i, name in enumerate(self._state_joint_names()) } for name, joint_names in control_parts.items(): # convert joint_names which is a regular expression to a list of joint names @@ -1125,12 +1155,16 @@ def set_joint_drive( max_velocity: torch.Tensor | None = None, friction: torch.Tensor | None = None, armature: torch.Tensor | None = None, - drive_type: str = "force", + drive_type: str | None = "force", joint_ids: Sequence[int] | None = None, env_ids: Sequence[int] | None = None, + *, + target_mode: str | int | None = None, ) -> None: """Set the drive properties for the robot. - Different from Articulation, default drive type is 'force' instead of 'none' + + With no explicit mode, robots retain their position+velocity force + drive default. Args: stiffness (torch.Tensor): The stiffness of the joint drive with shape (len(env_ids), len(joint_ids)). @@ -1139,9 +1173,10 @@ def set_joint_drive( max_velocity (torch.Tensor): The maximum velocity of the joint drive with shape (len(env_ids), len(joint_ids)). friction (torch.Tensor): The joint friction coefficient with shape (len(env_ids), len(joint_ids)). armature (torch.Tensor): The joint armature with shape (len(env_ids), len(joint_ids)). - drive_type (str, optional): The type of drive to apply. Defaults to "force". + drive_type: Drive type to apply. Defaults to ``"force"``. joint_ids (Sequence[int] | None, optional): The joint indices to apply the drive to. If None, applies to all joints. Defaults to None. env_ids (Sequence[int] | None, optional): The environment indices to apply the drive to. If None, applies to all environments. Defaults to None. + target_mode: Portable target mode name or integer value 0 through 4. """ super().set_joint_drive( stiffness=stiffness, @@ -1153,6 +1188,7 @@ def set_joint_drive( drive_type=drive_type, joint_ids=joint_ids, env_ids=env_ids, + target_mode=target_mode, ) def _set_default_joint_drive(self) -> None: @@ -1160,7 +1196,7 @@ def _set_default_joint_drive(self) -> None: import numbers from embodichain.utils.string import resolve_matching_names_values - drive_props = [ + joint_property_targets = [ ("damping", self.default_joint_damping), ("stiffness", self.default_joint_stiffness), ("max_effort", self.default_joint_max_effort), @@ -1169,8 +1205,8 @@ def _set_default_joint_drive(self) -> None: ("armature", self.default_joint_armature), ] - for prop_name, default_array in drive_props: - value = getattr(self.cfg.drive_pros, prop_name, None) + for prop_name, default_array in joint_property_targets: + value = getattr(self.cfg.joint_drive_props, prop_name, None) if value is None: continue if isinstance(value, numbers.Number): @@ -1210,11 +1246,19 @@ def _set_default_joint_drive(self) -> None: except Exception as e: logger.log_error(f"Failed to set {prop_name}: {e}") - drive_pros = self.cfg.drive_pros - if isinstance(drive_pros, dict): - drive_type = drive_pros.get("drive_type", "force") + joint_drive_props = self.cfg.joint_drive_props + if isinstance(joint_drive_props, dict): + drive_type = joint_drive_props.get("drive_type") + target_mode = joint_drive_props.get("target_mode") else: - drive_type = getattr(drive_pros, "drive_type", "force") + drive_type = getattr(joint_drive_props, "drive_type", None) + target_mode = getattr(joint_drive_props, "target_mode", None) + if isinstance(target_mode, dict): + logger.log_warning( + "Per-joint target_mode mappings require a Spawn-bound robot; " + "the retained raw-robot path preserves its current target modes." + ) + target_mode = None # Apply drive parameters to all articulations in the batch self.set_joint_drive( @@ -1225,6 +1269,7 @@ def _set_default_joint_drive(self) -> None: friction=self.default_joint_friction, armature=self.default_joint_armature, drive_type=drive_type, + target_mode=target_mode, ) def _sync_solver_limits(self, name: str | None = None) -> None: @@ -1385,21 +1430,29 @@ def _extract_control_group(self, joint_names: List[str]) -> ControlGroup: """ control_group = ControlGroup() joint_id_list = [] + state_joint_ids = { + name: index for index, name in enumerate(self._state_joint_names()) + } + source_joint_ids = { + name: index + for index, name in enumerate(self._entities[0].get_actived_joint_names()) + } for joint_name in joint_names: - if joint_name in self.joint_names: - joint_index = self.joint_names.index(joint_name) - joint_id_list.append(joint_index) + if joint_name in state_joint_ids and joint_name in source_joint_ids: + joint_id_list.append(state_joint_ids[joint_name]) control_group.joint_names.append(joint_name) # Set root link for first joint if len(control_group.link_names) == 0: parent_names = self._entities[0].get_ancestral_link_names( - joint_index + source_joint_ids[joint_name] ) control_group.link_names.extend(parent_names) - child_name = self._entities[0].get_child_link_name(joint_index) + child_name = self._entities[0].get_child_link_name( + source_joint_ids[joint_name] + ) control_group.link_names.append(child_name) control_group.joint_ids = joint_id_list @@ -1439,6 +1492,17 @@ def set_physical_visible( ) link_names = self.get_control_part_link_names(name=control_part) + if self.is_spawn_bound: + for env_idx in self._all_indices: + entity = self._entities[env_idx] + for link_name in link_names: + self._spawn_result.set_physical_visible( + (entity, link_name), rgba, visible + ) + for link_name in link_names: + self._has_collision_visible_node_dict[link_name] = True + return + # create collision visible node if not exist if visible: for i, env_idx in enumerate(self._all_indices): diff --git a/embodichain/lab/sim/objects/soft_object.py b/embodichain/lab/sim/objects/soft_object.py index 52344150c..e56592da7 100644 --- a/embodichain/lab/sim/objects/soft_object.py +++ b/embodichain/lab/sim/objects/soft_object.py @@ -14,522 +14,24 @@ # limitations under the License. # ---------------------------------------------------------------------------- -from __future__ import annotations +"""Compatibility exports for the volume-deformable object API.""" -import torch -import dexsim -import numpy as np -from functools import cached_property +from __future__ import annotations -from dataclasses import dataclass -from typing import List, Sequence, Union +from embodichain.lab.sim.cfg import SoftObjectCfg, VolumeDeformableObjectCfg -from dexsim.models import MeshObject -from dexsim.engine import PhysicsScene, SoftBody -from dexsim.types import SoftBodyGPUAPIReadWriteType -from scipy.spatial import ConvexHull, QhullError -from embodichain.lab.sim.common import ( - BatchEntity, -) -from embodichain.lab.sim.material import ( - VisualMaterial, - VisualMaterialInst, - _capture_render_materials, - _restore_render_materials, - _wrap_first_render_material, +from .deformable.volume import ( + SoftBodyData, + SoftObject, + VolumeDeformableData, + VolumeDeformableObject, ) -from embodichain.utils.math import ( - matrix_from_euler, -) -from embodichain.utils import logger -from embodichain.lab.sim.cfg import ( - SoftObjectCfg, -) -from embodichain.utils.math import xyz_quat_to_4x4_matrix - -__all__ = ["SoftBodyData", "SoftObject", "SoftObjectCfg"] - - -@dataclass -class SoftBodyData: - """Data manager for soft body - - Note: - 1. The pose data managed by dexsim is in the format of (qx, qy, qz, qw, x, y, z), but in EmbodiChain, we use (x, y, z, qw, qx, qy, qz) format. - """ - - def __init__( - self, entities: List[MeshObject], ps: PhysicsScene, device: torch.device - ) -> None: - """Initialize the SoftBodyData. - - Args: - entities (List[MeshObject]): List of MeshObjects representing the soft bodies. - ps (PhysicsScene): The physics scene. - device (torch.device): The device to use for the soft body data. - """ - self.entities = entities - # TODO: soft body data can only be stored in cuda device for now. - self.device = device - # TODO: inorder to retrieve arena position, we need to access the node of each entity. - self.ps = ps - self.num_instances = len(entities) - - self.soft_bodies: Sequence[SoftBody] = [ - self.entities[i].get_physical_body() for i in range(self.num_instances) - ] - self.n_collision_vertices = self.soft_bodies[0].get_num_vertices() - self.n_sim_vertices = self.soft_bodies[0].get_num_sim_vertices() - - self._rest_position_buffer = torch.empty( - (self.num_instances, self.n_collision_vertices, 4), - device=self.device, - dtype=torch.float32, - ) - for i, softbody in enumerate(self.soft_bodies): - self._rest_position_buffer[i] = softbody.get_position_inv_mass_buffer() - - self._rest_sim_position_buffer = torch.empty( - (self.num_instances, self.n_sim_vertices, 4), - device=self.device, - dtype=torch.float32, - ) - - for i, softbody in enumerate(self.soft_bodies): - self._rest_sim_position_buffer[i] = ( - softbody.get_sim_position_inv_mass_buffer() - ) - - self._collision_position = torch.zeros( - (self.num_instances, self.n_collision_vertices, 3), - device=self.device, - dtype=torch.float32, - ) - self._sim_vertex_velocity = torch.zeros( - (self.num_instances, self.n_sim_vertices, 3), - device=self.device, - dtype=torch.float32, - ) - self._sim_vertex_position = torch.zeros( - (self.num_instances, self.n_sim_vertices, 3), - device=self.device, - dtype=torch.float32, - ) - - @property - def rest_collision_vertices(self): - """Get the rest position buffer of the soft bodies.""" - return self._rest_position_buffer[:, :, :3].clone() - - @property - def rest_sim_vertices(self): - """Get the rest sim position buffer of the soft bodies.""" - return self._rest_sim_position_buffer[:, :, :3].clone() - - @property - def collision_position(self): - """Get the current vertex position buffer of the soft bodies.""" - for i, softbody in enumerate(self.soft_bodies): - self._collision_position[i] = softbody.get_position_inv_mass_buffer()[:, :3] - return self._collision_position.clone() - - @property - def sim_vertex_position(self): - """Get the current sim vertex position buffer of the soft bodies.""" - for i, softbody in enumerate(self.soft_bodies): - self._sim_vertex_position[i] = softbody.get_sim_position_inv_mass_buffer()[ - :, :3 - ] - return self._sim_vertex_position.clone() - - @property - def sim_vertex_velocity(self): - """Get the current vertex velocity buffer of the soft bodies.""" - for i, softbody in enumerate(self.soft_bodies): - self._sim_vertex_velocity[i] = softbody.get_sim_velocity_buffer()[:, :3] - return self._sim_vertex_velocity.clone() - - @cached_property - def collision_surface_triangles(self) -> torch.Tensor: - """Build a stable surface approximation for collision vertices. - - DexSim exposes live PhysX collision vertices but not their triangle - connectivity. The convex hull provides a stable topology whose indices - continue to reference the live collision-vertex buffer. - - Returns: - Cached convex-hull triangle indices. - """ - vertices = self.rest_collision_vertices[0].detach().cpu().numpy() - if vertices.shape[0] < 4: - logger.log_warning( - "Soft-body collision geometry has fewer than four vertices; " - "its visualization surface will be empty." - ) - triangles = np.empty((0, 3), dtype=np.int32) - else: - try: - triangles = np.asarray( - ConvexHull(vertices).simplices, - dtype=np.int32, - ) - except QhullError as error: - try: - triangles = np.asarray( - ConvexHull(vertices, qhull_options="QJ").simplices, - dtype=np.int32, - ) - except QhullError: - logger.log_warning( - "Unable to build a soft-body visualization surface from " - f"collision vertices: {error!r}" - ) - triangles = np.empty((0, 3), dtype=np.int32) - return torch.as_tensor( - triangles, - dtype=torch.int32, - device=self.device, - ) - - -class SoftObject(BatchEntity): - """SoftObject represents a batch of soft body in the simulation.""" - - def __init__( - self, - cfg: SoftObjectCfg, - entities: List[MeshObject] = None, - device: torch.device = torch.device("cpu"), - ) -> None: - self._world: dexsim.World = dexsim.default_world() - self._ps = self._world.get_physics_scene() - self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist() - - self._data = SoftBodyData(entities=entities, ps=self._ps, device=device) - - self._world.update(0.001) - - self._visual_material: List[VisualMaterialInst | None] = [None] * len(entities) - self.is_shared_visual_material = False - - super().__init__(cfg=cfg, entities=entities, device=device) - - self._initialize_existing_visual_material() - - # set default collision filter - self._set_default_collision_filter() - - def _initialize_existing_visual_material(self) -> None: - """Wrap asset-parsed materials during soft-object construction. - - For a multi-segment render body, the first segment with a valid - material is registered as the environment's representative material. - """ - self._original_visual_material = [[] for _ in self._entities] - self._original_visual_material_inst = [None] * len(self._entities) - for env_idx, entity in enumerate(self._entities): - render_body = entity.get_render_body() - if render_body is None: - continue - original_materials = _capture_render_materials(render_body) - self._original_visual_material[env_idx] = original_materials - wrapped = _wrap_first_render_material(original_materials) - if wrapped is not None: - self._visual_material[env_idx] = wrapped - self._original_visual_material_inst[env_idx] = wrapped - - def set_visual_material( - self, - mat: VisualMaterial, - env_ids: Sequence[int] | None = None, - shared: bool = False, - ) -> None: - """Set visual material for the soft object. - - Args: - mat: The material template to assign. - env_ids: Environment indices. If None, all instances are used. - shared: Whether selected environments share one material instance. - """ - local_env_ids = self._all_indices if env_ids is None else env_ids - if shared: - if len(local_env_ids) != self.num_instances: - logger.log_error("Cannot share material instance for partial env_ids.") - mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}") - for env_idx in local_env_ids: - self._entities[env_idx].set_material(mat_inst.mat) - self._visual_material[env_idx] = mat_inst - self.is_shared_visual_material = True - else: - for env_idx in local_env_ids: - mat_inst = mat.create_instance(f"{mat.uid}_{self.uid}_{env_idx}") - self._entities[env_idx].set_material(mat_inst.mat) - self._visual_material[env_idx] = mat_inst - self.is_shared_visual_material = False - - def restore_visual_material(self, env_ids: Sequence[int] | None = None) -> None: - """Restore visual materials captured when the soft object was created. - - Args: - env_ids: Environment indices. If None, all instances are restored. - """ - if not hasattr(self, "_original_visual_material"): - return - local_env_ids = self._all_indices if env_ids is None else env_ids - for env_idx in local_env_ids: - render_body = self._entities[env_idx].get_render_body() - if render_body is None: - continue - _restore_render_materials( - render_body, self._original_visual_material[env_idx] - ) - self._visual_material[env_idx] = self._original_visual_material_inst[ - env_idx - ] - self.is_shared_visual_material = False - - def get_visual_material_inst( - self, env_ids: Sequence[int] | None = None - ) -> List[VisualMaterialInst | None]: - """Get the material instance registered for each selected environment. - - Args: - env_ids: Environment indices. If None, all instances are returned. - - Returns: - The existing material wrappers, or None where an asset has no material. - """ - ids = env_ids if env_ids is not None else range(self.num_instances) - return [self._visual_material[i] for i in ids] - - def _set_default_collision_filter(self) -> None: - collision_filter_data = torch.zeros( - size=(self.num_instances, 4), dtype=torch.int32 - ) - for i in range(self.num_instances): - collision_filter_data[i, 0] = i - collision_filter_data[i, 1] = 1 - self.set_collision_filter(collision_filter_data) - - def set_collision_filter( - self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None - ) -> None: - """Set collision filter data for the soft object. - - Args: - filter_data (torch.Tensor): [N, 4] of int. - First element of each object is arena id. - If 2nd element is 0, the object will collision with all other objects in world. - 3rd and 4th elements are not used currently. - - 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 - - if len(local_env_ids) != len(filter_data): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(filter_data)}." - ) - - filter_data_np = filter_data.cpu().numpy().astype(np.uint32) - for i, env_idx in enumerate(local_env_ids): - self._entities[env_idx].get_physical_body().set_collision_filter_data( - filter_data_np[i] - ) - - @property - def body_data(self) -> SoftBodyData | None: - """Get the soft body data manager for this soft object. - - Returns: - SoftBodyData | None: The soft body data manager. - """ - return self._data - - def set_local_pose( - self, pose: torch.Tensor, env_ids: Sequence[int] | None = None - ) -> None: - """Set local pose of the soft object. - - Args: - pose (torch.Tensor): The local pose of the soft object with shape (N, 7) or (N, 4, 4). - env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used. - """ - from embodichain.lab.sim import SimulationManager - - sim = SimulationManager.get_instance() - local_env_ids = self._all_indices if env_ids is None else env_ids - - if len(local_env_ids) != len(pose): - logger.log_error( - f"Length of env_ids {len(local_env_ids)} does not match pose length {len(pose)}." - ) - - if pose.dim() == 2 and pose.shape[1] == 7: - pose4x4 = xyz_quat_to_4x4_matrix(pose) - elif pose.dim() == 3 and pose.shape[1:3] == (4, 4): - pose4x4 = pose - else: - logger.log_error( - f"Invalid pose shape {pose.shape}. Expected (N, 7) or (N, 4, 4)." - ) - - arena_offsets = sim.arena_offsets - for i, env_idx in enumerate(local_env_ids): - # TODO: soft body cannot directly set by `set_local_pose` currently. - rest_collision_vertices = self.body_data.rest_collision_vertices[i] - rest_sim_vertices = self.body_data.rest_sim_vertices[i] - rotation = pose4x4[i][:3, :3] - translation = pose4x4[i][:3, 3] - - # apply transformation to local rest vertices and back - rest_collision_vertices_local = rest_collision_vertices - arena_offsets[i] - transformed_collision_vertices = ( - rest_collision_vertices_local @ rotation.T + translation - ) - transformed_collision_vertices = ( - transformed_collision_vertices + arena_offsets[i] - ) - - rest_sim_vertices_local = rest_sim_vertices - arena_offsets[i] - transformed_sim_vertices = ( - rest_sim_vertices_local @ rotation.T + translation - ) - transformed_sim_vertices = transformed_sim_vertices + arena_offsets[i] - - # apply vertices to soft body - soft_body: SoftBody = self._entities[env_idx].get_physical_body() - collision_position_buffer = soft_body.get_position_inv_mass_buffer() - sim_position_buffer = soft_body.get_sim_position_inv_mass_buffer() - sim_velocity_buffer = soft_body.get_sim_velocity_buffer() - - collision_position_buffer[:, :3] = transformed_collision_vertices - sim_position_buffer[:, :3] = transformed_sim_vertices - sim_velocity_buffer[:, :3] = 0.0 - - soft_body.mark_dirty(SoftBodyGPUAPIReadWriteType.ALL) - # TODO: currently soft body has no wake up interface, use set_wake_counter and pass in a positive value to wake it up - soft_body.set_wake_counter(0.4) - - def get_rest_collision_vertices(self) -> torch.Tensor: - """Get the rest collision vertices of the soft object. - - Returns: - torch.Tensor: The rest collision vertices with shape (N, num_collision_vertices, 3). - """ - return self.body_data.rest_collision_vertices - - def get_rest_sim_vertices(self) -> torch.Tensor: - """Get the rest sim vertices of the soft object. - - Returns: - torch.Tensor: The rest sim vertices with shape (N, num_sim_vertices, 3). - """ - return self.body_data.rest_sim_vertices - - def get_current_collision_vertices(self) -> torch.Tensor: - """Get the current collision vertices of the soft object. - - Returns: - torch.Tensor: The current collision vertices with shape (N, num_collision_vertices, 3). - """ - return self.body_data.collision_position - - def get_current_sim_vertices(self) -> torch.Tensor: - """Get the current sim vertices of the soft object. - - Returns: - torch.Tensor: The current sim vertices with shape (N, num_sim_vertices, 3). - """ - return self.body_data.sim_vertex_position - - def get_current_sim_vertex_velocities(self) -> torch.Tensor: - """Get the current sim vertex velocities of the soft object. - - Returns: - torch.Tensor: The current sim vertex velocities with shape (N, num_sim_vertices, 3). - """ - return self.body_data.sim_vertex_velocity - - def get_collision_surface_triangles( - self, env_ids: Sequence[int] | None = None - ) -> torch.Tensor: - """Get approximate collision-surface triangles for selected instances. - - DexSim currently exposes live soft-body collision vertices without - their topology. This method returns a cached convex-hull topology, so - it is suitable for low-frequency external visualization but does not - preserve concave details of the render mesh. - - Args: - env_ids: Environment indices. If ``None``, returns all instances. - - Returns: - Triangle indices with shape ``(N, num_triangles, 3)``. - """ - ids = self._all_indices if env_ids is None else env_ids - return ( - self.body_data.collision_surface_triangles.unsqueeze(0) - .expand(len(ids), -1, -1) - .clone() - ) - - def get_triangles(self, env_ids: Sequence[int] | None = None) -> torch.Tensor: - """Get approximate surface triangles for generic mesh consumers. - - Args: - env_ids: Environment indices. If ``None``, returns all instances. - - Returns: - Triangle indices with shape ``(N, num_triangles, 3)``. - """ - return self.get_collision_surface_triangles(env_ids=env_ids) - - def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: - """Get local pose of the soft 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. - - Returns: - torch.Tensor: The local pose of the soft object with shape (N, 7) or (N, 4, 4) depending on `to_matrix`. - """ - raise NotImplementedError("Getting local pose for SoftObject is not supported.") - - 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.restore_visual_material(env_ids=local_env_ids) - - # TODO: set attr for soft body after loading in physics scene. - - # rest soft body to init_pos - pos = torch.as_tensor( - self.cfg.init_pos, dtype=torch.float32, device=self.device - ) - rot = ( - torch.as_tensor(self.cfg.init_rot, dtype=torch.float32, device=self.device) - * torch.pi - / 180.0 - ) - pos = pos.unsqueeze(0).repeat(num_instances, 1) - rot = rot.unsqueeze(0).repeat(num_instances, 1) - mat = matrix_from_euler(rot, "XYZ") - pose = ( - torch.eye(4, dtype=torch.float32, device=self.device) - .unsqueeze(0) - .repeat(num_instances, 1, 1) - ) - pose[:, :3, 3] = pos - pose[:, :3, :3] = mat - self.set_local_pose(pose, env_ids=local_env_ids) - def destroy(self) -> None: - # TODO: not tested yet - 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) +__all__ = [ + "SoftBodyData", + "SoftObject", + "SoftObjectCfg", + "VolumeDeformableData", + "VolumeDeformableObject", + "VolumeDeformableObjectCfg", +] diff --git a/embodichain/lab/sim/physics/__init__.py b/embodichain/lab/sim/physics/__init__.py new file mode 100644 index 000000000..deb91a875 --- /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 000000000..1eda6aa6c --- /dev/null +++ b/embodichain/lab/sim/physics/base.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Spawn-aware physics-backend abstraction for :class:`SimulationManager`. + +This module defines the contract that every physics backend (Default, Newton, +...) satisfies. The owning :class:`SimulationManager` +holds a single :class:`PhysicsBackend` instance as ``self.physics`` and +delegates backend-specific world configuration, compatibility scene access, +and capability queries to it. Scene topology and runtime readiness are owned +by DexSim's ``SceneBuilder`` and ``SpawnResult``. + +The design deliberately mirrors IsaacLab's split of an orchestrator +(``SimulationContext``) from a swappable physics manager (``PhysicsManager``), +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.sim_manager import SimulationManager, SimulationManagerCfg + +__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. + + Default configures the native DexSim globals. Newton is already + registered from ``WorldConfig.newton_cfg`` and therefore has no + additional activation work. + """ + + def sync_render_state(self, result: "dexsim.spawn.SpawnResult") -> None: + """Publish the current physics state to render resources without stepping. + + Backends whose physics and render state share native storage require no + work. Backends with a separate render bridge override this hook. + + Args: + result: The finalized Spawn result whose state should be published. + """ + del result + + def prepare_for_teardown(self) -> None: + """Release backend-owned views before Spawn releases their parents. + + :class:`SimulationManager` calls this during deferred destruction, + after render workers stop and before it closes the Spawn result. A + backend can use this boundary to synchronize device work and release + borrowed render or physics views while their World-owned native + parents are still alive. Backends without such views keep the default + no-op implementation. + """ + + # ------------------------------------------------------------------ # + # Scene access + # ------------------------------------------------------------------ # + @abstractmethod + def get_scene(self): + """Return a backend compatibility scene, or raise if none exists.""" + + @property + def newton_manager(self): + """Return ``None`` because Spawn does not use ``NewtonManager``. + + The Newton backend overrides this property with an actionable error so + callers do not accidentally mix the removed manager ownership domain + with the World-owned Spawn backend. + """ + return None + + @property + def differentiable_runtime(self): + """Return no differentiable runtime for non-Newton backends.""" + return None + + # ------------------------------------------------------------------ # + # Capabilities (override in subclasses; defaults are conservative) + # ------------------------------------------------------------------ # + @property + def supports_volume_deformables(self) -> bool: + """Whether this backend has a volume-deformable object adapter.""" + return False + + @property + def supports_surface_deformables(self) -> bool: + """Whether this backend has a surface-deformable object adapter.""" + return False + + @property + def supports_soft_bodies(self) -> bool: + """Compatibility alias for volume-deformable support.""" + return self.supports_volume_deformables + + @property + def supports_cloth(self) -> bool: + """Compatibility alias for surface-deformable support.""" + return self.supports_surface_deformables + + @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 000000000..5dc1e9381 --- /dev/null +++ b/embodichain/lab/sim/physics/default.py @@ -0,0 +1,78 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""Default physics backend implementation integrated through DexSim.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import dexsim + +from embodichain.lab.sim.cfg import DefaultPhysicsCfg + +from .base import PhysicsBackend + +if TYPE_CHECKING: + from embodichain.lab.sim.sim_manager import SimulationManagerCfg + +__all__ = ["DefaultPhysicsBackend"] + + +class DefaultPhysicsBackend(PhysicsBackend): + """Default backend using DexSim's native GPU or CPU physics path.""" + + name = "default" + + # -- 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()) + + # -- scene ---------------------------------------------------------- # + def get_scene(self): + """Return the Default backend's compatibility scene after Spawn is prepared.""" + self._manager.prepare() + return self._manager._world.get_physics_scene() + + # -- capabilities --------------------------------------------------- # + # The default backend supports deformables on GPU; the GPU + # precondition itself is enforced separately in SimulationManager. + @property + def supports_volume_deformables(self) -> bool: + return True + + @property + def supports_surface_deformables(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 000000000..22c3c28f8 --- /dev/null +++ b/embodichain/lab/sim/physics/newton.py @@ -0,0 +1,188 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- +"""World-owned Newton (Warp) physics backend configuration.""" + +from __future__ import annotations + +import importlib +from typing import TYPE_CHECKING +import weakref + +import warp as wp + +from .base import PhysicsBackend + +if TYPE_CHECKING: + import dexsim + + from embodichain.lab.sim.sim_manager import SimulationManagerCfg + +__all__ = ["NewtonPhysicsBackend"] + + +def is_newton_gradient_mode(result) -> bool: + """Return whether a finalized Spawn result uses Newton gradients.""" + if result is None or getattr(result, "backend", None) != "newton": + return False + from dexsim.engine.newton_physics.backend_registry import get_newton_backend + + backend = get_newton_backend(result.world) + if backend is None: + return False + return bool( + backend.cfg.requires_grad + or (backend.model is not None and backend.model.requires_grad) + ) + + +class NewtonPhysicsBackend(PhysicsBackend): + """The Warp-based Newton physics backend integrated through DexSim.""" + + name = "newton" + + def __init__(self, manager) -> None: + super().__init__(manager) + self._differentiable_runtime = None + self._runtime_device: str | None = None + self._configured_solver_type: str | None = None + + @property + def solver_type(self) -> str | None: + """Return the configured or scene-resolved Newton solver type.""" + world = getattr(self._manager, "_world", None) + if world is not None: + from dexsim.engine.newton_physics.backend_registry import ( + get_newton_backend, + ) + + backend = get_newton_backend(world) + if backend is not None: + return str(backend.solver_type) + return self._configured_solver_type + + # -- 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 + newton_cfg = newton_physics_cfg.to_dexsim_cfg( + gpu_id=sim_config.gpu_id, + ) + self._configured_solver_type = str(newton_cfg.solver_cfg.solver_type) + self._runtime_device = str(newton_cfg.device) + world_config.newton_cfg = newton_cfg + + def activate(self, sim_config: "SimulationManagerCfg") -> None: + del sim_config + # WorldConfig.newton_cfg registers the World-owned NewtonBackend. + # SceneBuilder.finalize() completes its model; no second manager-level + # activation or rebuild domain participates. + + def sync_render_state(self, result: "dexsim.spawn.SpawnResult") -> None: + """Publish Newton state through DexSim's render bridge without stepping.""" + from dexsim.engine.newton_physics.backend_registry import get_newton_backend + + backend = get_newton_backend(result.world) + if backend is None: + raise RuntimeError( + "Newton backend is unavailable for render-state synchronization." + ) + backend.sync_to_dexsim(result.world) + backend.sync_particle_fluids(result.world) + + def prepare_for_teardown(self) -> None: + """Release Newton render views while Spawn still owns their parents.""" + if self._runtime_device is not None and self._runtime_device.startswith("cuda"): + wp.synchronize_device(self._runtime_device) + + world = getattr(self._manager, "_world", None) + if world is None: + return + + from dexsim.engine.newton_physics.backend_registry import get_newton_backend + + backend = get_newton_backend(world) + if backend is not None: + # NewtonRenderSync retains native link-node wrappers. They must be + # released before SpawnResult.close() drops the owning skeletons; + # otherwise pybind can destruct a child after its native parent. + backend.render_sync.clear() + + @property + def newton_manager(self): + """Reject access to the removed, independently owned Newton manager.""" + raise RuntimeError( + "NewtonManager is not part of Spawn scene ownership. Use " + "SimulationManager.spawn_result and its Spawned*/Batch APIs." + ) + + @property + def differentiable_runtime(self): + """Return the differentiable facade over the Spawn-owned runtime.""" + if self._differentiable_runtime is None: + from embodichain.lab.sim.diff.runtime import NewtonDifferentiableRuntime + + owner_ref = weakref.ref(self) + + def backend_provider(): + owner = owner_ref() + if owner is None: + return None + result = owner._manager.spawn_result + if result is None: + return None + from dexsim.engine.newton_physics.backend_registry import ( + get_newton_backend, + ) + + return get_newton_backend(result.world) + + self._differentiable_runtime = NewtonDifferentiableRuntime(backend_provider) + return self._differentiable_runtime + + # -- scene ---------------------------------------------------------- # + def get_scene(self): + raise RuntimeError( + "Newton Spawn scenes do not expose a PhysicsScene. Use " + "SimulationManager.spawn_result and its Spawned*/Batch APIs." + ) + + # -- capabilities --------------------------------------------------- # + @property + def supports_volume_deformables(self) -> bool: + # Reserved entry point: add a Newton volume adapter before enabling. + return False + + @property + def supports_surface_deformables(self) -> bool: + # Reserved entry point: add a Newton surface adapter before enabling. + return False + + @property + def supports_robot(self) -> bool: + # Robots are SpawnedArticulations in the World-owned Newton model. + return True + + @property + def supports_rigid_object_group(self) -> bool: + # Groups are env-major views over the Spawn rigid-body batch, which + # provides the same state and mass-property API on Newton. + 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/planners/curobo/curobo_planner.py b/embodichain/lab/sim/planners/curobo/curobo_planner.py index d310c26a2..b25dd364f 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_planner.py +++ b/embodichain/lab/sim/planners/curobo/curobo_planner.py @@ -46,7 +46,7 @@ import yaml from embodichain.utils import configclass, logger -from embodichain.utils.math import pose_inv, quat_from_matrix +from embodichain.utils.math import convert_quat, pose_inv, quat_from_matrix from embodichain.lab.sim.planners.base_planner import ( BasePlanner, @@ -520,7 +520,9 @@ def _matrix_to_position_quaternion( # so materialize them at the adapter boundary rather than relying on a # caller-specific layout. position = matrix[:, :3, 3].contiguous() - quaternion = quat_from_matrix(matrix[:, :3, :3]).contiguous() # wxyz + quaternion = convert_quat( + quat_from_matrix(matrix[:, :3, :3]), to="wxyz" + ).contiguous() return position, quaternion diff --git a/embodichain/lab/sim/planners/curobo/curobo_yaml.py b/embodichain/lab/sim/planners/curobo/curobo_yaml.py index 65015695e..d328cc2d2 100644 --- a/embodichain/lab/sim/planners/curobo/curobo_yaml.py +++ b/embodichain/lab/sim/planners/curobo/curobo_yaml.py @@ -34,7 +34,7 @@ import torch from embodichain.utils import logger -from embodichain.utils.math import matrix_from_quat, quat_from_matrix +from embodichain.utils.math import convert_quat, matrix_from_quat, quat_from_matrix if TYPE_CHECKING: from embodichain.lab.sim.objects import RigidObject, Robot @@ -404,7 +404,8 @@ def _mesh_to_obstacle_entry( name: Obstacle name (cuRobo key under ``cuboid``/``mesh``/``sphere``). vertices: Mesh vertices ``(V, 3)`` in the object's local frame. faces: Triangle indices ``(F, 3)`` (any integer dtype). - pose: Object pose as ``(x, y, z, qw, qx, qy, qz)`` ``(7,)`` or a + pose: EmbodiChain object pose as ``(x, y, z, qx, qy, qz, qw)`` + ``(7,)`` or a homogeneous ``(4, 4)`` matrix, expressed in the cuRobo world/base frame (the same frame static collision YAMLs are authored in). representation: ``"cuboid"`` (local-frame AABB -> OBB via ``pose``, @@ -441,13 +442,16 @@ def _mesh_to_obstacle_entry( pose = torch.as_tensor(pose, dtype=torch.float32).detach().to("cpu") if pose.shape == (4, 4): position = pose[:3, 3] - quaternion = quat_from_matrix(pose[:3, :3]) # wxyz + quaternion = quat_from_matrix(pose[:3, :3]) pose = torch.cat([position, quaternion]) if pose.shape != (7,): raise ValueError( - f"pose must be (7,) [x,y,z,qw,qx,qy,qz] or (4, 4), got {tuple(pose.shape)}." + f"pose must be (7,) [x,y,z,qx,qy,qz,qw] or (4, 4), got {tuple(pose.shape)}." ) + # cuRobo world YAML stores 7D poses as xyz+wxyz. + curobo_pose = torch.cat([pose[:3], convert_quat(pose[3:7], to="wxyz")]) + if representation == "mesh": if vertices.numel() == 0 or faces.numel() == 0: raise ValueError( @@ -460,7 +464,7 @@ def _mesh_to_obstacle_entry( { "vertices": vertices.tolist(), "faces": faces.reshape(-1).to(torch.int64).tolist(), - "pose": pose.tolist(), + "pose": curobo_pose.tolist(), }, ) ] @@ -476,9 +480,9 @@ def _mesh_to_obstacle_entry( vmax = vertices.amax(dim=0) dims = vmax - vmin center_local = (vmin + vmax) / 2.0 - rotation = matrix_from_quat(pose[3:7]) # (3, 3), wxyz + rotation = matrix_from_quat(pose[3:7]) center_world = rotation @ center_local + pose[:3] - cuboid_pose = torch.cat([center_world, pose[3:7]]) + cuboid_pose = torch.cat([center_world, curobo_pose[3:7]]) return [("cuboid", name, {"dims": dims.tolist(), "pose": cuboid_pose.tolist()})] # representation == "sphere": fit spheres in the local frame, then transform diff --git a/embodichain/lab/sim/planners/neural_planner.py b/embodichain/lab/sim/planners/neural_planner.py index f5c3d8c90..a86f17667 100644 --- a/embodichain/lab/sim/planners/neural_planner.py +++ b/embodichain/lab/sim/planners/neural_planner.py @@ -31,7 +31,7 @@ ) from embodichain.lab.sim.planners.utils import MoveType, PlanResult, PlanState from embodichain.utils import configclass, logger -from embodichain.utils.math import convert_quat, quat_error_magnitude, quat_from_matrix +from embodichain.utils.math import quat_error_magnitude, quat_from_matrix __all__ = [ "NeuralPlanner", @@ -508,9 +508,7 @@ def _parse_waypoints( if xpos.dim() == 2: xpos = xpos.unsqueeze(0) waypoint_pos[:, idx] = xpos[:, :3, 3] - waypoint_quat[:, idx] = convert_quat( - quat_from_matrix(xpos[:, :3, :3]), to="xyzw" - ) + waypoint_quat[:, idx] = quat_from_matrix(xpos[:, :3, :3]) valid_mask[:, idx] = 1.0 return waypoint_pos, waypoint_quat, valid_mask, len(target_states) @@ -536,9 +534,7 @@ def _fk_matrix(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: def _fk_pose_xyzw(self, qpos: torch.Tensor, control_part: str) -> torch.Tensor: fk = self.robot.compute_fk(qpos=qpos, name=control_part, to_matrix=False) - pos = fk[:, :3] - quat_xyzw = convert_quat(fk[:, 3:7], to="xyzw") - return torch.cat([pos, quat_xyzw], dim=-1) + return fk def _build_obs( self, @@ -589,11 +585,9 @@ def _is_active_reached( idx = torch.arange(b, device=self.device) active_idx_clamped = torch.clamp(active_idx, max=self._num_waypoints - 1) active_pos = waypoint_pos[idx, active_idx_clamped] - active_quat_xyzw = waypoint_quat[idx, active_idx_clamped] + active_quat = waypoint_quat[idx, active_idx_clamped] pos_dist = (ee_pose[:, :3] - active_pos).norm(dim=-1) - ee_quat_wxyz = convert_quat(ee_pose[:, 3:7], to="wxyz") - active_quat_wxyz = convert_quat(active_quat_xyzw, to="wxyz") - rot_dist = quat_error_magnitude(ee_quat_wxyz, active_quat_wxyz) + rot_dist = quat_error_magnitude(ee_pose[:, 3:7], active_quat) orientation_required = self._intermediate_orientation | ( active_idx >= episode_k - 1 ) diff --git a/embodichain/lab/sim/planners/toppra_planner.py b/embodichain/lab/sim/planners/toppra_planner.py index 5d612348c..06fcf4381 100644 --- a/embodichain/lab/sim/planners/toppra_planner.py +++ b/embodichain/lab/sim/planners/toppra_planner.py @@ -250,7 +250,7 @@ class ToppraPlannerCfg(BasePlannerCfg): clears the inherited atexit registry and installs ``prctl(PR_SET_PDEATHSIG)`` so workers are reaped when the parent dies (incl. the ``os._exit`` path). ``'spawn'`` is the safer choice when the parent has initialized CUDA - physics (``sim_device='cuda'``) — fork-after-CUDA-init is the officially + physics (``device='cuda'``) — fork-after-CUDA-init is the officially unsupported case — or if fork deadlocks are observed, at the cost of re-importing modules per worker. """ diff --git a/embodichain/lab/sim/robots/cobotmagic.py b/embodichain/lab/sim/robots/cobotmagic.py index 66041772d..ab51fb1c6 100644 --- a/embodichain/lab/sim/robots/cobotmagic.py +++ b/embodichain/lab/sim/robots/cobotmagic.py @@ -22,10 +22,13 @@ from typing import TYPE_CHECKING, Dict, List, Union from embodichain.lab.sim.cfg import ( + ArticulationRootPropertiesCfg, + CollisionPropertiesCfg, RobotCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, URDFCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, ) from embodichain.lab.sim.solvers import SolverCfg, OPWSolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg @@ -122,9 +125,8 @@ def _build_defaults(self, init_dict: dict | None = None) -> None: ), ), } - self.min_position_iters = 8 - self.min_velocity_iters = 2 - self.drive_pros = JointDrivePropertiesCfg( + self.joint_drive_props = JointDrivePropertiesCfg( + drive_type="force", stiffness={ "left_joint[1-6]": 7e4, "right_joint[1-6]": 7e4, @@ -144,10 +146,19 @@ def _build_defaults(self, init_dict: dict | None = None) -> None: "right_joint[7-8]": 3e3, }, ) - self.attrs = RigidBodyAttributesCfg( - static_friction=0.95, - dynamic_friction=0.9, - contact_offset=0.001, + self.root_props = ArticulationRootPropertiesCfg( + min_position_iters=8, + min_velocity_iters=2, + ) + self.attrs = RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg( + contact_offset=0.001, + rest_offset=0.0, + ), + material_props=RigidBodyMaterialCfg( + static_friction=0.95, + dynamic_friction=0.9, + ), ) @property @@ -185,27 +196,42 @@ def build_pk_serial_chain( if __name__ == "__main__": + import argparse + from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain.lab.sim.cfg import RenderCfg + from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend from embodichain.lab.sim.robots import CobotMagicCfg + parser = argparse.ArgumentParser(description="Launch the CobotMagic robot") + parser.add_argument( + "--physics", + choices=("default", "newton"), + default="default", + help="Physics backend to launch (default: default).", + ) + args = parser.parse_args() + torch.set_printoptions(precision=5, sci_mode=False) config = SimulationManagerCfg( headless=True, - sim_device="cpu", + device="cpu", num_envs=2, + physics_cfg=physics_cfg_for_backend(args.physics), 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() + sim.prepare() + sim.open_window() + from IPython import embed - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + embed() # noqa: E702 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 2cea0df17..e18af9046 100644 --- a/embodichain/lab/sim/robots/dexforce_w1/cfg.py +++ b/embodichain/lab/sim/robots/dexforce_w1/cfg.py @@ -21,6 +21,15 @@ import numpy as np import torch +if __name__ == "__main__" and not __package__: + # Support running this example by file path from an uninstalled source tree. + import sys + from pathlib import Path + + # Replace the script directory so its ``types.py`` cannot shadow the + # standard-library ``types`` module in compiler subprocesses. + sys.path[0] = str(Path(__file__).resolve().parents[5]) + from typing import TYPE_CHECKING, Dict from embodichain.lab.sim.robots.dexforce_w1.types import ( @@ -39,9 +48,12 @@ ) from embodichain.lab.sim.robots.dexforce_w1.specs import get_w1_version_spec from embodichain.lab.sim.cfg import ( + ArticulationRootPropertiesCfg, + CollisionPropertiesCfg, RobotCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg from embodichain.utils import configclass @@ -164,7 +176,7 @@ def _build_defaults(self, init_dict: dict | None = None) -> None: Reads ``version``/``with_default_eef`` from ``init_dict``, sets them on ``self``, then populates ``urdf_cfg``, ``control_parts``, - ``solver_cfg``, ``drive_pros`` and ``attrs``. + ``solver_cfg``, ``joint_drive_props`` and ``attrs``. """ init_dict = init_dict or {} self.version = DexforceW1Version.parse( @@ -272,28 +284,38 @@ def _build_default_physics_cfgs( "damping": {ARM_JOINTS: 1e3, BODY_JOINTS: 1e4, HEAD_JOINTS: 1e3}, "max_effort": {ARM_JOINTS: 1e5, BODY_JOINTS: 1e10, HEAD_JOINTS: 1e5}, } - drive_pros = JointDrivePropertiesCfg(**joint_params) + joint_drive_props = JointDrivePropertiesCfg( + drive_type="force", + **joint_params, + ) if with_default_eef: eef_joint_names = DEFAULT_EEF_HAND_JOINT_NAMES - drive_pros.stiffness.update( + joint_drive_props.stiffness.update( {eef_joint_names: DEFAULT_EEF_JOINT_DRIVE_PARAMS["stiffness"]} ) - drive_pros.damping.update( + joint_drive_props.damping.update( {eef_joint_names: DEFAULT_EEF_JOINT_DRIVE_PARAMS["damping"]} ) - drive_pros.max_effort.update( + joint_drive_props.max_effort.update( {eef_joint_names: DEFAULT_EEF_JOINT_DRIVE_PARAMS["max_effort"]} ) return { - "min_position_iters": 32, - "min_velocity_iters": 8, - "drive_pros": drive_pros, - "attrs": RigidBodyAttributesCfg( - static_friction=0.95, - dynamic_friction=0.9, - contact_offset=0.001, + "joint_drive_props": joint_drive_props, + "root_props": ArticulationRootPropertiesCfg( + min_position_iters=32, + min_velocity_iters=8, + ), + "attrs": RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg( + contact_offset=0.001, + rest_offset=0.0, + ), + material_props=RigidBodyMaterialCfg( + static_friction=0.95, + dynamic_friction=0.9, + ), ), } @@ -329,17 +351,37 @@ def build_pk_serial_chain( if __name__ == "__main__": - # Example usage - import numpy as np + import argparse np.set_printoptions(precision=5, suppress=True) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - - config = SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=4) + from embodichain.lab.sim.cfg import physics_cfg_for_backend + + parser = argparse.ArgumentParser(description="Launch the Dexforce W1 robot") + parser.add_argument( + "--physics", + choices=("default", "newton"), + default="newton", + help="Physics backend to launch (default: newton).", + ) + args = parser.parse_args() + + config = SimulationManagerCfg( + headless=True, + device="cpu", + num_envs=4, + physics_cfg=physics_cfg_for_backend(args.physics), + ) sim = SimulationManager(config) cfg = DexforceW1Cfg.from_dict({"uid": "dexforce_w1", "version": "v021"}) robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.update(step=1) - print("DexforceW1 robot added to the simulation.") + print("DexforceW1 robot added to the simulation.", flush=True) + sim.open_window() + from IPython import embed + + embed() # noqa: E702 + sim.destroy() diff --git a/embodichain/lab/sim/robots/dual_arm.py b/embodichain/lab/sim/robots/dual_arm.py index c42c85396..a6d73d5a8 100644 --- a/embodichain/lab/sim/robots/dual_arm.py +++ b/embodichain/lab/sim/robots/dual_arm.py @@ -265,7 +265,7 @@ def _resolve_base_cfg(base_robot: str | dict) -> RobotCfg: # --------------------------------------------------------------------------- # -def _mirror_drive_pros( +def _mirror_joint_drive_props( base_drive: JointDrivePropertiesCfg, name_case: dict[str, str] | None = None ) -> JointDrivePropertiesCfg: """Mirror a single-arm drive config across left/right arms. @@ -288,13 +288,14 @@ def _mirror_drive_pros( Returns: A fresh :class:`JointDrivePropertiesCfg` for the dual arm. """ - new = JointDrivePropertiesCfg(drive_type=base_drive.drive_type) - for prop in _DRIVE_PROPS: + new = type(base_drive)(drive_type=base_drive.drive_type) + properties = [*_DRIVE_PROPS, "target_mode"] + for prop in properties: val = getattr(base_drive, prop, None) if val is None: continue if isinstance(val, dict): - mirrored: Dict[str, float] = {} + mirrored: Dict[str, object] = {} for pattern, v in val.items(): mirrored[_prefixed_name(str(pattern), "left_", "joint", name_case)] = v mirrored[_prefixed_name(str(pattern), "right_", "joint", name_case)] = v @@ -404,13 +405,11 @@ def _populate_dual_cfg( ) cfg.solver_cfg = new_solver - cfg.drive_pros = _mirror_drive_pros(base_cfg.drive_pros, name_case) + cfg.joint_drive_props = _mirror_joint_drive_props( + base_cfg.joint_drive_props, name_case + ) cfg.attrs = base_cfg.attrs.copy() - cfg.min_position_iters = base_cfg.min_position_iters - cfg.min_velocity_iters = base_cfg.min_velocity_iters - cfg.fix_base = base_cfg.fix_base - cfg.disable_self_collision = base_cfg.disable_self_collision - cfg.sleep_threshold = base_cfg.sleep_threshold + cfg.root_props = base_cfg.root_props.copy() def build_dual_arm_cfg( @@ -455,7 +454,7 @@ class DualArmRobotCfg(RobotCfg): Two identical arms (the ``base_robot``) are mounted on a shared synthetic ``base_link``. The left/right ``control_parts``, per-arm ``solver_cfg`` and - mirrored ``drive_pros`` are derived automatically by + mirrored ``joint_drive_props`` are derived automatically by :func:`build_dual_arm_cfg`. Example: @@ -574,15 +573,27 @@ def build_pk_serial_chain( if __name__ == "__main__": + import argparse + np.set_printoptions(precision=5, suppress=True) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain.lab.sim.cfg import RenderCfg + from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend + + parser = argparse.ArgumentParser(description="Launch a dual-arm robot") + parser.add_argument( + "--physics", + choices=("default", "newton"), + default="default", + help="Physics backend to launch (default: default).", + ) + args = parser.parse_args() config = SimulationManagerCfg( headless=True, - sim_device="cpu", + device="cpu", num_envs=1, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg(renderer="fast-rt"), ) sim = SimulationManager(config) @@ -608,11 +619,9 @@ def build_pk_serial_chain( } ) robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.open_window() - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - # Round-trip check: from_dict(to_dict()) reproduces the cfg. cfg2 = DualArmRobotCfg.from_dict(cfg.to_dict()) assert cfg2.base_robot == cfg.base_robot diff --git a/embodichain/lab/sim/robots/franka_panda.py b/embodichain/lab/sim/robots/franka_panda.py index f80aeea96..8bedcb773 100644 --- a/embodichain/lab/sim/robots/franka_panda.py +++ b/embodichain/lab/sim/robots/franka_panda.py @@ -24,7 +24,6 @@ from embodichain.data import get_data_path from embodichain.lab.sim.cfg import ( JointDrivePropertiesCfg, - RigidBodyAttributesCfg, RobotCfg, URDFCfg, ) @@ -141,7 +140,8 @@ def _build_defaults(self, init_dict: dict | None = None) -> None: ), } - self.drive_pros = JointDrivePropertiesCfg( + self.joint_drive_props = JointDrivePropertiesCfg( + drive_type="force", stiffness={ "fr3_joint[1-7]": 1e4, "fr3_finger_joint[1-2]": 1e3, @@ -188,26 +188,36 @@ def build_pk_serial_chain( if __name__ == "__main__": + import argparse + np.set_printoptions(precision=5, suppress=True) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain.lab.sim.cfg import RenderCfg + from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend + + parser = argparse.ArgumentParser(description="Launch the Franka Panda robot") + parser.add_argument( + "--physics", + choices=("default", "newton"), + default="default", + help="Physics backend to launch (default: default).", + ) + args = parser.parse_args() config = SimulationManagerCfg( - headless=False, - sim_device="cpu", + headless=True, + device="cpu", num_envs=1, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg(renderer="hybrid"), ) sim = SimulationManager(config) cfg = FrankaPandaCfg.from_dict({"robot_type": "panda"}) robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.open_window() - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - from IPython import embed embed() # noqa: F401 diff --git a/embodichain/lab/sim/robots/ur_robot.py b/embodichain/lab/sim/robots/ur_robot.py index b6a4f9135..7d3e30edd 100644 --- a/embodichain/lab/sim/robots/ur_robot.py +++ b/embodichain/lab/sim/robots/ur_robot.py @@ -23,7 +23,6 @@ RobotCfg, URDFCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, ) from embodichain.lab.sim.solvers import URSolverCfg from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg @@ -139,7 +138,8 @@ def _build_defaults(self, init_dict: dict | None = None) -> None: ), } - self.drive_pros = JointDrivePropertiesCfg( + self.joint_drive_props = JointDrivePropertiesCfg( + drive_type="force", stiffness={"arm": 1e4}, damping={"arm": 1e3}, max_effort={"arm": _UR_MAX_EFFORT[robot_type]}, @@ -180,17 +180,27 @@ def build_pk_serial_chain( if __name__ == "__main__": - import numpy as np + import argparse np.set_printoptions(precision=5, suppress=True) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain.lab.sim.cfg import RenderCfg + from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend + + parser = argparse.ArgumentParser(description="Launch a Universal Robot") + parser.add_argument( + "--physics", + choices=("default", "newton"), + default="default", + help="Physics backend to launch (default: default).", + ) + args = parser.parse_args() config = SimulationManagerCfg( - headless=False, - sim_device="cpu", + headless=True, + device="cpu", num_envs=1, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg(renderer="fast-rt"), ) sim = SimulationManager(config) @@ -200,11 +210,9 @@ def build_pk_serial_chain( {"robot_type": "ur10e", "init_qpos": [0.0, -1.57, 1.57, -1.57, -1.57, 0.0]} ) robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.open_window() - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - from IPython import embed embed() # noqa: F401 diff --git a/embodichain/lab/sim/sensors/base_sensor.py b/embodichain/lab/sim/sensors/base_sensor.py index 3fb932f0d..a8e43866d 100644 --- a/embodichain/lab/sim/sensors/base_sensor.py +++ b/embodichain/lab/sim/sensors/base_sensor.py @@ -54,8 +54,8 @@ class OffsetCfg: pos: Tuple[float, float, float] = (0.0, 0.0, 0.0) """Position of the sensor in the parent frame. Defaults to (0.0, 0.0, 0.0).""" - quat: Tuple[float, float, float, float] = (1.0, 0.0, 0.0, 0.0) - """Orientation of the sensor in the parent frame as a quaternion (w, x, y, z). Defaults to (1.0, 0.0, 0.0, 0.0).""" + quat: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 1.0) + """Orientation in the parent frame as ``(x, y, z, w)``. Defaults to identity.""" parent: str | None = None """Name of the parent frame. If not specified, the sensor will be placed in the arena frame. @@ -171,10 +171,18 @@ class BaseSensor(BatchEntity): SUPPORTED_DATA_TYPES = [] def __init__( - self, config: SensorCfg, device: torch.device = torch.device("cpu") + self, + config: SensorCfg, + device: torch.device = torch.device("cpu"), + *, + num_instances: int | None = None, ) -> None: - - num_envs = get_dexsim_arena_num() + num_envs = ( + get_dexsim_arena_num() if num_instances is None else int(num_instances) + ) + if num_envs <= 0: + raise ValueError("A sensor requires at least one simulation instance.") + self._num_instances = num_envs self._data_buffer: TensorDict[str, torch.Tensor] = TensorDict( {}, batch_size=[num_envs], device=device ) @@ -186,7 +194,7 @@ def __init__( @cached_property def num_instances(self) -> int: - return get_dexsim_arena_num() + return self._num_instances @abstractmethod def _build_sensor_from_config( @@ -216,7 +224,8 @@ def get_arena_pose(self, to_matrix: bool = False) -> torch.Tensor: """Get the pose of the sensor in the arena frame. Args: - to_matrix: If True, return the pose as a 4x4 transformation matrix. + to_matrix: If True, return the pose as a 4x4 transformation matrix; + otherwise return ``(x, y, z, qx, qy, qz, qw)``. Returns: A tensor representing the pose of the sensor in the arena frame. diff --git a/embodichain/lab/sim/sensors/camera.py b/embodichain/lab/sim/sensors/camera.py index cc9a7aa44..3d9bc00d1 100644 --- a/embodichain/lab/sim/sensors/camera.py +++ b/embodichain/lab/sim/sensors/camera.py @@ -21,12 +21,15 @@ import dexsim.render as dr from functools import cached_property -from typing import List, Literal, Sequence, Tuple +from typing import TYPE_CHECKING, List, Literal, Sequence, Tuple from embodichain.lab.sim.sensors import BaseSensor, SensorCfg from embodichain.utils.math import matrix_from_quat, quat_from_matrix, look_at_to_pose from embodichain.utils import logger, configclass +if TYPE_CHECKING: + from embodichain.lab.sim.sim_manager import SimulationManager + __all__ = ["Camera", "CameraCfg"] @@ -134,27 +137,32 @@ class Camera(BaseSensor): SUPPORTED_DATA_TYPES = ["color", "depth", "mask", "normal", "position"] def __init__( - self, config: CameraCfg, device: torch.device = torch.device("cpu") + self, + config: CameraCfg, + device: torch.device = torch.device("cpu"), + *, + owner: SimulationManager, ) -> None: - super().__init__(config, device) + self._world = owner.get_world() + self._arenas = [owner.get_env(i) for i in range(owner.num_envs)] + if len(self._arenas) == 0: + raise ValueError("Camera requires at least one materialized Arena.") + self._camera_names: list[tuple[dexsim.environment.Arena, str]] = [] + self._is_attached = False + self._is_destroyed = False + super().__init__(config, device, num_instances=len(self._arenas)) + self.reset() def _build_sensor_from_config( self, config: CameraCfg, device: torch.device ) -> None: - self._world = dexsim.default_world() - env = self._world.get_env() - arenas = env.get_all_arenas() - if len(arenas) == 0: - arenas = [env] - num_instances = len(arenas) - self._frame_buffer = self._world.create_camera_group( - [config.width, config.height], num_instances, True + [config.width, config.height], self.num_instances, True ) view_attrib = config.get_view_attrib() - for i, arena in enumerate(arenas): - view_name = f"{self.uid}_view{i + 1}" + for i, arena in enumerate(self._arenas): + view_name = f"{config.uid}_view{i + 1}" view = arena.create_camera( view_name, config.width, @@ -167,6 +175,7 @@ def _build_sensor_from_config( view.set_near(config.near) view.set_far(config.far) self._entities[i] = view + self._camera_names.append((arena, view_name)) # Define a mapping of data types to their respective shapes and dtypes buffer_specs = { @@ -202,8 +211,6 @@ def _build_sensor_from_config( ) self.cfg: CameraCfg = config - if self.cfg.extrinsics.parent is not None: - self._attach_to_entity() @cached_property def group_id(self) -> int: @@ -216,12 +223,12 @@ def group_id(self) -> int: @property def is_attached(self) -> bool: - """Check if the camera is attached to a parent entity. + """Return whether all camera views are attached to parent nodes. Returns: - bool: True if the camera is attached to a parent entity, False otherwise. + True after parent attachment and extrinsics application succeed. """ - return self.cfg.extrinsics.parent is not None + return self._is_attached def update(self, **kwargs) -> None: """Update the sensor data. @@ -268,22 +275,28 @@ def update(self, **kwargs) -> None: self._frame_buffer.get_position_gpu_buffer().to(self.device)[..., :3] ) - def _attach_to_entity(self) -> None: - """Attach the sensor to the parent entity in each environment.""" - env = self._world.get_env() - for i, entity in enumerate(self._entities): + def attach_to_parent_nodes(self, parent_nodes: Sequence[object]) -> None: + """Attach camera views to one resolved parent node per environment. - parent = None - if i == 0: - parent = env.find_node(f"{self.cfg.extrinsics.parent}") - else: - parent = env.find_node(f"{self.cfg.extrinsics.parent}.{i-1}") - if parent is None: - logger.log_error( - f"Failed to find parent entity {self.cfg.extrinsics.parent} for sensor {self.cfg.uid}." - ) + Args: + parent_nodes: Parent render nodes ordered by environment index. + Raises: + RuntimeError: If the number of parent nodes does not match the + number of camera instances. + """ + nodes = list(parent_nodes) + if len(nodes) != self.num_instances: + raise RuntimeError( + f"Camera attachment received {len(nodes)} parent nodes for " + f"{self.num_instances} camera instances." + ) + for entity, parent in zip(self._entities, nodes, strict=True): entity.attach_node(parent) + # Extrinsics are expressed in the parent frame. Reapply them after + # reparenting because the camera was initially reset in Arena space. + self.reset() + self._is_attached = True def set_local_pose( self, pose: torch.Tensor, env_ids: Sequence[int] | None = None @@ -293,7 +306,8 @@ def set_local_pose( Note: The pose should be in the OpenGL coordinate system, which means the Y is up and Z is forward. Args: - pose (torch.Tensor): The local pose to set, should be a 4x4 transformation matrix. + pose (torch.Tensor): The local pose as ``(N, 4, 4)`` matrices or + ``(N, 7)`` vectors in ``(x, y, z, qx, qy, qz, qw)`` order. env_ids (Sequence[int] | None): The environment IDs to set the pose for. If None, set for all environments. """ if env_ids is None: @@ -320,7 +334,8 @@ def get_local_pose(self, to_matrix: bool = False) -> torch.Tensor: """Get the local pose of the camera. Args: - to_matrix (bool): If True, return the pose as a 4x4 matrix. If False, return as a quaternion. + to_matrix (bool): If True, return the pose as a 4x4 matrix. If + False, return ``(x, y, z, qx, qy, qz, qw)``. Returns: torch.Tensor: The local pose of the camera. @@ -341,19 +356,16 @@ def get_arena_pose(self, to_matrix: bool = False) -> torch.Tensor: """Get the pose of the sensor in the arena frame. Args: - to_matrix (bool): If True, return the pose as a 4x4 transformation matrix. + to_matrix (bool): If True, return the pose as a 4x4 transformation + matrix. If False, return ``(x, y, z, qx, qy, qz, qw)``. Returns: A tensor representing the pose of the sensor in the arena frame. """ - from embodichain.lab.sim.utility import get_dexsim_arenas - - arenas = get_dexsim_arenas() - poses = [] for i, entity in enumerate(self._entities): pose = entity.get_world_pose() - pose[:2, 3] -= arenas[i].get_root_node().get_local_pose()[:2, 3] + pose[:2, 3] -= self._arenas[i].get_root_node().get_local_pose()[:2, 3] poses.append(torch.as_tensor(pose, dtype=torch.float32)) poses = torch.stack(poses, dim=0).to(self.device) @@ -363,6 +375,28 @@ def get_arena_pose(self, to_matrix: bool = False) -> torch.Tensor: return torch.cat((xyz, quat), dim=-1) return poses + def destroy(self) -> None: + """Remove render cameras before releasing their World-owned group.""" + if self._is_destroyed: + return + self._is_destroyed = True + for arena, camera_name in self._camera_names: + try: + arena.remove_camera(camera_name) + except Exception as error: + logger.log_warning( + f"Failed to remove camera {camera_name!r}: {error!r}" + ) + self._entities = [] + self._camera_names = [] + # DexSim currently has no public remove_camera_group API. The group is + # World-owned; dropping this borrowed facade after removing all views + # is the narrowest safe lifetime boundary available to EmbodiChain. + self._frame_buffer = None + self._is_attached = False + self._arenas = [] + self._world = None + def look_at( self, eye: torch.Tensor, diff --git a/embodichain/lab/sim/sensors/contact_sensor.py b/embodichain/lab/sim/sensors/contact_sensor.py index 49ebbe1c8..1cdb82b43 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/sensors/stereo.py b/embodichain/lab/sim/sensors/stereo.py index 999bedca9..34b53bfce 100644 --- a/embodichain/lab/sim/sensors/stereo.py +++ b/embodichain/lab/sim/sensors/stereo.py @@ -21,13 +21,16 @@ import numpy as np import dexsim.render as dr -from typing import Dict, Tuple, List, Sequence +from typing import TYPE_CHECKING, Dict, Tuple, List, Sequence from dexsim.utility import inv_transform from embodichain.lab.sim.sensors import Camera, CameraCfg from embodichain.utils.math import matrix_from_euler from embodichain.utils import logger, configclass +if TYPE_CHECKING: + from embodichain.lab.sim.sim_manager import SimulationManager + @configclass class StereoCameraCfg(CameraCfg): @@ -155,8 +158,14 @@ def __init__( self, config: StereoCameraCfg, device: torch.device = torch.device("cpu"), + *, + owner: SimulationManager, ) -> None: - super().__init__(config, device) + super().__init__( + config, + device, + owner=owner, + ) # check valid config if self.cfg.enable_disparity and not self.cfg.enable_depth: @@ -165,21 +174,14 @@ def __init__( def _build_sensor_from_config( self, config: StereoCameraCfg, device: torch.device ) -> None: - self._world = dexsim.default_world() - env = self._world.get_env() - arenas = env.get_all_arenas() - if len(arenas) == 0: - arenas = [env] - num_instances = len(arenas) - self._frame_buffer = self._world.create_camera_group( - [config.width, config.height], num_instances * 2, True + [config.width, config.height], self.num_instances * 2, True ) view_attrib = config.get_view_attrib() left_list = [] right_list = [] - for i, arena in enumerate(arenas): - left_view_name = f"{self.uid}_left_view{i + 1}" + for i, arena in enumerate(self._arenas): + left_view_name = f"{config.uid}_left_view{i + 1}" left_view = arena.create_camera( left_view_name, config.width, @@ -192,9 +194,10 @@ def _build_sensor_from_config( left_view.set_near(config.near) left_view.set_far(config.far) left_list.append(left_view) + self._camera_names.append((arena, left_view_name)) - for i, arena in enumerate(arenas): - right_view_name = f"{self.uid}_right_view{i + 1}" + for i, arena in enumerate(self._arenas): + right_view_name = f"{config.uid}_right_view{i + 1}" right_view = arena.create_camera( right_view_name, config.width, @@ -207,8 +210,9 @@ def _build_sensor_from_config( right_view.set_near(config.near) right_view.set_far(config.far) right_list.append(right_view) + self._camera_names.append((arena, right_view_name)) - for i in range(num_instances): + for i in range(self.num_instances): self._entities[i] = PairCameraView( left_list[i], right_list[i], config.left_to_right.cpu().numpy() ) @@ -277,8 +281,6 @@ def _build_sensor_from_config( ][:, :, config.width :, :] self.cfg: CameraCfg = config - if self.cfg.extrinsics.parent is not None: - self._attach_to_entity() def update(self, **kwargs) -> None: """Update the sensor data. @@ -343,14 +345,10 @@ def get_left_right_arena_pose(self) -> torch.Tensor: Returns: torch.Tensor: The local pose of the left camera with shape (num_envs, 4, 4). """ - from embodichain.lab.sim.utility import get_dexsim_arenas - - arenas = get_dexsim_arenas() - left_poses = [] right_poses = [] for i, entity in enumerate(self._entities): - arena_pose = arenas[i].get_root_node().get_local_pose() + arena_pose = self._arenas[i].get_root_node().get_local_pose() left_pose = entity._left_view.get_world_pose() left_pose[:2, 3] -= arena_pose[:2, 3] left_poses.append( diff --git a/embodichain/lab/sim/shapes.py b/embodichain/lab/sim/shapes.py index 08e44f587..1fa7c606b 100755 --- a/embodichain/lab/sim/shapes.py +++ b/embodichain/lab/sim/shapes.py @@ -16,13 +16,278 @@ from __future__ import annotations -from typing import List, Dict, Union, TYPE_CHECKING, Any +import math +import warnings from dataclasses import MISSING +from numbers import Integral +from typing import Any, Dict, List, Literal, TYPE_CHECKING + from embodichain.utils import configclass, is_configclass, logger if TYPE_CHECKING: from embodichain.lab.sim.material import VisualMaterialCfg +__all__ = [ + "MeshCollisionApproximation", + "MeshCollisionCfg", + "LoadOption", + "ShapeCfg", + "MeshCfg", + "CubeCfg", + "SphereCfg", +] + + +MeshCollisionApproximation = Literal[ + "convex_hull", + "convex_decomposition", + "triangle_mesh", + "sdf", +] +"""Supported collision representations for a triangle mesh.""" + + +@configclass +class MeshCollisionCfg: + """Collision-geometry construction for :class:`MeshCfg`. + + The approximation is explicit. Strategy-specific fields are rejected when + they do not apply, so changing a numerical cooking value cannot silently + select a different collision representation. + """ + + approximation: MeshCollisionApproximation = "convex_hull" + """Collision representation built from the source triangle mesh.""" + + max_hulls: int | None = None + """Maximum hull count for ``convex_decomposition``; must be at least two.""" + + acd_method: Literal["coacd", "vhacd"] | None = None + """Approximate-convex-decomposition implementation.""" + + sdf_resolution: int | None = None + """Maximum SDF grid resolution; valid only for the ``sdf`` strategy.""" + + is_hydroelastic: bool | None = None + """Whether Newton uses the generated SDF for hydroelastic contact.""" + + sdf_narrow_band_range: tuple[float, float] | None = None + """Inner and outer signed-distance limits of the Newton SDF band [m].""" + + sdf_target_voxel_size: float | None = None + """Target Newton sparse-SDF voxel size [m], alternative to resolution.""" + + sdf_texture_format: Literal["uint16", "float32", "uint8"] | None = None + """Newton SDF voxel storage format.""" + + sdf_padding: float | None = None + """Extra padding used while Newton builds the mesh SDF [m].""" + + @property + def max_convex_hull_num(self) -> int: + """Deprecated compatibility view of :attr:`max_hulls`.""" + return self.max_hulls or 1 + + def __post_init__(self) -> None: + """Validate strategy-specific mesh-cooking fields.""" + supported = { + "convex_hull", + "convex_decomposition", + "triangle_mesh", + "sdf", + } + if self.approximation not in supported: + raise ValueError( + "MeshCollisionCfg.approximation must be one of " + f"{sorted(supported)}, got {self.approximation!r}." + ) + + if self.approximation == "convex_decomposition": + if ( + not isinstance(self.max_hulls, Integral) + or isinstance(self.max_hulls, bool) + or self.max_hulls < 2 + ): + raise ValueError( + "convex_decomposition requires max_hulls to be an integer " + "of at least 2." + ) + if self.acd_method not in (None, "coacd", "vhacd"): + raise ValueError("acd_method must be 'coacd' or 'vhacd'.") + elif self.max_hulls is not None or self.acd_method is not None: + raise ValueError( + "max_hulls and acd_method are valid only for convex_decomposition." + ) + + sdf_values = { + "sdf_resolution": self.sdf_resolution, + "is_hydroelastic": self.is_hydroelastic, + "sdf_narrow_band_range": self.sdf_narrow_band_range, + "sdf_target_voxel_size": self.sdf_target_voxel_size, + "sdf_texture_format": self.sdf_texture_format, + "sdf_padding": self.sdf_padding, + } + configured_sdf_fields = [ + name for name, value in sdf_values.items() if value is not None + ] + if self.approximation != "sdf" and configured_sdf_fields: + raise ValueError( + f"{configured_sdf_fields} are valid only for the sdf approximation." + ) + if self.sdf_resolution is not None and ( + not isinstance(self.sdf_resolution, Integral) + or isinstance(self.sdf_resolution, bool) + or self.sdf_resolution <= 0 + ): + raise ValueError("sdf_resolution must be a positive integer.") + if self.sdf_target_voxel_size is not None and ( + not math.isfinite(self.sdf_target_voxel_size) + or self.sdf_target_voxel_size <= 0.0 + ): + raise ValueError("sdf_target_voxel_size must be finite and positive.") + if self.sdf_resolution is not None and self.sdf_target_voxel_size is not None: + raise ValueError( + "Configure only one of sdf_resolution and sdf_target_voxel_size." + ) + if self.sdf_padding is not None and ( + not math.isfinite(self.sdf_padding) or self.sdf_padding < 0.0 + ): + raise ValueError("sdf_padding must be finite and non-negative.") + if self.is_hydroelastic is not None and not isinstance( + self.is_hydroelastic, bool + ): + raise TypeError("is_hydroelastic must be a boolean when configured.") + if self.sdf_texture_format not in (None, "uint16", "float32", "uint8"): + raise ValueError( + "sdf_texture_format must be 'uint16', 'float32', or 'uint8'." + ) + if self.sdf_narrow_band_range is not None: + if len(self.sdf_narrow_band_range) != 2: + raise ValueError("sdf_narrow_band_range must contain two values.") + inner, outer = (float(value) for value in self.sdf_narrow_band_range) + if not math.isfinite(inner) or not math.isfinite(outer): + raise ValueError("sdf_narrow_band_range values must be finite.") + if inner > outer: + raise ValueError( + "sdf_narrow_band_range inner value cannot exceed the outer value." + ) + + @classmethod + def from_dict(cls, init_dict: Dict[str, Any]) -> MeshCollisionCfg: + """Parse a mesh-collision mapping, including deprecated field names.""" + data = dict(init_dict) + legacy_fields = { + "max_convex_hull_num", + "force_sdf", + "sdf_max_resolution", + } + has_legacy_fields = bool(legacy_fields.intersection(data)) + if has_legacy_fields: + warnings.warn( + "Legacy mesh collision fields are deprecated; use an explicit " + "approximation with max_hulls or sdf_resolution.", + DeprecationWarning, + stacklevel=2, + ) + + legacy_max_hulls = data.pop("max_convex_hull_num", None) + legacy_force_sdf = data.pop("force_sdf", None) + legacy_sdf_resolution = data.pop("sdf_max_resolution", None) + if legacy_sdf_resolution is not None: + if "sdf_resolution" in data: + raise ValueError( + "sdf_max_resolution and sdf_resolution cannot both be configured." + ) + data["sdf_resolution"] = legacy_sdf_resolution + + if "approximation" not in data and has_legacy_fields: + sdf_requested = bool(legacy_force_sdf) or ( + data.get("sdf_resolution") is not None + and int(data["sdf_resolution"]) > 0 + ) + if sdf_requested: + data["approximation"] = "sdf" + data.pop("max_hulls", None) + data.pop("acd_method", None) + elif legacy_max_hulls is not None and int(legacy_max_hulls) > 1: + data["approximation"] = "convex_decomposition" + data["max_hulls"] = int(legacy_max_hulls) + else: + data["approximation"] = "convex_hull" + data.pop("acd_method", None) + elif legacy_max_hulls is not None: + if "max_hulls" in data: + raise ValueError( + "max_convex_hull_num and max_hulls cannot both be configured." + ) + data["max_hulls"] = int(legacy_max_hulls) + + if data.get("sdf_resolution") == 0: + data.pop("sdf_resolution") + return cls(**data) + + +_mesh_collision_cfg_init = MeshCollisionCfg.__init__ + + +def _mesh_collision_cfg_init_with_legacy_max_hulls( + self: MeshCollisionCfg, + approximation: MeshCollisionApproximation | None = None, + max_hulls: int | None = None, + acd_method: Literal["coacd", "vhacd"] | None = None, + sdf_resolution: int | None = None, + is_hydroelastic: bool | None = None, + sdf_narrow_band_range: tuple[float, float] | None = None, + sdf_target_voxel_size: float | None = None, + sdf_texture_format: Literal["uint16", "float32", "uint8"] | None = None, + sdf_padding: float | None = None, + *, + max_convex_hull_num: int | None = None, +) -> None: + """Initialize with the deprecated hull-count spelling at the API boundary.""" + if max_convex_hull_num is not None: + warnings.warn( + "max_convex_hull_num is deprecated; use max_hulls with an explicit " + "approximation.", + DeprecationWarning, + stacklevel=2, + ) + if max_hulls is not None: + raise ValueError( + "max_convex_hull_num and max_hulls cannot both be configured." + ) + if ( + not isinstance(max_convex_hull_num, Integral) + or isinstance(max_convex_hull_num, bool) + or max_convex_hull_num < 1 + ): + raise ValueError("max_convex_hull_num must be a positive integer.") + if approximation is None: + approximation = ( + "convex_decomposition" if max_convex_hull_num > 1 else "convex_hull" + ) + max_hulls = ( + None + if approximation == "convex_hull" and max_convex_hull_num == 1 + else max_convex_hull_num + ) + + _mesh_collision_cfg_init( + self, + approximation="convex_hull" if approximation is None else approximation, + max_hulls=max_hulls, + acd_method=acd_method, + sdf_resolution=sdf_resolution, + is_hydroelastic=is_hydroelastic, + sdf_narrow_band_range=sdf_narrow_band_range, + sdf_target_voxel_size=sdf_target_voxel_size, + sdf_texture_format=sdf_texture_format, + sdf_padding=sdf_padding, + ) + + +MeshCollisionCfg.__init__ = _mesh_collision_cfg_init_with_legacy_max_hulls + @configclass class LoadOption: @@ -70,13 +335,35 @@ def from_dict(cls, init_dict: Dict[str, Any]) -> ShapeCfg: """Initialize the configuration from a dictionary.""" from embodichain.utils.utility import get_class_instance - if "shape_type" not in init_dict: + data = dict(init_dict) + if "shape_type" not in data: logger.log_error("shape type must be specified in the configuration.") cfg = get_class_instance( - "embodichain.lab.sim.shapes", init_dict["shape_type"] + "Cfg" + "embodichain.lab.sim.shapes", data["shape_type"] + "Cfg" )() - for key, value in init_dict.items(): + legacy_mesh_fields = { + "max_convex_hull_num", + "acd_method", + "sdf_resolution", + } + if isinstance(cfg, MeshCfg): + configured_legacy = legacy_mesh_fields.intersection(data) + if configured_legacy: + if data.get("collision") is not None: + raise ValueError( + "MeshCfg collision cannot be combined with deprecated flat " + f"mesh fields {sorted(configured_legacy)}." + ) + legacy_collision = { + key: data.pop(key) for key in tuple(configured_legacy) + } + # Route through the legacy normalizer. Presence of this old hull + # name also makes the deprecation warning deterministic. + legacy_collision.setdefault("max_convex_hull_num", 1) + data["collision"] = legacy_collision + + for key, value in data.items(): if hasattr(cfg, key): attr = getattr(cfg, key) if key == "visual_material" and isinstance(value, dict): @@ -87,6 +374,15 @@ def from_dict(cls, init_dict: Dict[str, Any]) -> ShapeCfg: key, VisualMaterialCfg.from_dict(value), ) + elif key == "collision" and isinstance(cfg, MeshCfg): + if value is not None and not isinstance(value, MeshCollisionCfg): + if not isinstance(value, dict): + raise TypeError( + "MeshCfg.collision must be a mapping, " + "MeshCollisionCfg, or None." + ) + value = MeshCollisionCfg.from_dict(value) + setattr(cfg, key, value) elif is_configclass(attr): setattr(cfg, key, attr.from_dict(value)) else: @@ -119,28 +415,12 @@ class MeshCfg(ShapeCfg): project_direction: List[float] = [1.0, 1.0, 1.0] """Direction to project the UV coordinates. Defaults to [1.0, 1.0, 1.0].""" - max_convex_hull_num: int = 1 - """The maximum number of convex hulls that will be created for the mesh. - - If set to larger than 1, the mesh will be decomposed into multiple convex hulls - using the approximate convex decomposition method specified by :attr:`acd_method`. - Reference: https://github.com/SarahWeiii/CoACD - """ - - acd_method: str = "coacd" - """The method used for approximate convex decomposition (ACD) of the mesh. - - Currently, ``"coacd"`` and ``"vhacd"`` are supported. Only used when - :attr:`max_convex_hull_num` is set to larger than 1. - """ - - sdf_resolution: int = 0 - """Resolution for the signed distance field (SDF) of the mesh. + collision: MeshCollisionCfg | None = None + """Optional collision representation and cooking parameters. - The spacing of the uniformly sampled SDF is equal to the largest AABB extent - of the mesh, divided by the resolution. If ``sdf_resolution`` is set to larger - than 0, an SDF will be generated for collision detection. SDF increases the - accuracy of collision, but also takes more time to initialize and simulate. + ``None`` uses a single convex hull. Mesh collision construction belongs to + the geometry because it cannot be applied meaningfully to primitive shapes + or to articulation links without a named source-shape overlay. """ diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index b044575f6..a49273b03 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -22,17 +22,17 @@ import queue import time import threading +from contextlib import contextmanager import dexsim import torch import numpy as np import warp as wp -from tqdm import tqdm from pathlib import Path from copy import deepcopy from datetime import datetime -from functools import cached_property -from typing import TYPE_CHECKING, Callable, Dict, List, Sequence, Union +from functools import cached_property, partial +from typing import TYPE_CHECKING, Callable, Dict, Iterator, List, Sequence, Union from dataclasses import dataclass, asdict, field, MISSING # Global cache directories @@ -41,24 +41,29 @@ CONVEX_DECOMP_DIR = SIM_CACHE_DIR / "convex_decomposition" REACHABLE_XPOS_DIR = SIM_CACHE_DIR / "robot_reachable_xpos" + +def _is_usd_path(path: object | None) -> bool: + """Return whether a source path is a USD stage.""" + return path is not None and str(path).lower().endswith((".usd", ".usda", ".usdc")) + + from dexsim.types import ( + ActorType, Backend, ThreadMode, - PhysicalAttr, - ActorType, - RigidBodyShape, - RigidBodyGPUAPIReadType, - ArticulationGPUAPIReadType, ) from dexsim.core import TASK_RETURN -from dexsim.engine import CudaArray, Material +from dexsim.engine import Material from dexsim.models import MeshObject -from dexsim.render import Light as _Light, LightType, Windows +from dexsim.render import LightType, Windows from dexsim.engine import GizmoController, ObjectManipulator from embodichain.lab.sim.objects import ( RigidObject, RigidObjectGroup, + DeformableObject, + SurfaceDeformableObject, + VolumeDeformableObject, SoftObject, ClothObject, Articulation, @@ -76,27 +81,56 @@ ) from embodichain.lab.sim.cfg import ( RenderCfg, - PhysicsCfg, - MarkerCfg, + PhysicsBackendCfg, GPUMemoryCfg, + DefaultPhysicsCfg, + NewtonPhysicsCfg, + validate_physics_cfg, + MarkerCfg, WindowRecordCfg, WindowCameraPoseCfg, LightCfg, RigidObjectCfg, + DeformableObjectCfg, + SurfaceDeformableObjectCfg, + VolumeDeformableObjectCfg, SoftObjectCfg, ClothObjectCfg, RigidObjectGroupCfg, ArticulationCfg, + ArticulationRootPropertiesCfg, RobotCfg, + RobotPresetCfg, RigidConstraintCfg, ) +from embodichain.lab.sim.physics import NewtonPhysicsBackend, make_physics_backend +from embodichain.lab.sim.spawn.descriptors import ( + articulation_desc_from_cfg, + configure_articulation_desc, + rigid_desc_from_cfg, + surface_deformable_desc_from_cfg, + volume_deformable_desc_from_cfg, +) +from embodichain.lab.sim.spawn.usd import ( + articulation_desc_from_usd, + rigid_desc_from_usd, +) +from embodichain.lab.sim.spawn.scene import SpawnScene from embodichain.lab.sim import VisualMaterial, VisualMaterialCfg from embodichain.lab.sim.profiler import Profiler, ProfilerCfg from embodichain.lab.visualization.cfg import VisualizationCfg from embodichain.utils import configclass, logger -from embodichain.utils.math import look_at_to_pose, matrix_from_quat, pose_inv +from embodichain.utils.math import ( + convert_quat, + look_at_to_pose, + matrix_from_quat, + pose_inv, +) if TYPE_CHECKING: + from dexsim.engine import PhysicsScene + from dexsim.spawn import SpawnResult + from embodichain.lab.visualization import ( RuntimeHealth, RuntimeStats, @@ -108,6 +142,7 @@ __all__ = [ "SimulationManager", "SimulationManagerCfg", + "get_physics_scene", "SIM_CACHE_DIR", "MATERIAL_CACHE_DIR", "CONVEX_DECOMP_DIR", @@ -115,10 +150,131 @@ ] +@contextmanager +def _temporary_warp_kernel_log_suppression( + physics_cfg: PhysicsBackendCfg, +) -> Iterator[None]: + """Temporarily suppress informational Warp logs for Newton operations.""" + if not ( + isinstance(physics_cfg, NewtonPhysicsCfg) + and physics_cfg.suppress_warp_kernel_logs + ): + yield + return + + previous_log_level = wp.config.log_level + try: + # Warp emits its startup banner and module-load timers at INFO level. + # Keep warnings and errors visible. + wp.config.log_level = wp.LOG_WARNING + yield + finally: + wp.config.log_level = previous_log_level + + +def _initialize_warp_runtime(physics_cfg: PhysicsBackendCfg) -> None: + """Initialize Warp while honoring Newton startup-log suppression.""" + with _temporary_warp_kernel_log_suppression(physics_cfg): + wp.init() + + +# Deformable implementations remain backend-specific even though their public +# object/data contract is shared. Newton is an explicit empty placeholder until +# its native object adapters are integrated and validated. +_DEFORMABLE_BACKEND_IMPLEMENTATIONS = { + "default": { + "volume": ( + VolumeDeformableObjectCfg, + VolumeDeformableObject, + volume_deformable_desc_from_cfg, + "soft_object", + ), + "surface": ( + SurfaceDeformableObjectCfg, + SurfaceDeformableObject, + surface_deformable_desc_from_cfg, + "cloth_object", + ), + }, + "newton": {}, +} + + @configclass 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: PhysicsBackendCfg | None = None, + sim_device: str | torch.device | None = None, + physics_config: DefaultPhysicsCfg | None = None, + gpu_memory_config: GPUMemoryCfg | None = None, + profiler: ProfilerCfg | None = None, + visualization: VisualizationCfg | None = None, + window_record: WindowRecordCfg | None = None, + window_camera_pose: WindowCameraPoseCfg | 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 + if physics_cfg is None: + physics_cfg = ( + DefaultPhysicsCfg() if physics_config is None else physics_config + ) + self.physics_cfg = physics_cfg + if gpu_memory_config is not None: + if not isinstance(self.physics_cfg, DefaultPhysicsCfg): + logger.log_error( + "gpu_memory_config is only supported by the default physics backend.", + ValueError, + ) + self.physics_cfg.gpu_memory = gpu_memory_config + self.profiler = profiler + self.visualization = ( + VisualizationCfg() if visualization is None else visualization + ) + self.window_record = ( + WindowRecordCfg() if window_record is None else window_record + ) + self.window_camera_pose = ( + WindowCameraPoseCfg() if window_camera_pose is None else window_camera_pose + ) + if physics_dt is not None: + self.physics_cfg.physics_dt = physics_dt + runtime_device = device if device is not None else sim_device + if runtime_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(runtime_device) + if isinstance(runtime_device, str) + else runtime_device + ) + if torch_device.type != "cpu": + self.physics_cfg.device = runtime_device + else: + self.physics_cfg.device = runtime_device + + self.__post_init__() + width: int = 1920 """The width of the simulation window.""" @@ -158,8 +314,8 @@ 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.""" + physics_cfg: PhysicsBackendCfg = field(default_factory=DefaultPhysicsCfg) + """Physics backend configuration (type selects default vs Newton backend).""" profiler: ProfilerCfg | None = None """Optional simulation profiler. ``None`` disables profiling. @@ -169,14 +325,6 @@ class SimulationManagerCfg: profiler instance composes with the environment's step/reset hierarchy. """ - 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.""" - window_record: WindowRecordCfg = field(default_factory=WindowRecordCfg) """Viewer window recording settings (hotkey, paths, FPS, memory budget).""" @@ -187,10 +335,63 @@ class SimulationManagerCfg: """Live browser visualization settings.""" def __post_init__(self) -> None: - """Apply visualization-dependent simulation defaults.""" + """Validate physics and apply visualization-dependent defaults.""" + validate_physics_cfg(self.physics_cfg) if self.visualization.backend == "viser": self.headless = True + @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 + + @property + def sim_device(self) -> str | torch.device: + """Legacy alias for :attr:`device`.""" + return self.device + + @sim_device.setter + def sim_device(self, value: str | torch.device) -> None: + self.device = value + + @property + def physics_config(self) -> PhysicsBackendCfg: + """Legacy alias for :attr:`physics_cfg`.""" + return self.physics_cfg + + @physics_config.setter + def physics_config(self, value: PhysicsBackendCfg) -> None: + validate_physics_cfg(value) + self.physics_cfg = value + + @property + def gpu_memory_config(self) -> GPUMemoryCfg | None: + """Legacy alias for the default backend GPU-memory configuration.""" + if not isinstance(self.physics_cfg, DefaultPhysicsCfg): + return None + return self.physics_cfg.gpu_memory + + @gpu_memory_config.setter + def gpu_memory_config(self, value: GPUMemoryCfg) -> None: + if not isinstance(self.physics_cfg, DefaultPhysicsCfg): + raise AttributeError( + "gpu_memory_config is unavailable for the Newton physics backend." + ) + self.physics_cfg.gpu_memory = value + @dataclass class _WindowRecordState: @@ -285,11 +486,18 @@ 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) self.profiler = Profiler(sim_config.profiler, self.device) - # Initialize warp runtime context before creating the world. - wp.init() + # Initialize Warp before creating the world. For Newton, honor the + # configured startup/kernel-log suppression from the very first init. + _initialize_warp_runtime(sim_config.physics_cfg) self._world: dexsim.World = dexsim.World(world_config) self._window: Windows | None = None @@ -316,14 +524,11 @@ def __init__( ) self._window_camera_pose_input_control: ObjectManipulator | None = None - 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) @@ -342,13 +547,22 @@ def __init__( self._rigid_objects: Dict[str, RigidObject] = dict() self._constraints: Dict[str, RigidConstraint] = dict() self._rigid_object_groups: Dict[str, RigidObjectGroup] = dict() - self._soft_objects: Dict[str, SoftObject] = dict() - self._cloth_objects: Dict[str, ClothObject] = dict() + self._deformable_objects: Dict[str, DeformableObject] = dict() self._articulations: Dict[str, Articulation] = dict() self._robots: Dict[str, Robot] = dict() self._sensors: Dict[str, BaseSensor] = dict() - self._lights: Dict[str, _Light] = dict() + self._pending_sensor_attachments: list[Camera] = [] + self._lights: Dict[str, Light] = dict() + + self._spawn_scene = SpawnScene( + self._world, + num_envs=sim_config.num_envs, + spacing=(sim_config.arena_space, sim_config.arena_space, 0.0), + ) + self._arenas = list(self._spawn_scene.builder.prepare_arenas()) + self._prepared_spawn_topology_revision = -1 + self._synced_spawn_render_topology_revision = -1 self._visualization_runtime = None self._visualization_overlays: SceneOverlays | None = None @@ -367,15 +581,16 @@ def __init__( self._init_sim_resources() - self._create_default_plane() + # The plane material and visibility are authored before declaration so + # both eager Default loading and deferred Newton loading see them. + self._spawn_default_plane_visibility = True + self._default_plane = None self.set_default_background() + self._declare_spawn_default_plane() self.set_default_global_lighting() # Set physics to manual update mode by default. self.set_manual_update(True) - self._build_multiple_arenas(sim_config.num_envs) - self.start_visualization() - if sim_config.headless is False: self._window = self._world.get_windows() @@ -475,13 +690,65 @@ def num_envs(self) -> int: Returns: int: number of arenas. """ - return len(self._arenas) if len(self._arenas) > 0 else 1 + return self.sim_config.num_envs + + @property + def spawn_result(self) -> "SpawnResult | None": + """Return the current SpawnResult, or ``None`` before first prepare.""" + spawn_scene = getattr(self, "_spawn_scene", None) + if spawn_scene is None or not spawn_scene.builder.is_finalized: + return None + return spawn_scene.builder.result @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 Default physics backend is active.""" + return self.physics.name == "default" + + @property + def is_newton_backend(self) -> bool: + """Whether the Newton physics backend is active.""" + return self.physics.name == "newton" + + @property + def _active_newton_solver_type(self) -> str | None: + """Return the resolved Newton solver without widening the base contract.""" + if isinstance(self.physics, NewtonPhysicsBackend): + return self.physics.solver_type + return None + + @property + def newton_manager(self): + """Compatibility accessor for the removed NewtonManager API. + + A non-Newton backend still returns ``None``. The Newton backend raises + an actionable error because Spawn owns its World-level runtime and no + independent NewtonManager exists. + """ + if not self.is_newton_backend: + logger.log_warning("Newton backend is not active.") + return None + return self.physics.newton_manager + + @property + def differentiable_runtime(self): + """Return the differentiable facade over the Spawn-owned Newton runtime.""" + if not self.is_newton_backend: + raise RuntimeError( + "differentiable_runtime requires the Newton physics backend." + ) + return self.physics.differentiable_runtime + @property def is_physics_manually_update(self) -> bool: return self._world.is_physics_manually_update() @@ -501,8 +768,7 @@ def asset_uids(self) -> List[str]: uid_list.extend(list(self._robots.keys())) uid_list.extend(list(self._rigid_objects.keys())) uid_list.extend(list(self._rigid_object_groups.keys())) - uid_list.extend(list(self._soft_objects.keys())) - uid_list.extend(list(self._cloth_objects.keys())) + uid_list.extend(list(self._deformable_objects.keys())) uid_list.extend(list(self._articulations.keys())) return uid_list @@ -566,6 +832,8 @@ def start_visualization(self) -> VisualizationRuntime | None: """Start the configured live visualizer and publish the current scene.""" if self.sim_config.visualization.backend == "none": return None + if getattr(self, "_spawn_scene", None) is not None: + self.prepare() if getattr(self, "is_window_opened", False): raise RuntimeError( "Cannot start the Viser backend while the native DexSim window " @@ -709,8 +977,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 ( @@ -725,15 +991,12 @@ def _convert_sim_config( sim_config.render_cfg.apply_to_dexsim_config(world_config) - 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." @@ -744,6 +1007,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: @@ -752,6 +1019,56 @@ def _init_sim_resources(self) -> None: self._default_resources = SimResources() + def prepare(self) -> None: + """Materialize declarations, bind state, and resolve sensor parents.""" + scene = self._spawn_scene + result = scene.builder.result + if ( + not scene.builder.is_finalized + or result is None + or result.needs_rebuild + or scene.builder.has_pending_changes + ): + result = scene.commit() + self._env = result.get_arena("default") + self._arenas = [result.get_arena(name) for name in scene.arena_names] + self.__dict__.pop("arena_offsets", None) + if self._default_plane is None: + self._bind_default_plane(scene.handles("default_plane")[0]) + + # Runtime readiness belongs to the SimulationManager. Keep this and + # facade binding outside the topology-change branch so a failed call + # remains retryable without rematerializing the scene. + scene.prepare_runtime_config(result) + self._prepare_spawn_runtime(result) + scene.bind() + self._sync_spawn_render_state(result) + + while self._pending_sensor_attachments: + sensor = self._pending_sensor_attachments[0] + self._attach_camera_parent(sensor) + self._pending_sensor_attachments.pop(0) + + def _prepare_spawn_runtime(self, result: dexsim.spawn.SpawnResult) -> None: + """Prepare backend runtime buffers for one Spawn topology revision.""" + topology_revision = int(result.topology_revision) + if getattr(self, "_prepared_spawn_topology_revision", -1) == topology_revision: + return + if self.is_default_backend and self.device.type == "cuda": + self._world.init_gpu_physics() + self._prepared_spawn_topology_revision = topology_revision + + def _sync_spawn_render_state(self, result: dexsim.spawn.SpawnResult) -> None: + """Publish newly bound state once for each Spawn topology revision.""" + topology_revision = int(result.topology_revision) + if ( + getattr(self, "_synced_spawn_render_topology_revision", -1) + == topology_revision + ): + return + self.physics.sync_render_state(result) + self._synced_spawn_render_topology_revision = topology_revision + def enable_physics(self, enable: bool) -> None: """Enable or disable physics simulation. @@ -769,31 +1086,78 @@ 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." + """Prepare the Spawn-owned physics runtime. + + This backwards-compatible alias now has the same backend-neutral + behavior as :meth:`prepare`. + """ + self.prepare() + + def finalize_newton_physics(self) -> None: + """Prepare the Spawn-owned physics runtime. + + This backwards-compatible alias now has the same backend-neutral + behavior as :meth:`prepare`. + """ + self.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.differentiable_runtime.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.differentiable_runtime.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. @@ -806,22 +1170,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. """ with self.profiler.section("sim_update", is_root=True): with self.profiler.section("gpu_physics_check"): - if self.is_use_gpu_physics and not self._is_initialized_gpu_physics: - logger.log_warning( - "Using GPU physics, but not initialized yet. " - "Forcing initialization." - ) - with self.profiler.section("gpu_physics_init"): - self.init_gpu_physics() + self.prepare() if self.is_physics_manually_update: with self.profiler.section("manual_update"): @@ -832,7 +1190,10 @@ def update(self, physics_dt: float | None = None, step: int = 10) -> None: with self.profiler.section("gizmo_update"): self.update_gizmos() with self.profiler.section("world_update"): - self._world.update(physics_dt) + with _temporary_warp_kernel_log_suppression( + self.sim_config.physics_cfg + ): + self._world.update(physics_dt) self._visualization_sim_step += 1 self._visualization_sim_time += physics_dt if ( @@ -951,6 +1312,14 @@ def visualize_point_cloud( def get_world(self) -> dexsim.World: return self._world + def get_physics_scene(self) -> "PhysicsScene": + """Return the Default backend's compatibility scene after Spawn preparation. + + Newton has no ``PhysicsScene`` facade and raises with guidance to use + :attr:`spawn_result` instead. + """ + return self.physics.get_scene() + def can_open_native_window(self) -> bool: """Return whether the native DexSim window may be opened. @@ -1009,32 +1378,6 @@ def close_window(self) -> None: self._window_camera_pose_input_control = None self.is_window_opened = False - def _build_multiple_arenas(self, num: int, space: float | None = None) -> None: - """Build multiple arenas in a grid pattern. - - This interface is used for vectorized simulation. - - Args: - num (int): number of arenas to build. - space (float | None, optional): The distance between each arena. Defaults to the arena_space in sim_config. - """ - - if space is None: - space = self.sim_config.arena_space - - if num <= 0: - logger.log_warning("Number of arenas must be greater than 0.") - return - - scene_grid_length = int(np.ceil(np.sqrt(num))) - - for i in range(num): - arena = self._env.add_arena(f"arena_{i}") - - id_x, id_y = i % scene_grid_length, i // scene_grid_length - arena.set_root_node_position([id_x * space, id_y * space, 0]) - self._arenas.append(arena) - def set_indirect_lighting(self, name: str) -> None: """Set indirect lighting. @@ -1064,23 +1407,67 @@ def set_emission_light( if intensity is not None: self._env.set_env_light_intensity(intensity) - def _create_default_plane(self): - default_length = 1000 - repeat_uv_size = int(default_length / 2) - self._default_plane = self._env.create_plane( - 0, default_length, repeat_uv_size, repeat_uv_size + def _declare_spawn_default_plane(self) -> None: + """Declare the global ground in the World's Spawn scene.""" + + from dexsim.spawn import ( + CollisionApproximation, + CollisionDesc, + DexsimCollisionDesc, + GeometryDesc, + NewtonCollisionDesc, + ObjectDesc, + RenderDesc, + RigidBodyPhysicsDesc, ) - self._default_plane.set_name("default_plane") - plane_collision = self._env.create_cube( - default_length, default_length, default_length / 10 + + default_length = 1000.0 + geometry = GeometryDesc.plane(default_length) + repeat_uv_size = default_length / 2.0 + render = RenderDesc.from_geometry( + geometry, + material=self._spawn_default_plane_material, + ) + render.uv_coords = np.asarray( + [ + [0.0, 0.0], + [repeat_uv_size, 0.0], + [repeat_uv_size, repeat_uv_size], + [0.0, repeat_uv_size], + ], + dtype=np.float32, + ) + collision = CollisionDesc.from_geometry( + geometry, + approximation=CollisionApproximation.NONE, + ) + collision.dexsim = DexsimCollisionDesc( + dynamic_friction=0.5, + static_friction=0.5, + ) + collision.newton = NewtonCollisionDesc(mu=0.5) + collision.render_source_index = 0 + descriptor = ObjectDesc( + name="default_plane", + renders=[render], + collisions=[collision], + physics=RigidBodyPhysicsDesc.static(), + per_env=False, + ) + + self._spawn_scene.declare( + "rigid_object", + "default_plane", + descriptor, ) - 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) + handles = self._spawn_scene.handles("default_plane") + if handles: + self._bind_default_plane(handles[0]) - # TODO: add default physics attributes for the plane. + def _bind_default_plane(self, plane: Any) -> None: + """Retain the spawned ground plane and apply its visibility.""" + self._default_plane = plane + plane.set_visible(self._spawn_default_plane_visibility) def set_default_global_lighting(self) -> None: """Set default global lighting for the scene. @@ -1097,7 +1484,6 @@ def set_default_background(self) -> None: """Set default background.""" mat_name = "plane_mat" - mat = None mat_path = self._default_resources.get_material_path("PlaneDark") color_texture = os.path.join(mat_path, "PlaneDark_2K_Color.jpg") roughness_texture = os.path.join(mat_path, "PlaneDark_2K_Roughness.jpg") @@ -1110,7 +1496,11 @@ def set_default_background(self) -> None: ) ) - self._default_plane.set_material(mat.get_instance("plane_mat").mat) + material = mat.get_instance("plane_mat").mat + # Consumed by _declare_spawn_default_plane(). Keeping the native + # material in the descriptor preserves the VisualMaterial registry + # used by visual randomization without forcing finalization. + self._spawn_default_plane_material = material self._visual_materials[mat_name] = mat def set_ground_plane_visibility(self, visible: bool) -> None: @@ -1119,10 +1509,10 @@ def set_ground_plane_visibility(self, visible: bool) -> None: Args: visible (bool): _description_ """ - if visible: - self._default_plane.set_visible(True) - else: - self._default_plane.set_visible(False) + self._spawn_default_plane_visibility = bool(visible) + if self._default_plane is None: + return + self._default_plane.set_visible(bool(visible)) def set_texture_cache( self, key: str, texture: Union[torch.Tensor, List[torch.Tensor]] @@ -1156,16 +1546,26 @@ def get_texture_cache( def get_asset( self, uid: str - ) -> Light | BaseSensor | Robot | RigidObject | Articulation | None: + ) -> ( + Light + | BaseSensor + | Robot + | RigidObject + | RigidObjectGroup + | DeformableObject + | Articulation + | None + ): """Get an asset by its UID. - The asset can be a light, sensor, robot, rigid object or articulation. + The asset can be a light, sensor, robot, rigid object, deformable, or + articulation. Args: uid (str): The UID of the asset. Returns: - Light | BaseSensor | Robot | RigidObject | Articulation | None: The asset instance if found, otherwise None. + The asset instance if found, otherwise ``None``. """ if uid in self._lights: return self._lights[uid] @@ -1177,17 +1577,14 @@ def get_asset( return self._rigid_objects[uid] if uid in self._rigid_object_groups: return self._rigid_object_groups[uid] - if uid in self._soft_objects: - return self._soft_objects[uid] - if uid in self._cloth_objects: - return self._cloth_objects[uid] + if uid in self._deformable_objects: + return self._deformable_objects[uid] if uid in self._articulations: return self._articulations[uid] logger.log_warning(f"Asset {uid} not found.") return None - # Light type string → dexsim LightType enum mapping _LIGHT_TYPE_MAP: dict[str, LightType] = { "point": LightType.POINT, "sun": LightType.SUN, @@ -1196,8 +1593,6 @@ def get_asset( "rect": LightType.RECT, "mesh": LightType.MESH, } - - # Light types that are created as a single global scene light (not per-environment). _GLOBAL_LIGHT_TYPES: tuple[str, ...] = ("sun", "direction") def add_light(self, cfg: LightCfg) -> Light: @@ -1222,7 +1617,7 @@ def add_light(self, cfg: LightCfg) -> Light: Light: The created light instance. Raises: - RuntimeError: If ``cfg.light_type`` is not one of the supported types. + ValueError: If ``cfg.light_type`` is not supported. """ if cfg.uid is None: uid = "light" @@ -1233,45 +1628,41 @@ def add_light(self, cfg: LightCfg) -> Light: if uid in self._lights: logger.log_error(f"Light {uid} already exists.") - light_type_str = cfg.light_type - light_type = self._LIGHT_TYPE_MAP.get(light_type_str) + light_type = self._LIGHT_TYPE_MAP.get(cfg.light_type) if light_type is None: - supported = ", ".join(self._LIGHT_TYPE_MAP.keys()) - logger.log_error( - f"Unsupported light type: '{light_type_str}'. " + supported = ", ".join(self._LIGHT_TYPE_MAP) + raise ValueError( + f"Unsupported light type {cfg.light_type!r}. " f"Supported types: {supported}." ) - # Validation warnings for type-specific constraints - if light_type_str == "mesh" and not cfg.mesh_path: + if cfg.light_type == "mesh" and not cfg.mesh_path: logger.log_warning( f"Mesh light '{uid}' has no mesh_path set. " f"Use set_mesh() to assign a MeshObject." ) - if light_type_str == "rect" and (cfg.rect_width <= 0 or cfg.rect_height <= 0): + if cfg.light_type == "rect" and (cfg.rect_width <= 0 or cfg.rect_height <= 0): logger.log_warning( f"Rect light '{uid}' has zero or negative dimensions " f"(width={cfg.rect_width}, height={cfg.rect_height})." ) if cfg.light_type in self._GLOBAL_LIGHT_TYPES: - # Global scene light: create a single instance on the root - # environment. Infinite-distance lights (sun, direction) are - # physically scene-global and should not be duplicated per arena. - light = self._env.create_light(uid, light_type) - batch_lights = Light(cfg=cfg, entities=[light]) + batch_lights = Light( + cfg=cfg, + entities=[self._env.create_light(uid, light_type)], + ) else: - # Per-environment batched light: one instance per arena. - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - light_list = [] - for i, env in enumerate(env_list): - light_name = f"{uid}_{i}" - light = env.create_light(light_name, light_type) - light_list.append(light) - batch_lights = Light(cfg=cfg, entities=light_list) + batch_lights = Light( + cfg=cfg, + entities=[ + arena.create_light(f"{uid}_{index}", light_type) + for index, arena in enumerate(self._arenas) + ], + ) self._lights[uid] = batch_lights - + self.notify_visualization_topology_changed() return batch_lights def get_light(self, uid: str) -> Light | None: @@ -1296,107 +1687,289 @@ def get_light_uid_list(self) -> List[str]: """ return list(self._lights.keys()) - def add_rigid_object( + def add_usd( self, - cfg: RigidObjectCfg, - ) -> RigidObject: - """Add a rigid object to the scene. + name: str, + file_path: str, + *, + pose: np.ndarray | None = None, + robot_cfgs: dict[str, RobotCfg] | None = None, + ) -> dict[str, RigidObject | Articulation | Robot]: + """Declare the supported entities in a USD scene. + + The returned facades are keyed by their USD prim paths. They remain in + declared state until :meth:`prepare` finalizes the shared Spawn scene, + then bind in place to the resulting DexSim handles. + + USD does not identify which articulations should expose EmbodiChain's + robot interface. Pass those explicitly through ``robot_cfgs``; all + other articulation descriptions become :class:`Articulation` objects. Args: - cfg (RigidObjectCfg): Configuration for the rigid object. + name: Name passed to DexSim's USD scene parser. + file_path: USD, USDA, or USDC file path. + pose: Optional scene-root transform. + robot_cfgs: Robot configurations keyed by USD prim path. These + provide robot-side metadata while physics remains authored by + the USD scene. Returns: - RigidObject: The added rigid object instance handle. + Supported EmbodiChain facades keyed by USD prim path. + + Raises: + RuntimeError: If called after the Spawn scene was finalized. """ - from embodichain.lab.sim.utility.sim_utils import ( - load_mesh_objects_from_cfg, - ) + if self.spawn_result is not None: + raise RuntimeError( + "add_usd() must be called before SimulationManager.prepare()." + ) - uid = cfg.uid - if uid is None: - logger.log_error("Rigid object uid must be specified.") - if uid in self._rigid_objects: - logger.log_error(f"Rigid object {uid} already exists.") + from dexsim.spawn import ArticulationDesc, MeshObjectDesc - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = load_mesh_objects_from_cfg( - cfg=cfg, - env_list=env_list, - cache_dir=self._convex_decomp_dir, + descriptors = self._spawn_scene.builder.add_usd( + name, + file_path, + pose=pose, + per_env=True, ) + assets: dict[str, RigidObject | Articulation | Robot] = {} + robot_cfgs = robot_cfgs or {} + + for descriptor in descriptors: + source_path = ( + descriptor.usd.prim_path + if descriptor.usd is not None and descriptor.usd.prim_path + else descriptor.name + ) - rigid_obj = RigidObject(cfg=cfg, entities=obj_list, device=self.device) + if type(descriptor) is MeshObjectDesc: + body_type = "static" + if descriptor.physics is not None: + body_type = { + ActorType.DYNAMIC: "dynamic", + ActorType.KINEMATIC: "kinematic", + ActorType.STATIC: "static", + }[descriptor.physics.actor_type] + cfg = RigidObjectCfg( + uid=descriptor.name, + init_local_pose=descriptor.pose.copy(), + body_type=body_type, + body_scale=tuple(float(value) for value in descriptor.body_scale), + asset_physics_mode="preserve", + ) + facade = RigidObject( + cfg=cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) - if cfg.shape.visual_material: - mat = self.create_visual_material(cfg.shape.visual_material) - rigid_obj.set_visual_material(mat, update_default=True) + self._spawn_scene.track( + "rigid_object", + descriptor.name, + descriptor, + facade=facade, + ) + self._rigid_objects[descriptor.name] = facade + assets[source_path] = facade + continue - self._rigid_objects[uid] = rigid_obj - self.notify_visualization_topology_changed() + if isinstance(descriptor, ArticulationDesc): + robot_cfg = robot_cfgs.get(source_path) + facade_type: type[Articulation] = ( + Robot if robot_cfg is not None else Articulation + ) + cfg = ( + deepcopy(robot_cfg) + if robot_cfg is not None + else ArticulationCfg(uid=descriptor.name) + ) + cfg.uid = descriptor.name + cfg.fpath = file_path + cfg.init_local_pose = descriptor.pose.copy() + cfg.asset_physics_mode = "preserve" + if robot_cfg is None: + cfg.root_props = ArticulationRootPropertiesCfg() + else: + cfg.root_props = cfg.root_props.copy() + cfg.root_props.fixed_base = bool(descriptor.fixed_base) + cfg.root_props.self_collision_enabled = descriptor.enable_self_collision + cfg.body_scale = tuple(float(value) for value in descriptor.body_scale) + cfg.build_pk_chain = False + facade = facade_type( + cfg=cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) - return rigid_obj + self._spawn_scene.track( + "articulation", + descriptor.name, + descriptor, + facade=facade, + ) + registry = ( + self._robots if robot_cfg is not None else self._articulations + ) + registry[descriptor.name] = facade + assets[source_path] = facade - def add_soft_object(self, cfg: SoftObjectCfg) -> SoftObject: - """Add a soft object to the scene. + self.notify_visualization_topology_changed() + return assets + + def add_rigid_object( + self, + cfg: RigidObjectCfg, + ) -> RigidObject: + """Add a rigid object to the scene. Args: - cfg (SoftObjectCfg): Configuration for the soft object. + cfg (RigidObjectCfg): Configuration for the rigid object. Returns: - SoftObject: The added soft object instance handle. + RigidObject: The added rigid object instance handle. """ - if not self.is_use_gpu_physics: - logger.log_error("Soft object requires GPU physics to be enabled.") - - from embodichain.lab.sim.utility import ( - load_soft_object_from_cfg, - ) - uid = cfg.uid if uid is None: - logger.log_error("Soft object uid must be specified.") + raise ValueError("Rigid object uid must be specified.") + if uid in self._rigid_objects: + raise ValueError(f"Rigid object {uid!r} already exists.") + source_path = getattr(cfg.shape, "fpath", None) + if _is_usd_path(source_path): + descriptor, materials = rigid_desc_from_usd( + cfg, + per_env=True, + newton_solver_type=self._active_newton_solver_type, + ) + else: + descriptor, materials = rigid_desc_from_cfg( + cfg, + per_env=True, + newton_solver_type=self._active_newton_solver_type, + ) + self._spawn_scene.builder.materials.update(materials) - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = load_soft_object_from_cfg( + rigid_obj = RigidObject( cfg=cfg, - env_list=env_list, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, ) - soft_obj = SoftObject(cfg=cfg, entities=obj_list, device=self.device) - self._soft_objects[uid] = soft_obj + was_materialized = self.spawn_result is not None + self._spawn_scene.declare( + "rigid_object", + uid, + descriptor, + facade=rigid_obj, + ) + self._rigid_objects[uid] = rigid_obj self.notify_visualization_topology_changed() - return soft_obj - def add_cloth_object(self, cfg: ClothObjectCfg) -> ClothObject: - """Add a cloth object to the scene. + # Preserve the legacy immediate-availability behavior for runtime + # additions. Initial environment construction still batches all + # declarations into one finalize at BaseEnv's prepare boundary. + if was_materialized: + self.prepare() + return rigid_obj + + def add_deformable_object(self, cfg: DeformableObjectCfg) -> DeformableObject: + """Declare a volume or surface deformable in the scene. + + DexSim is the only deformable implementation currently registered. + Backend capability flags and the dispatch boundary are intentionally + explicit so a future Newton adapter can be added without changing this + public method or its callers. Args: - cfg (ClothObjectCfg): Configuration for the cloth object. + cfg: Volume- or surface-deformable configuration. Returns: - ClothObject: The added cloth object instance handle. - """ - if not self.is_use_gpu_physics: - logger.log_error("Cloth object requires GPU physics to be enabled.") + The declared deformable facade. - from embodichain.lab.sim.utility import ( - load_cloth_object_from_cfg, - ) + Raises: + NotImplementedError: If the active backend or device cannot host + the requested deformable type. + ValueError: If the discriminator or UID is invalid. + """ + deformable_type = cfg.deformable_type + if deformable_type == "volume": + supported = self.physics.supports_volume_deformables + elif deformable_type == "surface": + supported = self.physics.supports_surface_deformables + else: + raise ValueError( + f"Unsupported deformable_type {deformable_type!r}; expected " + "'volume' or 'surface'." + ) + if not supported: + raise NotImplementedError( + f"The {self.physics.name} backend does not yet provide a " + f"{deformable_type}-deformable object adapter." + ) + if self.device.type != "cuda": + raise NotImplementedError( + "DexSim deformable objects currently require a CUDA device." + ) + if self.spawn_result is not None: + raise NotImplementedError( + "DexSim Spawn does not yet support adding deformables after " + "finalization." + ) uid = cfg.uid if uid is None: - logger.log_error("Cloth object uid must be specified.") + raise ValueError("Deformable object uid must be specified.") + if uid in self._deformable_objects: + raise ValueError(f"Deformable object {uid!r} already exists.") - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = load_cloth_object_from_cfg( - cfg=cfg, - env_list=env_list, + backend_implementations = _DEFORMABLE_BACKEND_IMPLEMENTATIONS.get( + self.physics.name ) + if not backend_implementations: + raise NotImplementedError( + f"No deformable implementation is registered for the " + f"{self.physics.name} backend." + ) - cloth_obj = ClothObject(cfg=cfg, entities=obj_list, device=self.device) - self._cloth_objects[uid] = cloth_obj + config_cls, object_cls, descriptor_factory, spawn_kind = ( + backend_implementations[deformable_type] + ) + if not isinstance(cfg, config_cls): + raise TypeError( + f"A {deformable_type} deformable requires " + f"{config_cls.__name__}, got {type(cfg).__name__}." + ) + descriptor, materials = descriptor_factory(cfg, per_env=True) + self._spawn_scene.builder.materials.update(materials) + deformable = object_cls( + cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) + self._spawn_scene.declare( + spawn_kind, + uid, + descriptor, + facade=deformable, + ) + self._deformable_objects[uid] = deformable self.notify_visualization_topology_changed() - return cloth_obj + return deformable + + def add_soft_object(self, cfg: SoftObjectCfg) -> SoftObject: + """Compatibility wrapper for adding a volume deformable.""" + deformable = self.add_deformable_object(cfg) + assert isinstance(deformable, VolumeDeformableObject) + return deformable + + def add_cloth_object(self, cfg: ClothObjectCfg) -> ClothObject: + """Compatibility wrapper for adding a surface deformable.""" + deformable = self.add_deformable_object(cfg) + assert isinstance(deformable, SurfaceDeformableObject) + return deformable def get_rigid_object(self, uid: str) -> RigidObject | None: """Get a rigid object by its unique ID. @@ -1412,33 +1985,28 @@ def get_rigid_object(self, uid: str) -> RigidObject | None: return None return self._rigid_objects[uid] - def get_soft_object(self, uid: str) -> SoftObject | None: - """Get a soft object by its unique ID. - - Args: - uid (str): The unique ID of the soft object. + def get_deformable_object(self, uid: str) -> DeformableObject | None: + """Get a deformable object by its unique ID.""" + if uid not in self._deformable_objects: + logger.log_warning(f"Deformable object {uid} not found.") + return None + return self._deformable_objects[uid] - Returns: - SoftObject | None: The soft object instance if found, otherwise None. - """ - if uid not in self._soft_objects: + def get_soft_object(self, uid: str) -> SoftObject | None: + """Get a volume deformable through the legacy soft-object API.""" + deformable = self._deformable_objects.get(uid) + if not isinstance(deformable, VolumeDeformableObject): logger.log_warning(f"Soft object {uid} not found.") return None - return self._soft_objects[uid] + return deformable def get_cloth_object(self, uid: str) -> ClothObject | None: - """Get a cloth object by its unique ID. - - Args: - uid (str): The unique ID of the cloth object. - - Returns: - ClothObject | None: The cloth object instance if found, otherwise None. - """ - if uid not in self._cloth_objects: + """Get a surface deformable through the legacy cloth-object API.""" + deformable = self._deformable_objects.get(uid) + if not isinstance(deformable, SurfaceDeformableObject): logger.log_warning(f"Cloth object {uid} not found.") return None - return self._cloth_objects[uid] + return deformable def get_rigid_object_uid_list(self) -> List[str]: """Get current rigid body uid list @@ -1455,20 +2023,7 @@ def _broadcast_frame( env_ids: Sequence[int], name: str, ) -> list[np.ndarray]: - """Broadcast a local-frame spec to one matrix per target env. - - Args: - frame: None -> identity; (4,4) -> repeated; (N,4,4) -> indexed per env. - num_envs: Total number of arenas (used to validate (N,4,4)). - env_ids: Target env indices to produce frames for. - name: Constraint name (for error messages). - - Returns: - A list of (4,4) numpy arrays, one per env in env_ids. - - Raises: - RuntimeError: If an (N,4,4) frame's N != num_envs, or shape is invalid. - """ + """Broadcast a local constraint frame to the selected environments.""" if frame is None: identity = np.eye(4, dtype=np.float32) return [identity for _ in env_ids] @@ -1518,15 +2073,11 @@ def create_rigid_constraint( cfg: RigidConstraintCfg, env_ids: Sequence[int] | torch.Tensor | None = None, ) -> RigidConstraint: - """Create a fixed constraint between two RigidObjects. + """Create a fixed constraint between two rigid objects. - Binds ``rigid_object_a``'s entity[i] to ``rigid_object_b``'s entity[i] - within arena[i], for each env in ``env_ids``. Local frames default to - welding the objects at their *current* relative pose: - ``local_frame_a`` defaults to identity (object A's origin) and - ``local_frame_b`` defaults to ``inv(pose_B) @ pose_A`` (computed per env), - so the offset is preserved rather than the two origins being pulled - together. Pass explicit frames to define a specific joint frame. + Constraints are native Default-backend resources owned by each Arena. + Spawn owns the two actors; this method only borrows their native actor + handles while creating the constraint. Args: cfg: The constraint configuration. @@ -1534,20 +2085,18 @@ def create_rigid_constraint( the :class:`EventManager`) or a sequence of ints. None -> all arenas. Returns: - The created :class:`RigidConstraint`. - - Raises: - RuntimeError: If either object is missing, the name is already in use, - a frame shape is invalid, or dexsim fails to create a handle. + The created constraint batch. """ - # validate constraint type (only fixed supported in v1) + if hasattr(self, "physics") and not self.is_default_backend: + raise NotImplementedError( + "Rigid constraints are currently supported only by the Default " + "backend." + ) if cfg.constraint_type != "fixed": logger.log_error( f"Constraint '{cfg.name}' has unsupported type " - f"'{cfg.constraint_type}'. Only 'fixed' is supported in v1." + f"'{cfg.constraint_type}'. Only 'fixed' is supported." ) - - # resolve objects if cfg.rigid_object_a_uid not in self._rigid_objects: logger.log_error( f"RigidObject '{cfg.rigid_object_a_uid}' not found for constraint " @@ -1558,16 +2107,16 @@ def create_rigid_constraint( f"RigidObject '{cfg.rigid_object_b_uid}' not found for constraint " f"'{cfg.name}'. Available: {list(self._rigid_objects.keys())}." ) - rigid_object_a = self._rigid_objects[cfg.rigid_object_a_uid] - rigid_object_b = self._rigid_objects[cfg.rigid_object_b_uid] - - # validate duplicate name if cfg.name in self._constraints: logger.log_error( f"Constraint '{cfg.name}' already exists. Remove it before recreating." ) - # validate object entity counts match num_envs + rigid_object_a = self._rigid_objects[cfg.rigid_object_a_uid] + rigid_object_b = self._rigid_objects[cfg.rigid_object_b_uid] + if hasattr(self, "_spawn_scene"): + self.prepare() + num_envs = self.num_envs if rigid_object_a.num_instances != num_envs: logger.log_error( @@ -1580,50 +2129,52 @@ def create_rigid_constraint( f"{rigid_object_b.num_instances} instances but num_envs is {num_envs}." ) - # resolve target env_ids (accepts None / tensor / sequence) target_env_ids = self._normalize_env_ids(env_ids, num_envs) - - # broadcast local frames. - # local_frame_a defaults to identity (object A's origin). - # local_frame_b defaults to the current relative pose of A w.r.t. B - # (inv(pose_B) @ pose_A), so that with both frames left as None the - # constraint welds the objects at their *current* relative pose instead - # of pulling their origins together. frames_a = self._broadcast_frame( cfg.local_frame_a, num_envs, target_env_ids, cfg.name ) if cfg.local_frame_b is None: pose_a = rigid_object_a.get_local_pose(to_matrix=True) pose_b = rigid_object_b.get_local_pose(to_matrix=True) - frame_b = torch.bmm(pose_inv(pose_b), pose_a) # (N, 4, 4) - frame_b = frame_b.cpu().numpy().astype(np.float32) + frame_b = ( + torch.bmm(pose_inv(pose_b), pose_a).cpu().numpy().astype(np.float32) + ) frames_b = [frame_b[i] for i in target_env_ids] else: frames_b = self._broadcast_frame( cfg.local_frame_b, num_envs, target_env_ids, cfg.name ) - # pre-size handles list with None, fill target envs handles: list = [None] * num_envs try: - for idx, env_id in enumerate(target_env_ids): + for index, env_id in enumerate(target_env_ids): + actor_a = rigid_object_a._entities[env_id] + actor_b = rigid_object_b._entities[env_id] + if getattr(rigid_object_a, "is_spawn_bound", False) is True: + actor_a = actor_a.native + if getattr(rigid_object_b, "is_spawn_bound", False) is True: + actor_b = actor_b.native + if actor_a is None or actor_b is None: + logger.log_error( + f"Constraint '{cfg.name}' references a released Spawn actor " + f"in environment {env_id}." + ) + arena = self.get_env(env_id) - name_i = cfg.name if num_envs <= 1 else f"{cfg.name}_{env_id}" + name = cfg.name if num_envs <= 1 else f"{cfg.name}_{env_id}" handle = arena.create_fixed_constraint( - name_i, - rigid_object_a._entities[env_id], - rigid_object_b._entities[env_id], - frames_a[idx], - frames_b[idx], + name, + actor_a, + actor_b, + frames_a[index], + frames_b[index], ) if handle is None: logger.log_error( - f"Failed to create constraint '{name_i}' in arena {env_id}." + f"Failed to create constraint '{name}' in arena {env_id}." ) handles[env_id] = handle except Exception: - # Ensure partially created per-arena constraints are removed if a later - # arena fails, so create/remove semantics stay consistent. RigidConstraint( cfg=cfg, constraint_handles=handles, @@ -1643,21 +2194,25 @@ def create_rigid_constraint( self._constraints[cfg.name] = constraint return constraint - def get_soft_object_uid_list(self) -> List[str]: - """Get current soft body uid list + def get_deformable_object_uid_list(self) -> List[str]: + """Return all deformable object UIDs in declaration order.""" + return list(self._deformable_objects.keys()) - Returns: - List[str]: list of soft body uid. - """ - return list(self._soft_objects.keys()) + def get_soft_object_uid_list(self) -> List[str]: + """Return volume-deformable UIDs through the legacy soft API.""" + return [ + uid + for uid, asset in self._deformable_objects.items() + if asset.deformable_type == "volume" + ] def get_cloth_object_uid_list(self) -> List[str]: - """Get current cloth body uid list - - Returns: - List[str]: list of cloth body uid. - """ - return list(self._cloth_objects.keys()) + """Return surface-deformable UIDs through the legacy cloth API.""" + return [ + uid + for uid, asset in self._deformable_objects.items() + if asset.deformable_type == "surface" + ] def remove_rigid_constraint( self, @@ -1719,43 +2274,74 @@ def add_rigid_object_group(self, cfg: RigidObjectGroupCfg) -> RigidObjectGroup: Args: cfg (RigidObjectGroupCfg): Configuration for the rigid object group. - """ - from embodichain.lab.sim.utility.sim_utils import ( - load_mesh_objects_from_cfg, - ) + Returns: + The stable Group facade. During initial scene construction it is + bound to Spawn handles by :meth:`prepare`. + """ + if not self.physics.supports_rigid_object_group: + raise NotImplementedError( + f"The {self.physics.name} backend does not support rigid object groups." + ) uid = cfg.uid if uid is None: - logger.log_error("Rigid object group uid must be specified.") + raise ValueError("Rigid object group uid must be specified.") if uid in self._rigid_object_groups: - logger.log_error(f"Rigid object group {uid} already exists.") - + raise ValueError(f"Rigid object group {uid!r} already exists.") if cfg.body_type == "static": - logger.log_error("Rigid object group cannot be static.") - - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - - obj_group_list = [] - for key, rigid_cfg in tqdm( - cfg.rigid_objects.items(), desc="Loading rigid objects" - ): - obj_list = load_mesh_objects_from_cfg( - cfg=rigid_cfg, - env_list=env_list, - cache_dir=self._convex_decomp_dir, - ) - obj_group_list.append(obj_list) + raise ValueError("Rigid object group cannot be static.") + if not cfg.rigid_objects: + raise ValueError("Rigid object group must contain at least one object.") + + actor_type = { + "dynamic": ActorType.DYNAMIC, + "kinematic": ActorType.KINEMATIC, + }[cfg.body_type] + descriptors = [] + for index, member in enumerate(cfg.rigid_objects.values()): + member_cfg = deepcopy(member) + member_cfg.uid = f"{uid}__member_{index}" + member_cfg.body_type = cfg.body_type + source_path = getattr(member_cfg.shape, "fpath", None) + if _is_usd_path(source_path): + descriptor, materials = rigid_desc_from_usd( + member_cfg, + per_env=True, + newton_solver_type=self._active_newton_solver_type, + ) + else: + descriptor, materials = rigid_desc_from_cfg( + member_cfg, + per_env=True, + newton_solver_type=self._active_newton_solver_type, + ) + if descriptor.physics is None: + raise ValueError( + f"Rigid object group member {index} has no rigid-body physics." + ) + descriptor.physics.actor_type = actor_type + self._spawn_scene.builder.materials.update(materials) + descriptors.append(descriptor) - # Convert [a1, a2, ...], [b1, b2, ...] to [(a1, b1, ...), (a2, b2, ...), ...] - obj_group_list = list(zip(*obj_group_list)) - rigid_obj_group = RigidObjectGroup( - cfg=cfg, entities=obj_group_list, device=self.device + group = RigidObjectGroup( + cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, ) - self._rigid_object_groups[uid] = rigid_obj_group + was_materialized = self.spawn_result is not None + self._spawn_scene.declare( + "rigid_object_group", + uid, + tuple(descriptors), + facade=group, + ) + self._rigid_object_groups[uid] = group self.notify_visualization_topology_changed() - - return rigid_obj_group + if was_materialized: + self.prepare() + return group def get_rigid_object_group(self, uid: str) -> RigidObjectGroup | None: """Get a rigid object group by its unique ID. @@ -1824,54 +2410,23 @@ def add_articulation( Returns: Articulation: The added articulation instance handle. """ - uid = cfg.uid if uid is None: + if cfg.fpath is None: + raise ValueError( + "Articulation configuration must provide fpath when uid " + "is not specified." + ) uid = os.path.splitext(os.path.basename(cfg.fpath))[0] cfg.uid = uid if uid in self._articulations: - logger.log_error(f"Articulation {uid} already exists.") - - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = [] - - is_usd = cfg.fpath.endswith((".usd", ".usda", ".usdc")) - if is_usd: - # 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 - ) - # 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]) - 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) - - articulation = Articulation(cfg=cfg, entities=obj_list, device=self.device) + raise ValueError(f"Articulation {uid!r} already exists.") + was_materialized = self.spawn_result is not None + articulation = self._declare_spawn_articulation(cfg, Articulation) self._articulations[uid] = articulation - self.notify_visualization_topology_changed() - + if was_materialized: + self.prepare() return articulation def get_articulation(self, uid: str) -> Articulation | None: @@ -1896,15 +2451,29 @@ def get_articulation_uid_list(self) -> List[str]: """ return list(self._articulations.keys()) - def add_robot(self, cfg: RobotCfg) -> Robot | None: + def add_robot(self, cfg: RobotCfg | RobotPresetCfg) -> Robot | None: """Add a Robot to the scene. Args: - cfg (RobotCfg): Configuration for the robot. + cfg: A concrete robot configuration or a replace-only backend + preset. Presets are resolved from ``physics_cfg`` before the + robot is declared. 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, + ) + + if isinstance(cfg, RobotPresetCfg): + cfg = cfg.resolve( + self.sim_config.physics_cfg, + newton_solver_type=self._active_newton_solver_type, + ) uid = cfg.uid if cfg.fpath is None: @@ -1929,45 +2498,61 @@ def add_robot(self, cfg: RobotCfg) -> Robot | None: logger.log_error(f"Robot {uid} already exists.") return self._robots[uid] - env_list = [self._env] if len(self._arenas) == 0 else self._arenas - obj_list = [] - - 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]) - else: - # non-usd file does not support this option, will be forced set False to avoid potential issues. - cfg.use_usd_properties = False + was_materialized = self.spawn_result is not None + robot = self._declare_spawn_articulation(cfg, Robot) + self._robots[uid] = robot + if was_materialized: + self.prepare() + return robot - for env in env_list: - art = env.load_urdf(cfg.fpath) - obj_list.append(art) + def _declare_spawn_articulation( + self, + cfg: ArticulationCfg, + facade_type: type[Articulation], + ) -> Articulation: + """Declare an articulation facade and bind its Batch after finalize. + + DexSim remains the sole articulation source loader. EmbodiChain applies + regex/group configuration to the resolved descriptor before either + backend materializes it. Runtime Batch data is created at the shared + prepare boundary. + """ + if _is_usd_path(cfg.fpath): + descriptor, materials = articulation_desc_from_usd( + cfg, + per_env=True, + newton_solver_type=self._active_newton_solver_type, + ) + self._spawn_scene.builder.materials.update(materials) + else: + descriptor = articulation_desc_from_cfg( + cfg, + per_env=True, + newton_solver_type=self._active_newton_solver_type, + ) + if cfg.uid is None: + cfg.uid = descriptor.name - robot = Robot(cfg=cfg, entities=obj_list, device=self.device) + facade = facade_type( + cfg=cfg, + entities=None, + device=self.device, + declared_num_instances=self.sim_config.num_envs, + ) - self._robots[uid] = robot + self._spawn_scene.declare( + "articulation", + descriptor.name, + descriptor, + facade=facade, + configure_source=partial( + configure_articulation_desc, + cfg=cfg, + newton_solver_type=self._active_newton_solver_type, + ), + ) self.notify_visualization_topology_changed() - - return robot + return facade def get_robot(self, uid: str) -> Robot | None: """Get a Robot by its unique ID. @@ -2196,17 +2781,20 @@ def process_visualization_commands(self) -> int: device=self.device, ) position = position - self.arena_offsets[0] - wxyz = torch.as_tensor( - command.wxyz, - dtype=torch.float32, - device=self.device, + xyzw = convert_quat( + torch.as_tensor( + command.wxyz, + dtype=torch.float32, + device=self.device, + ), + to="xyzw", ).unsqueeze(0) pose = torch.eye( 4, dtype=torch.float32, device=self.device, ).unsqueeze(0) - pose[0, :3, :3] = matrix_from_quat(wxyz)[0] + pose[0, :3, :3] = matrix_from_quat(xyzw)[0] pose[0, :3, 3] = position if not gizmo.request_local_pose(pose, source_id=source_id): continue @@ -2245,7 +2833,12 @@ def set_gizmo_visibility( gizmo.set_visible(visible) def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: - """General interface to add a sensor to the scene and returns a handle. + """Create a sensor on the pre-created simulation Arenas. + + Cameras keep EmbodiChain's native CameraGroup implementation. A camera + attached to an articulation link is created immediately and attached + after the physical Spawn scene is prepared. ContactSensor still + requires the Default backend scene and therefore prepares physics first. Args: sensor_cfg (SensorCfg): configuration for the sensor. @@ -2254,28 +2847,122 @@ def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: BaseSensor: The added sensor instance handle. """ sensor_type = sensor_cfg.sensor_type - if sensor_type not in self.SUPPORTED_SENSOR_TYPES: - logger.log_warning(f"Unsupported sensor type: {sensor_type}") - return None + uid = sensor_cfg.uid + if uid is None: + uid = f"{sensor_type.lower()}_{len(self._sensors)}" + sensor_cfg.uid = uid + if uid in self._sensors: + raise ValueError(f"Sensor {uid!r} already exists.") - sensor_uid = sensor_cfg.uid - if sensor_uid is None: - sensor_uid = f"{sensor_type.lower()}_{len(self._sensors)}" - sensor_cfg.uid = sensor_uid + sensor_factory = self.SUPPORTED_SENSOR_TYPES.get(sensor_type) + if sensor_factory is None: + raise ValueError( + f"Unsupported sensor type {sensor_type!r}. Supported types: " + f"{sorted(self.SUPPORTED_SENSOR_TYPES)}." + ) + if sensor_type == "ContactSensor" and self.is_newton_backend: + raise NotImplementedError( + "ContactSensor currently requires the Default backend PhysicsScene. " + "Newton needs a public backend-neutral contact query API in DexSim." + ) - if sensor_uid in self._sensors: - logger.log_warning(f"Sensor {sensor_uid} already exists.") - return None + if isinstance(sensor_factory, type) and issubclass(sensor_factory, Camera): + if len(self._arenas) != self.num_envs: + raise RuntimeError( + "Camera creation requires all Spawn Arenas to be " + f"prepared ({len(self._arenas)} of {self.num_envs} ready)." + ) + sensor = sensor_factory( + sensor_cfg, + self.device, + owner=self, + ) + if sensor_cfg.extrinsics.parent is not None: + scene = self._spawn_scene + if scene.builder.result is not None: + self._attach_camera_parent(sensor) + else: + self._pending_sensor_attachments.append(sensor) + else: + # ContactSensor and custom native sensors require a prepared + # physics scene; cameras only depend on the pre-created Arenas. + self.prepare() + # Preserve custom test/plugin factories whose two-argument + # constructor predates the manager-owned render context. + sensor = sensor_factory(sensor_cfg, self.device) + + self._sensors[uid] = sensor + self.notify_visualization_topology_changed() + return sensor + + def _attach_camera_parent(self, sensor: Camera) -> None: + """Resolve and attach one camera to its configured parent nodes.""" + parent = sensor.cfg.extrinsics.parent + if parent is None: + return + parent_nodes = self._resolve_spawn_sensor_parent_nodes(parent) + sensor.attach_to_parent_nodes(parent_nodes) - sensor = self.SUPPORTED_SENSOR_TYPES[sensor_type](sensor_cfg, self.device) + def _resolve_spawn_sensor_parent_nodes(self, parent: str) -> list[object]: + """Resolve one canonical articulation link to a render node per Arena. - self._sensors[sensor_uid] = sensor - if isinstance(sensor, Camera): - self.notify_visualization_topology_changed() + A plain link name remains compatible with existing CameraCfg values. + When more than one robot/articulation owns that link, callers can use + ``"/"`` to disambiguate without introducing + backend clone suffixes. + """ + assets: dict[str, Articulation] = { + **self._articulations, + **self._robots, + } + asset_uid: str | None = None + link_name = parent + if "/" in parent: + candidate_uid, candidate_link = parent.split("/", maxsplit=1) + if candidate_uid in assets: + asset_uid = candidate_uid + link_name = candidate_link + + matches: list[tuple[str, list[object]]] = [] + for uid, asset in assets.items(): + if asset_uid is not None and uid != asset_uid: + continue + handles = list(getattr(asset, "_entities", ())) + if len(handles) != self.num_envs: + continue + if link_name not in handles[0].get_link_names(): + continue - # Check if the sensor needs to change the parent frame. + nodes: list[object] = [] + for handle in handles: + if link_name not in handle.get_link_names(): + raise RuntimeError( + f"Articulation {uid!r} has heterogeneous link topology; " + f"link {link_name!r} is missing in one Arena." + ) + render_body = handle.get_render_body(link_name) + if render_body is None: + raise RuntimeError( + f"Articulation {uid!r} link {link_name!r} has no public " + "render node for camera attachment." + ) + nodes.append(render_body.render_node()) + matches.append((uid, nodes)) - return sensor + if len(matches) == 1: + return matches[0][1] + if len(matches) > 1: + owners = ", ".join(uid for uid, _ in matches) + raise ValueError( + f"Camera parent link {link_name!r} is ambiguous across assets " + f"[{owners}]; use '/{link_name}'." + ) + scope = f" on asset {asset_uid!r}" if asset_uid is not None else "" + raise ValueError( + f"Camera parent link {link_name!r} was not found{scope} in any " + "Spawn-bound Robot or Articulation. Attachment to arbitrary render " + "nodes is not yet supported by the Spawn-only bridge." + ) def get_sensor(self, uid: str) -> BaseSensor | None: """Get a sensor by its UID. @@ -2302,53 +2989,42 @@ def get_sensor_uid_list(self) -> List[str]: def remove_asset(self, uid: str) -> bool: """Remove an asset by its UID. - The asset can be a light, sensor, robot, rigid object or articulation. - - Note: - Currently, lights and sensors are not supported to be removed. + Native render lights are not removed by this method. Sensors and + Spawn-owned physical assets are supported. Args: uid (str): The UID of the asset. Returns: bool: True if the asset is removed successfully, otherwise False. """ - if uid in self._rigid_objects: - obj = self._rigid_objects.pop(uid) - obj.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._soft_objects: - obj = self._soft_objects.pop(uid) - obj.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._cloth_objects: - obj = self._cloth_objects.pop(uid) - obj.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._rigid_object_groups: - group = self._rigid_object_groups.pop(uid) - group.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._articulations: - art = self._articulations.pop(uid) - art.destroy() - self.notify_visualization_topology_changed() - return True - - if uid in self._robots: - robot = self._robots.pop(uid) - robot.destroy() + if uid in self._sensors: + sensor = self._sensors.pop(uid) + if sensor in self._pending_sensor_attachments: + self._pending_sensor_attachments.remove(sensor) + destroy = getattr(sensor, "destroy", None) + if callable(destroy): + destroy() self.notify_visualization_topology_changed() return True - return False + scene = self._spawn_scene + if uid not in scene: + return False + if uid == "default_plane": + raise ValueError("The Spawn-owned default plane cannot be removed.") + + was_materialized = scene.builder.is_finalized + scene.remove(uid) + if was_materialized: + self.prepare() + + self._rigid_objects.pop(uid, None) + self._rigid_object_groups.pop(uid, None) + self._deformable_objects.pop(uid, None) + self._articulations.pop(uid, None) + self._robots.pop(uid, None) + self.notify_visualization_topology_changed() + return True def draw_marker( self, @@ -3094,12 +3770,9 @@ def reset_objects_state( for uid, rigid_obj_group in self._rigid_object_groups.items(): if uid not in excluded_uids: rigid_obj_group.reset(env_ids) - for uid, soft_obj in self._soft_objects.items(): + for uid, deformable_obj in self._deformable_objects.items(): if uid not in excluded_uids: - soft_obj.reset(env_ids) - for uid, cloth_obj in self._cloth_objects.items(): - if uid not in excluded_uids: - cloth_obj.reset(env_ids) + deformable_obj.reset(env_ids) for uid, light in self._lights.items(): if uid not in excluded_uids: light.reset(env_ids) @@ -3202,6 +3875,48 @@ def _deferred_destroy(self) -> None: import sys, gc + # Release backend-owned views before SpawnResult closes the native + # resources that back them. Newton also synchronizes its device here. + self.physics.prepare_for_teardown() + # Run wrapper destructors while their World is still alive. The later + # collections continue to break cycles left by the native teardown. + gc.collect() + + # Render-only cameras may be attached to Spawn articulation link + # nodes. Remove their Arena views before closing SpawnResult, which + # releases those parent nodes, and before World.quit releases their + # CameraGroups. + for sensor in list(getattr(self, "_sensors", {}).values()): + try: + sensor.destroy() + except Exception as error: + logger.log_warning( + f"Failed to destroy sensor {getattr(sensor, 'uid', None)!r}: " + f"{error!r}" + ) + + if self._spawn_scene is not None: + # Release result-scoped batches/facades before closing the + # SpawnResult and, finally, the World that owns native resources. + for registry_name in ( + "_rigid_objects", + "_rigid_object_groups", + "_deformable_objects", + "_articulations", + "_robots", + ): + for asset in getattr(self, registry_name, {}).values(): + if hasattr(asset, "_data"): + asset._data = None + if hasattr(asset, "_spawn_result"): + asset._spawn_result = None + if hasattr(asset, "_entities"): + asset._entities = [] + try: + self._spawn_scene.close() + finally: + self._spawn_scene = None + self.clean_materials() if self._env: @@ -3236,15 +3951,13 @@ def _sever_wrapper_refs(obj_registry): _sever_wrapper_refs("_rigid_objects") _sever_wrapper_refs("_constraints") _sever_wrapper_refs("_rigid_object_groups") - _sever_wrapper_refs("_soft_objects") - _sever_wrapper_refs("_cloth_objects") + _sever_wrapper_refs("_deformable_objects") _sever_wrapper_refs("_articulations") _sever_wrapper_refs("_robots") _sever_wrapper_refs("_sensors") _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 @@ -3301,3 +4014,12 @@ def flush_cleanup_queue() -> None: # 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/skills/calls.py b/embodichain/lab/sim/skills/calls.py index 3935f74bc..2b8e803da 100644 --- a/embodichain/lab/sim/skills/calls.py +++ b/embodichain/lab/sim/skills/calls.py @@ -184,52 +184,52 @@ def _validate_static_skill_descriptor( @dataclass(frozen=True, slots=True, init=False, eq=False) class SemanticPose: - """Object-space pose expressed as position and a WXYZ quaternion. + """Object-space pose expressed as position and an XYZW quaternion. The value owns normalized tensor snapshots and never exposes its internal tensors directly. A single pose or an environment batch is accepted. Args: position: Shape ``(3,)`` or ``(B, 3)``. - quaternion_wxyz: Shape ``(4,)`` or ``(B, 4)``. Finite, non-zero + quaternion_xyzw: Shape ``(4,)`` or ``(B, 4)``. Finite, non-zero quaternions are normalized at construction. """ _position: torch.Tensor = field(repr=False) - _quaternion_wxyz: torch.Tensor = field(repr=False) + _quaternion_xyzw: torch.Tensor = field(repr=False) def __init__( self, position: torch.Tensor | tuple[float, float, float] | list[float], - quaternion_wxyz: torch.Tensor | tuple[float, float, float, float] | list[float], + quaternion_xyzw: torch.Tensor | tuple[float, float, float, float] | list[float], ) -> None: position_tensor = torch.as_tensor(position, dtype=torch.float32) - quaternion_tensor = torch.as_tensor(quaternion_wxyz, dtype=torch.float32) + quaternion_tensor = torch.as_tensor(quaternion_xyzw, dtype=torch.float32) if position_tensor.dim() not in (1, 2) or position_tensor.shape[-1] != 3: raise ValueError("position must have shape (3,) or (B, 3).") if quaternion_tensor.dim() not in (1, 2) or quaternion_tensor.shape[-1] != 4: - raise ValueError("quaternion_wxyz must have shape (4,) or (B, 4).") + raise ValueError("quaternion_xyzw must have shape (4,) or (B, 4).") if position_tensor.dim() != quaternion_tensor.dim(): raise ValueError( - "position and quaternion_wxyz must both be unbatched or batched." + "position and quaternion_xyzw must both be unbatched or batched." ) if position_tensor.dim() == 2 and ( position_tensor.shape[0] != quaternion_tensor.shape[0] ): - raise ValueError("position and quaternion_wxyz batch sizes must match.") + raise ValueError("position and quaternion_xyzw batch sizes must match.") if position_tensor.dim() == 2 and position_tensor.shape[0] == 0: raise ValueError("SemanticPose batches must contain at least one pose.") if not torch.isfinite(position_tensor).all(): raise ValueError("position must contain only finite values.") if not torch.isfinite(quaternion_tensor).all(): - raise ValueError("quaternion_wxyz must contain only finite values.") + raise ValueError("quaternion_xyzw must contain only finite values.") norms = torch.linalg.vector_norm(quaternion_tensor, dim=-1, keepdim=True) if torch.any(norms <= torch.finfo(torch.float32).eps): - raise ValueError("quaternion_wxyz must be non-zero.") + raise ValueError("quaternion_xyzw must be non-zero.") object.__setattr__(self, "_position", position_tensor.clone()) object.__setattr__( self, - "_quaternion_wxyz", + "_quaternion_xyzw", (quaternion_tensor / norms).clone(), ) @@ -239,9 +239,9 @@ def position(self) -> torch.Tensor: return self._position.clone() @property - def quaternion_wxyz(self) -> torch.Tensor: + def quaternion_xyzw(self) -> torch.Tensor: """Return an independent normalized quaternion tensor.""" - return self._quaternion_wxyz.clone() + return self._quaternion_xyzw.clone() @property def batch_size(self) -> int | None: @@ -250,7 +250,7 @@ def batch_size(self) -> int | None: def snapshot(self) -> SemanticPose: """Return an independently owned pose value.""" - return SemanticPose(self._position, self._quaternion_wxyz) + return SemanticPose(self._position, self._quaternion_xyzw) def to_matrix(self) -> torch.Tensor: """Convert the semantic pose to a homogeneous transform. @@ -259,14 +259,14 @@ def to_matrix(self) -> torch.Tensor: Shape ``(4, 4)`` for an unbatched pose or ``(B, 4, 4)`` for a batched pose. """ - quaternion = self._quaternion_wxyz + quaternion = self._quaternion_xyzw was_unbatched = quaternion.dim() == 1 if was_unbatched: quaternion = quaternion.unsqueeze(0) position = self._position.unsqueeze(0) else: position = self._position - w, x, y, z = quaternion.unbind(dim=-1) + x, y, z, w = quaternion.unbind(dim=-1) output = torch.zeros( quaternion.shape[0], 4, @@ -291,7 +291,7 @@ def to_metadata(self) -> dict[str, object]: """Return the pose as deterministic JSON-safe semantic data.""" return { "position": self._position.detach().cpu().tolist(), - "quaternion_wxyz": self._quaternion_wxyz.detach().cpu().tolist(), + "quaternion_xyzw": self._quaternion_xyzw.detach().cpu().tolist(), } diff --git a/embodichain/lab/sim/solvers/differential_solver.py b/embodichain/lab/sim/solvers/differential_solver.py index 12e51bcbd..239ee332e 100644 --- a/embodichain/lab/sim/solvers/differential_solver.py +++ b/embodichain/lab/sim/solvers/differential_solver.py @@ -139,7 +139,7 @@ def action_dim(self) -> int: elif self.cfg.command_type == "pose" and self.cfg.use_relative_mode: return 6 # (dx, dy, dz, droll, dpitch, dyaw) else: - return 7 # (x, y, z, qw, qx, qy, qz) + return 7 # (x, y, z, qx, qy, qz, qw) def reset(self, env_ids: torch.Tensor | None = None): """Reset the internal buffers for the specified environments. @@ -151,7 +151,7 @@ def reset(self, env_ids: torch.Tensor | None = None): env_ids = torch.arange(self.num_envs, device=self.device) self.ee_pos_des[env_ids] = 0 - self.ee_quat_des[env_ids] = torch.tensor([1.0, 0, 0, 0], device=self.device) + self.ee_quat_des[env_ids] = torch.tensor([0.0, 0, 0, 1.0], device=self.device) self._command[env_ids] = 0 def set_command( @@ -412,9 +412,8 @@ def _matrix_to_pos_quat(mat): rot_matrices = mat[:, :3, :3].cpu().numpy() # Convert to NumPy for scipy quats = Rotation.from_matrix(rot_matrices).as_quat() # (N, 4), [x, y, z, w] - # Convert quaternion back to torch.Tensor and reorder to [w, x, y, z] + # SciPy's xyzw convention matches EmbodiChain's quaternion contract. quats = torch.tensor(quats, device=mat.device, dtype=mat.dtype) # (N, 4) - quats = quats[:, [3, 0, 1, 2]] # Reorder to [w, x, y, z] # Concatenate position and quaternion return torch.cat([pos, quats], dim=1) diff --git a/embodichain/lab/sim/solvers/neural_ik_solver.py b/embodichain/lab/sim/solvers/neural_ik_solver.py index 7f1cb1d12..cdf0d27a0 100644 --- a/embodichain/lab/sim/solvers/neural_ik_solver.py +++ b/embodichain/lab/sim/solvers/neural_ik_solver.py @@ -19,11 +19,7 @@ import torch.nn as nn from embodichain.utils import configclass -from embodichain.utils.math import ( - convert_quat, - quat_error_magnitude, - quat_from_matrix, -) +from embodichain.utils.math import quat_error_magnitude, quat_from_matrix from embodichain.lab.sim.solvers import SolverCfg, BaseSolver from embodichain.lab.sim.solvers.qpos_seed_sampler import QposSeedSampler @@ -194,7 +190,7 @@ def _run_policy( for _ in range(self._max_steps): ee_xpos = self.get_fk(qpos) ee_pos = ee_xpos[:, :3, 3] - ee_quat = convert_quat(quat_from_matrix(ee_xpos[:, :3, :3]), to="xyzw") + ee_quat = quat_from_matrix(ee_xpos[:, :3, :3]) obs = self._build_obs( qpos, ee_pos, ee_quat, target_pos, target_quat, last_action @@ -212,9 +208,9 @@ def _run_policy( # Convergence check ik_xpos = self.get_fk(qpos) pos_err = (ik_xpos[:, :3, 3] - target_pos).norm(dim=-1) - ik_quat_wxyz = quat_from_matrix(ik_xpos[:, :3, :3]) - target_quat_wxyz = quat_from_matrix(target_xpos[:, :3, :3]) - rot_err = quat_error_magnitude(target_quat_wxyz, ik_quat_wxyz) + ik_quat_xyzw = quat_from_matrix(ik_xpos[:, :3, :3]) + target_quat_xyzw = quat_from_matrix(target_xpos[:, :3, :3]) + rot_err = quat_error_magnitude(target_quat_xyzw, ik_quat_xyzw) success = (pos_err < self._pos_eps) & (rot_err < self._rot_eps) return success, qpos @@ -253,7 +249,7 @@ def get_ik( B = target_xpos.shape[0] target_pos = target_xpos[:, :3, 3] - target_quat = convert_quat(quat_from_matrix(target_xpos[:, :3, :3]), to="xyzw") + target_quat = quat_from_matrix(target_xpos[:, :3, :3]) if qpos_seed is None: qpos_seed = torch.zeros(B, self.dof, device=self.device) @@ -279,9 +275,7 @@ def get_ik( ) target_xpos_repeated = sampler.repeat_target_xpos(target_xpos, n) target_pos_rep = target_xpos_repeated[:, :3, 3] - target_quat_rep = convert_quat( - quat_from_matrix(target_xpos_repeated[:, :3, :3]), to="xyzw" - ) + target_quat_rep = quat_from_matrix(target_xpos_repeated[:, :3, :3]) success_flat, ik_qpos_flat = self._run_policy( all_seeds, target_xpos_repeated, target_pos_rep, target_quat_rep diff --git a/embodichain/lab/sim/spawn/__init__.py b/embodichain/lab/sim/spawn/__init__.py new file mode 100644 index 000000000..4aa7b1513 --- /dev/null +++ b/embodichain/lab/sim/spawn/__init__.py @@ -0,0 +1,40 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Translate EmbodiChain asset configs into DexSim Spawn descriptors.""" + +from __future__ import annotations + +from .descriptors import ( + articulation_desc_from_cfg, + cloth_desc_from_cfg, + rigid_desc_from_cfg, + soft_desc_from_cfg, + surface_deformable_desc_from_cfg, + volume_deformable_desc_from_cfg, +) +from .usd import articulation_desc_from_usd, rigid_desc_from_usd + +__all__ = [ + "articulation_desc_from_cfg", + "articulation_desc_from_usd", + "cloth_desc_from_cfg", + "rigid_desc_from_cfg", + "rigid_desc_from_usd", + "soft_desc_from_cfg", + "surface_deformable_desc_from_cfg", + "volume_deformable_desc_from_cfg", +] diff --git a/embodichain/lab/sim/spawn/descriptors.py b/embodichain/lab/sim/spawn/descriptors.py new file mode 100644 index 000000000..4c64d124f --- /dev/null +++ b/embodichain/lab/sim/spawn/descriptors.py @@ -0,0 +1,1400 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Translate EmbodiChain asset configurations into DexSim Spawn descriptors. + +This module translates one EmbodiChain configuration into a canonical +descriptor carrying both the common physics values and the optional backend +extension blocks. The selected :mod:`dexsim.spawn` adapter remains the only +component that chooses between the Default and Newton backends. When supplied, the active +Newton solver type only prevents common contact values from being authored to +a solver that cannot consume them. + +Articulation source names come from the handles produced by normal backend +materialization. EmbodiChain owns regex/group selection, applies exact-name +typed properties, and explicitly rebuilds Newton once when those post-load +properties must be committed to its immutable model. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import MISSING, dataclass, field, fields +import math +import numbers +import os +import warnings +from typing import TYPE_CHECKING + +import numpy as np +from dexsim.spawn import ( + ArticulationDesc, + ClothDesc, + ClothPhysicsDesc, + CollisionApproximation, + CollisionDesc, + DexsimClothPhysicsDesc, + DexsimCollisionDesc, + DexsimJointDesc, + DexsimPhysicsDesc, + DexsimSoftBodyPhysicsDesc, + GeometryDesc, + MaterialDesc, + NewtonCollisionDesc, + NewtonJointDesc, + ObjectDesc, + RenderDesc, + RigidBodyPhysicsDesc, + SoftBodyDesc, + SoftBodyMeshingDesc, + SoftBodyPhysicsDesc, +) +from dexsim.spawn.descs import NEWTON_CONTACT_SOLVER_FIELDS +from dexsim.types import ActorType, DriveType, LoadOption as DexsimLoadOption + +from embodichain.lab.sim.cfg import ( + _normalize_joint_target_mode, + ArticulationCfg, + ClothObjectCfg, + CollisionPropertiesCfg, + DefaultCollisionPropertiesCfg, + DefaultRigidBodyPropertiesCfg, + MassPropertiesCfg, + NewtonCollisionPropertiesCfg, + NewtonRigidBodyMaterialCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + SoftObjectCfg, + SurfaceDeformableObjectCfg, + VolumeDeformableObjectCfg, +) +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, MeshCollisionCfg, SphereCfg +from embodichain.utils import logger +from embodichain.utils.math import convert_quat +from embodichain.utils.string import ( + resolve_matching_names, + resolve_matching_names_values, +) + +if TYPE_CHECKING: + from embodichain.lab.sim.material import VisualMaterialCfg + +__all__ = [ + "articulation_desc_from_cfg", + "cloth_desc_from_cfg", + "configure_articulation_desc", + "rigid_desc_from_cfg", + "soft_desc_from_cfg", + "surface_deformable_desc_from_cfg", + "volume_deformable_desc_from_cfg", +] + + +@dataclass +class _RigidPhysicsSpec: + """Canonical, backend-partitioned rigid-physics values.""" + + mass_props: dict[str, object] = field(default_factory=dict) + recompute_inertia: bool | None = None + default_rigid_props: dict[str, object] = field(default_factory=dict) + collision_enabled: bool | None = None + contact_offset: float | None = None + rest_offset: float | None = None + default_collision_props: dict[str, object] = field(default_factory=dict) + newton_collision_props: dict[str, object] = field(default_factory=dict) + material_props: dict[str, object] = field(default_factory=dict) + newton_material_props: dict[str, object] = field(default_factory=dict) + + def merged(self, override: _RigidPhysicsSpec) -> _RigidPhysicsSpec: + """Return ``override`` layered onto this spec using non-None values.""" + result = _RigidPhysicsSpec( + mass_props=dict(self.mass_props), + recompute_inertia=self.recompute_inertia, + default_rigid_props=dict(self.default_rigid_props), + collision_enabled=self.collision_enabled, + contact_offset=self.contact_offset, + rest_offset=self.rest_offset, + default_collision_props=dict(self.default_collision_props), + newton_collision_props=dict(self.newton_collision_props), + material_props=dict(self.material_props), + newton_material_props=dict(self.newton_material_props), + ) + for name in ( + "mass_props", + "default_rigid_props", + "default_collision_props", + "newton_collision_props", + "material_props", + "newton_material_props", + ): + getattr(result, name).update(getattr(override, name)) + if "mass" in override.mass_props: + mass = float(override.mass_props["mass"]) + if mass > 0.0: + result.mass_props.pop("density", None) + elif mass == 0.0 and "density" in result.mass_props: + result.mass_props.pop("mass", None) + elif "density" in override.mass_props: + result.mass_props.pop("mass", None) + if override.recompute_inertia is not None: + result.recompute_inertia = override.recompute_inertia + if override.collision_enabled is not None: + result.collision_enabled = override.collision_enabled + if override.contact_offset is not None: + result.contact_offset = override.contact_offset + if override.rest_offset is not None: + result.rest_offset = override.rest_offset + return result + + +def _configured_values(cfg: object | None) -> dict[str, object]: + """Return non-None configclass fields without backend metadata.""" + if cfg is None: + return {} + return { + item.name: value + for item in fields(cfg) + if (value := getattr(cfg, item.name)) is not None + } + + +def _resolve_rigid_physics( + cfg: RigidBodyPhysicsCfg, + *, + newton_solver_type: str | None = None, +) -> _RigidPhysicsSpec: + """Normalize grouped rigid-body configuration into one internal spec.""" + if isinstance(cfg, RigidBodyPhysicsCfg): + mass_props = _configured_values(cfg.mass_props) + recompute_inertia = mass_props.pop("recompute_inertia", None) + if recompute_inertia is not None and not isinstance( + recompute_inertia, (bool, np.bool_) + ): + raise TypeError("recompute_inertia must be a boolean or None.") + spec = _RigidPhysicsSpec( + mass_props=mass_props, + recompute_inertia=( + None if recompute_inertia is None else bool(recompute_inertia) + ), + collision_enabled=( + None + if cfg.collision_props is None + else cfg.collision_props.collision_enabled + ), + contact_offset=( + None + if cfg.collision_props is None + else cfg.collision_props.contact_offset + ), + rest_offset=( + None if cfg.collision_props is None else cfg.collision_props.rest_offset + ), + material_props={ + name: getattr(cfg.material_props, name) + for name in ("static_friction", "dynamic_friction", "restitution") + if cfg.material_props is not None + and getattr(cfg.material_props, name) is not None + }, + ) + + rigid_props = cfg.rigid_props + if isinstance(rigid_props, DefaultRigidBodyPropertiesCfg): + spec.default_rigid_props = _configured_values(rigid_props) + elif rigid_props is not None: + raise TypeError( + f"Unsupported rigid_props type {type(rigid_props).__name__!r}." + ) + + collision_props = cfg.collision_props + if isinstance(collision_props, DefaultCollisionPropertiesCfg): + spec.default_collision_props = _configured_values(collision_props) + for name in ("collision_enabled", "contact_offset", "rest_offset"): + spec.default_collision_props.pop(name, None) + elif isinstance(collision_props, NewtonCollisionPropertiesCfg): + values = _configured_values(collision_props) + for name in ("collision_enabled", "contact_offset", "rest_offset"): + values.pop(name, None) + spec.newton_collision_props = values + elif ( + collision_props is not None + and type(collision_props) is not CollisionPropertiesCfg + ): + raise TypeError( + f"Unsupported collision_props type {type(collision_props).__name__!r}." + ) + + material_props = cfg.material_props + if isinstance(material_props, NewtonRigidBodyMaterialCfg): + values = _configured_values(material_props) + for name in ("static_friction", "dynamic_friction", "restitution"): + values.pop(name, None) + if "torsional_friction" in values: + values["mu_torsional"] = values.pop("torsional_friction") + if "rolling_friction" in values: + values["mu_rolling"] = values.pop("rolling_friction") + spec.newton_material_props = values + elif ( + material_props is not None + and type(material_props) is not RigidBodyMaterialCfg + ): + raise TypeError( + f"Unsupported material_props type {type(material_props).__name__!r}." + ) + + return spec + + raise AssertionError("Unhandled grouped rigid-body physics configuration.") + + +def rigid_desc_from_cfg( + cfg: RigidObjectCfg, + *, + per_env: bool = True, + newton_solver_type: str | None = None, +) -> tuple[ObjectDesc, dict[str, MaterialDesc]]: + """Translate a rigid-object config into a DexSim Spawn descriptor.""" + uid = _required_uid(cfg.uid, "Rigid object") + if isinstance(cfg.shape, MeshCfg) and _is_usd_path(cfg.shape.fpath): + raise NotImplementedError( + "USD files describe typed scenes; use rigid_desc_from_usd() to " + "select the sole rigid object." + ) + + physics = _resolve_rigid_physics( + cfg.attrs, + newton_solver_type=newton_solver_type, + ) + geometry, approximation, max_hulls = _compile_geometry(cfg) + material_ref, material_entry = _compile_visual_material( + uid, cfg.shape.visual_material + ) + collision = CollisionDesc.from_geometry( + geometry, + approximation=approximation, + ) + collision.enable_collision = physics.collision_enabled + collision.decomp_max_hulls = max_hulls + collision.dexsim = _compile_default_collision(physics) + collision.newton = _compile_newton_collision( + physics, + newton_solver_type=newton_solver_type, + author_shape_defaults=True, + mesh_collision=( + cfg.shape.collision if isinstance(cfg.shape, MeshCfg) else None + ), + ) + collision.render_source_index = 0 + + descriptor = ObjectDesc( + name=uid, + pose=_pose_from_cfg(cfg), + renders=[ + RenderDesc.from_geometry( + geometry, + load_option=_compile_load_option(cfg.shape), + material_ref=material_ref, + ) + ], + collisions=[collision], + physics=_compile_rigid_physics(physics, cfg.body_type), + per_env=per_env, + body_scale=_vector3(cfg.body_scale, field_name="body_scale"), + ) + materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} + return descriptor, materials + + +def volume_deformable_desc_from_cfg( + cfg: VolumeDeformableObjectCfg, + *, + per_env: bool = True, +) -> tuple[SoftBodyDesc, dict[str, MaterialDesc]]: + """Translate a volume-deformable config into a DexSim descriptor.""" + uid = _required_uid(cfg.uid, "Volume deformable") + if _is_missing(cfg.shape.fpath) or not str(cfg.shape.fpath).strip(): + raise ValueError( + "VolumeDeformableObjectCfg.shape.fpath must be a non-empty path." + ) + geometry = GeometryDesc.mesh(file_path=str(cfg.shape.fpath), segment_name=uid) + material_ref, material_entry = _compile_visual_material( + uid, cfg.shape.visual_material + ) + physical_attr = cfg.physical_attr + youngs = float(physical_attr.youngs) + poissons = float(physical_attr.poissons) + descriptor = SoftBodyDesc( + name=uid, + pose=_pose_from_cfg(cfg), + mesh=RenderDesc.from_geometry( + geometry, + load_option=_compile_load_option(cfg.shape), + material_ref=material_ref, + ), + physics=SoftBodyPhysicsDesc( + volume_density=float(physical_attr.density), + k_mu=youngs / (2.0 * (1.0 + poissons)), + k_lambda=(youngs * poissons / ((1.0 + poissons) * (1.0 - 2.0 * poissons))), + dexsim=DexsimSoftBodyPhysicsDesc(**_configured_values(physical_attr)), + ), + # DexSim's typed meshing contract currently exposes these three + # source-mesh controls; maximal_edge_length has no Spawn equivalent. + meshing=SoftBodyMeshingDesc( + proxy_simplify_target=cfg.voxel_attr.triangle_simplify_target, + proxy_remesh_resolution=cfg.voxel_attr.triangle_remesh_resolution, + voxel_resolution=cfg.voxel_attr.simulation_mesh_resolution, + ), + per_env=per_env, + ) + materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} + return descriptor, materials + + +def surface_deformable_desc_from_cfg( + cfg: SurfaceDeformableObjectCfg, + *, + per_env: bool = True, +) -> tuple[ClothDesc, dict[str, MaterialDesc]]: + """Translate a surface-deformable config into a DexSim descriptor.""" + uid = _required_uid(cfg.uid, "Surface deformable") + if _is_missing(cfg.shape.fpath) or not str(cfg.shape.fpath).strip(): + raise ValueError( + "SurfaceDeformableObjectCfg.shape.fpath must be a non-empty path." + ) + geometry = GeometryDesc.mesh(file_path=str(cfg.shape.fpath), segment_name=uid) + material_ref, material_entry = _compile_visual_material( + uid, cfg.shape.visual_material + ) + descriptor = ClothDesc( + name=uid, + pose=_pose_from_cfg(cfg), + mesh=RenderDesc.from_geometry( + geometry, + load_option=_compile_load_option(cfg.shape), + material_ref=material_ref, + ), + physics=ClothPhysicsDesc( + surface_density=float(cfg.physical_attr.density), + dexsim=DexsimClothPhysicsDesc(**_configured_values(cfg.physical_attr)), + ), + per_env=per_env, + ) + materials = {} if material_entry is None else {material_entry[0]: material_entry[1]} + return descriptor, materials + + +def soft_desc_from_cfg( + cfg: SoftObjectCfg, + *, + per_env: bool = True, +) -> tuple[SoftBodyDesc, dict[str, MaterialDesc]]: + """Compatibility wrapper for :func:`volume_deformable_desc_from_cfg`.""" + return volume_deformable_desc_from_cfg(cfg, per_env=per_env) + + +def cloth_desc_from_cfg( + cfg: ClothObjectCfg, + *, + per_env: bool = True, +) -> tuple[ClothDesc, dict[str, MaterialDesc]]: + """Compatibility wrapper for :func:`surface_deformable_desc_from_cfg`.""" + return surface_deformable_desc_from_cfg(cfg, per_env=per_env) + + +def articulation_desc_from_cfg( + cfg: ArticulationCfg, + *, + per_env: bool = True, + source_path: str | None = None, + newton_solver_type: str | None = None, +) -> ArticulationDesc: + """Translate an articulation config into a DexSim Spawn descriptor.""" + path = source_path if source_path is not None else cfg.fpath + if path is None or not str(path).strip(): + raise ValueError( + "No articulation source path is available. Assemble the robot URDF " + "before converting its configuration." + ) + if _is_usd_path(path): + raise NotImplementedError( + "USD files describe typed scenes; use articulation_desc_from_usd() " + "to select the sole articulation." + ) + if cfg.resolve_asset_physics_mode() == "overlay": + _validate_articulation_rigid_physics( + cfg, + newton_solver_type=newton_solver_type, + ) + fixed_base, self_collision_enabled = _articulation_root_values(cfg) + return ArticulationDesc( + name=_articulation_uid(cfg.uid, str(path)), + pose=_pose_from_cfg(cfg), + path=str(path), + urdf_path=str(path), + fixed_base=fixed_base, + enable_self_collision=self_collision_enabled, + urdf_fix_root_link=fixed_base, + # EmbodiChain's preserve/overlay policy starts from source-authored + # inertia. MassPropertiesCfg can request geometry-based recomputation + # after exact source names are available. + urdf_read_inertia=True, + per_env=per_env, + body_scale=_vector3(cfg.body_scale, field_name="body_scale"), + ) + + +def _validate_articulation_rigid_physics( + cfg: ArticulationCfg, + *, + newton_solver_type: str | None, +) -> None: + """Validate global and per-link physics before source materialization.""" + _resolve_rigid_physics( + cfg.attrs, + newton_solver_type=newton_solver_type, + ) + for group in (cfg.link_attrs or {}).values(): + _resolve_rigid_physics( + group.attrs, + newton_solver_type=newton_solver_type, + ) + + +def _articulation_root_values( + cfg: ArticulationCfg, + *, + fixed_base_default: bool = True, + self_collision_default: bool = False, +) -> tuple[bool, bool]: + """Resolve articulation-root values over source/import defaults.""" + props = cfg.root_props + fixed_base = ( + fixed_base_default if props.fixed_base is None else bool(props.fixed_base) + ) + self_collision_enabled = ( + self_collision_default + if props.self_collision_enabled is None + else bool(props.self_collision_enabled) + ) + return fixed_base, self_collision_enabled + + +def _configured_articulation_overlay_fields(cfg: ArticulationCfg) -> list[str]: + """Return physics overlay fields that preserve mode would ignore.""" + configured: list[str] = [] + if any( + _configured_values(group) + for group in ( + cfg.attrs.mass_props, + cfg.attrs.rigid_props, + cfg.attrs.collision_props, + cfg.attrs.material_props, + ) + ): + configured.append("attrs") + if cfg.link_attrs: + configured.append("link_attrs") + if _configured_values(cfg.joint_drive_props): + configured.append("joint_drive_props") + if cfg.qpos_limits is not None: + configured.append("qpos_limits") + return configured + + +def _compile_link_properties( + physics: _RigidPhysicsSpec, + *, + newton_solver_type: str | None, + author_newton_shape_defaults: bool, +) -> tuple[RigidBodyPhysicsDesc, CollisionDesc, bool]: + collision = CollisionDesc( + enable_collision=physics.collision_enabled, + dexsim=_compile_default_collision(physics), + newton=_compile_newton_collision( + physics, + newton_solver_type=newton_solver_type, + author_shape_defaults=author_newton_shape_defaults, + ), + ) + return ( + _compile_rigid_physics(physics, "dynamic"), + collision, + bool(physics.recompute_inertia), + ) + + +def configure_articulation_desc( + desc: ArticulationDesc, + cfg: ArticulationCfg, + *, + newton_solver_type: str | None = None, +) -> ArticulationDesc: + """Apply one EmbodiChain config to exact source-resolved names. + + Regex/default/group semantics remain private to EmbodiChain. The DexSim + descriptor receives only concrete link and joint properties. + """ + if not desc.links: + raise RuntimeError( + f"Articulation source {desc.name!r} must be resolved before " + "configuration." + ) + if cfg.resolve_asset_physics_mode() == "preserve": + configured_fields = _configured_articulation_overlay_fields(cfg) + if configured_fields: + warnings.warn( + "asset_physics_mode='preserve' ignores configured articulation " + f"physics overlays: {', '.join(configured_fields)}. Set " + "asset_physics_mode='overlay' to apply them.", + UserWarning, + stacklevel=2, + ) + return desc + default_physics = _resolve_rigid_physics( + cfg.attrs, + newton_solver_type=newton_solver_type, + ) + author_newton_shape_defaults = not _is_usd_path(cfg.fpath) + default_link_properties = _compile_link_properties( + default_physics, + newton_solver_type=newton_solver_type, + author_newton_shape_defaults=author_newton_shape_defaults, + ) + link_properties = {link.name: default_link_properties for link in desc.links} + + claimed_links: dict[str, str] = {} + link_names = [link.name for link in desc.links] + for group_name, group in (cfg.link_attrs or {}).items(): + _, matched_names = resolve_matching_names( + group.link_names_expr, + link_names, + ) + group_properties = _compile_link_properties( + default_physics.merged( + _resolve_rigid_physics( + group.attrs, + newton_solver_type=newton_solver_type, + ) + ), + newton_solver_type=newton_solver_type, + author_newton_shape_defaults=author_newton_shape_defaults, + ) + for link_name in matched_names: + previous = claimed_links.get(link_name) + if previous is not None: + raise ValueError( + f"Link {link_name!r} matches both {previous!r} and " + f"{group_name!r}." + ) + claimed_links[link_name] = group_name + link_properties[link_name] = group_properties + + ( + joint_properties, + joint_common, + joint_limits, + joint_target_modes, + ) = _compile_joint_properties( + desc, + cfg, + newton_solver_type=newton_solver_type, + ) + + # Commit only after every regex, value, and limit has been validated. Each + # source-resolved item receives one exact-name update. + for link_name, ( + rigid_body, + collision, + recompute_inertia, + ) in link_properties.items(): + link = desc.get_link_desc(link_name) + desc.set_link_properties( + link_name, + rigid_body=rigid_body, + # The URDF resolver intentionally keeps source-owned collision + # geometry outside LinkDesc. An attribute-only CollisionDesc is + # still required so the adapters can overlay properties onto the + # native source shapes; it does not synthesize geometry. Explicit + # descriptors, including collisionless links, remain unchanged. + collision=( + collision if link.collisions or desc.urdf_path is not None else None + ), + replace_inertial=recompute_inertia, + ) + for joint_name, (default_desc, newton_desc) in joint_properties.items(): + lower_limit, upper_limit = joint_limits.get(joint_name, (None, None)) + common = joint_common[joint_name] + desc.set_joint_properties( + joint_name, + lower_limit=lower_limit, + upper_limit=upper_limit, + effort_limit=common.get("effort_limit"), + velocity_limit=common.get("velocity_limit"), + armature=common.get("armature"), + dexsim=default_desc, + newton=newton_desc, + newton_target_mode=joint_target_modes.get(joint_name), + ) + return desc + + +def _compile_joint_properties( + desc: ArticulationDesc, + cfg: ArticulationCfg, + *, + newton_solver_type: str | None, +) -> tuple[ + dict[str, tuple[DexsimJointDesc, NewtonJointDesc]], + dict[str, dict[str, float]], + dict[str, tuple[object, object]], + dict[str, int], +]: + joint_names = [joint.name for joint in desc.joints] + control_parts = getattr(cfg, "control_parts", None) + target_mode_cfg: object = None + drive_type: str | None = None + if cfg.joint_drive_props is not None: + target_mode_cfg, drive_type = cfg.joint_drive_props._resolve_modes() + + joint_target_modes: dict[str, int] = {} + if target_mode_cfg is not None: + matches = _joint_property_matches( + target_mode_cfg, + joint_names, + property_name="target_mode", + numeric_only=False, + control_parts=control_parts, + ) + for joint_name, value in matches: + joint_target_modes[joint_name] = _normalize_joint_target_mode(value) + + # A scalar drive type remains the fallback for joints not selected by an + # explicit target-mode rule. The established force drive activates both + # position and velocity targets. + if drive_type is not None: + fallback_target_mode = 0 if drive_type == "none" else 3 + for joint_name in joint_names: + joint_target_modes.setdefault(joint_name, fallback_target_mode) + + active_joints = [ + name for name, mode in joint_target_modes.items() if mode in {1, 2, 3} + ] + if drive_type == "none" and active_joints: + raise ValueError( + "drive_type='none' conflicts with an active joint target_mode; " + "use target_mode='none' or 'effort'." + ) + if newton_solver_type is not None and drive_type == "acceleration": + if active_joints: + raise NotImplementedError( + "Newton Spawn does not have an exact acceleration-drive " + "equivalent; use drive_type='force' or disable the drive." + ) + + default_drive_mode = { + None: None, + "force": DriveType.FORCE, + "acceleration": DriveType.ACCELERATION, + "none": DriveType.NONE, + }[drive_type] + joint_properties = { + joint_name: ( + DexsimJointDesc( + drive_mode=( + DriveType.NONE + if joint_target_modes.get(joint_name) in {0, 4} + else ( + ( + default_drive_mode + if default_drive_mode is not None + else DriveType.FORCE + ) + if joint_target_modes.get(joint_name) in {1, 2, 3} + else None + ) + ) + ), + NewtonJointDesc(), + ) + for joint_name in joint_names + } + joint_common: dict[str, dict[str, float]] = { + joint_name: {} for joint_name in joint_names + } + property_fields = { + "stiffness": ("stiffness", "target_ke"), + "damping": ("damping", "target_kd"), + "friction": ("joint_friction", "friction"), + } + for property_name in ("stiffness", "damping"): + if cfg.joint_drive_props is None: + continue + configured = getattr(cfg.joint_drive_props, property_name) + if configured is None: + continue + matches = _joint_property_matches( + configured, + joint_names, + property_name=property_name, + control_parts=control_parts, + ) + for joint_name, value in matches: + if not isinstance(value, numbers.Number): + raise TypeError( + f"Articulation drive rule for {joint_name!r} and " + f"{property_name!r} must contain a numeric value." + ) + scalar = float(value) + default_desc, newton_desc = joint_properties[joint_name] + default_field, newton_field = property_fields[property_name] + setattr(default_desc, default_field, scalar) + setattr(newton_desc, newton_field, scalar) + + if cfg.joint_drive_props is not None: + source = cfg.joint_drive_props + for property_name in ( + "max_effort", + "max_velocity", + "friction", + "armature", + ): + configured = getattr(source, property_name) + if configured is None: + continue + matches = _joint_property_matches( + configured, + joint_names, + property_name=property_name, + control_parts=control_parts, + ) + for joint_name, value in matches: + if not isinstance(value, numbers.Number): + raise TypeError( + f"Articulation joint rule for {joint_name!r} and " + f"{property_name!r} must contain a numeric value." + ) + scalar = float(value) + default_desc, newton_desc = joint_properties[joint_name] + if property_name == "armature": + joint_common[joint_name]["armature"] = scalar + elif property_name == "max_effort": + default_desc.max_force = scalar + joint_common[joint_name]["effort_limit"] = scalar + elif property_name == "max_velocity": + default_desc.max_velocity = scalar + joint_common[joint_name]["velocity_limit"] = scalar + else: + default_field, newton_field = property_fields[property_name] + setattr(default_desc, default_field, scalar) + setattr(newton_desc, newton_field, scalar) + + # Solvers that ignore Newton's target-mode enum still consume drive gains. + # Masking inactive components makes NONE, EFFORT, and VELOCITY deterministic + # across the currently supported solver set. + for joint_name, target_mode in joint_target_modes.items(): + default_desc, newton_desc = joint_properties[joint_name] + if target_mode in {0, 4}: + default_desc.stiffness = 0.0 + default_desc.damping = 0.0 + newton_desc.target_ke = 0.0 + newton_desc.target_kd = 0.0 + elif target_mode == 2: + default_desc.stiffness = 0.0 + newton_desc.target_ke = 0.0 + + normalized_solver = ( + None + if newton_solver_type is None + else newton_solver_type.replace("-", "_").lower() + ) + if normalized_solver not in {None, "auto", "mujoco_warp", "mjwarp"} and any( + mode == 1 for mode in joint_target_modes.values() + ): + warnings.warn( + f"Newton solver {newton_solver_type!r} does not consume " + "joint_target_mode. POSITION is emulated with its configured " + "gains and assumes the velocity target remains zero.", + UserWarning, + stacklevel=3, + ) + + joint_limits = _compile_joint_limits(desc, cfg) + + return joint_properties, joint_common, joint_limits, joint_target_modes + + +def _joint_limit_array(value: object) -> np.ndarray: + """Convert a tensor/array/sequence limit value to a CPU NumPy array.""" + if hasattr(value, "detach"): + value = value.detach().cpu().numpy() + return np.asarray(value, dtype=np.float32) + + +def _compile_joint_limits( + desc: ArticulationDesc, + cfg: ArticulationCfg, +) -> dict[str, tuple[object, object]]: + """Compile regex or flattened-DOF joint limits before backend build.""" + joint_limits: dict[str, tuple[object, object]] = {} + if cfg.qpos_limits is None: + return joint_limits + + joint_names = [joint.name for joint in desc.joints] + if isinstance(cfg.qpos_limits, dict): + indices, _, values = resolve_matching_names_values( + cfg.qpos_limits, + joint_names, + ) + for index, limits in zip(indices, values): + limit_values = _joint_limit_array(limits).reshape(-1) + if limit_values.size != 2: + raise ValueError( + f"qpos_limits for {joint_names[index]!r} must contain " + "[lower, upper]." + ) + lower_limit, upper_limit = map(float, limit_values) + if not math.isfinite(lower_limit) or not math.isfinite(upper_limit): + raise ValueError( + f"qpos_limits for {joint_names[index]!r} must be finite." + ) + if lower_limit > upper_limit: + raise ValueError( + f"qpos_limits for {joint_names[index]!r} has lower limit " + f"{lower_limit} greater than upper limit {upper_limit}." + ) + joint_limits[joint_names[index]] = (lower_limit, upper_limit) + return joint_limits + + dof_joints = [joint for joint in desc.joints if joint.dof_count > 0] + dof_count = sum(joint.dof_count for joint in dof_joints) + limit_values = _joint_limit_array(cfg.qpos_limits) + expected_shape = (dof_count, 2) + if tuple(limit_values.shape) != expected_shape: + raise ValueError( + "Array qpos_limits must have flattened source-resolved DOF shape " + f"{expected_shape}, got {tuple(limit_values.shape)}." + ) + if not np.isfinite(limit_values).all(): + raise ValueError("Array qpos_limits must contain only finite values.") + if np.any(limit_values[:, 0] > limit_values[:, 1]): + raise ValueError( + "Array qpos_limits contains a lower limit greater than its upper limit." + ) + + dof_start = 0 + for joint in dof_joints: + dof_stop = dof_start + joint.dof_count + joint_values = limit_values[dof_start:dof_stop] + if joint.dof_count == 1: + lower_limit: object = float(joint_values[0, 0]) + upper_limit: object = float(joint_values[0, 1]) + else: + lower_limit = joint_values[:, 0].copy() + upper_limit = joint_values[:, 1].copy() + joint_limits[joint.name] = (lower_limit, upper_limit) + dof_start = dof_stop + return joint_limits + + +def _joint_property_matches( + configured: object, + joint_names: list[str], + *, + property_name: str, + numeric_only: bool = True, + control_parts: dict[str, Sequence[str]] | None = None, +) -> list[tuple[str, object]]: + """Resolve scalar, regex, and robot control-part drive rules.""" + scalar_types = (numbers.Number,) if numeric_only else (numbers.Number, str) + if isinstance(configured, scalar_types): + return [(name, configured) for name in joint_names] + if isinstance(configured, dict): + control_parts = control_parts or {} + part_rules = { + name: value for name, value in configured.items() if name in control_parts + } + direct_rules = { + name: value + for name, value in configured.items() + if name not in control_parts + } + + resolved: dict[str, object] = {} + owners: dict[str, str] = {} + for part_name, value in part_rules.items(): + expressions = list(control_parts[part_name]) + if not expressions: + raise ValueError(f"Robot control part {part_name!r} has no joints.") + indices, _, _ = resolve_matching_names_values( + {expression: value for expression in expressions}, + joint_names, + ) + for index in indices: + joint_name = joint_names[index] + previous = owners.get(joint_name) + if previous is not None: + raise ValueError( + f"Joint {joint_name!r} is selected by both control " + f"parts {previous!r} and {part_name!r} for drive " + f"property {property_name!r}." + ) + resolved[joint_name] = value + owners[joint_name] = part_name + + if direct_rules: + indices, _, values = resolve_matching_names_values( + direct_rules, + joint_names, + ) + # Exact/regex joint rules intentionally override a broader control + # part rule, matching RobotCfg's public configuration contract. + for index, value in zip(indices, values): + resolved[joint_names[index]] = value + return [(name, resolved[name]) for name in joint_names if name in resolved] + expected = "number" if numeric_only else "string/integer" + raise TypeError( + f"Articulation drive property {property_name!r} must be a {expected} " + f"or regex-to-{expected} mapping." + ) + + +def _compile_rigid_physics( + physics: _RigidPhysicsSpec, + body_type: str, +) -> RigidBodyPhysicsDesc: + actor_types = { + "dynamic": ActorType.DYNAMIC, + "kinematic": ActorType.KINEMATIC, + "static": ActorType.STATIC, + } + try: + actor_type = actor_types[body_type] + except KeyError as exc: + raise ValueError( + f"Unsupported rigid body_type {body_type!r}; expected one of " + f"{tuple(actor_types)}." + ) from exc + + mass_value = physics.mass_props.get("mass") + density_value = physics.mass_props.get("density") + if mass_value is not None and float(mass_value) < 0: + raise ValueError("Rigid-body mass cannot be negative.") + if density_value is not None and float(density_value) <= 0: + raise ValueError("Rigid-body density must be positive.") + if mass_value == 0 and density_value is None: + raise ValueError("Rigid-body density is required when mass is zero.") + + inertia = _rigid_array( + physics.mass_props.get("inertia"), + field_name="inertia", + allowed_sizes=(3, 9), + ) + if inertia is not None and physics.recompute_inertia: + raise ValueError( + "Rigid-body inertia cannot be explicit when recompute_inertia is true." + ) + com_position = _rigid_array( + physics.mass_props.get("com_position"), + field_name="com_position", + allowed_sizes=(3,), + ) + com_quaternion = _rigid_array( + physics.mass_props.get("com_quaternion"), + field_name="com_quaternion", + allowed_sizes=(4,), + ) + if inertia is not None: + if mass_value is None or float(mass_value) <= 0: + raise ValueError("Explicit rigid-body inertia requires a positive mass.") + if inertia.size == 3 and (np.any(inertia <= 0.0) or np.allclose(inertia, 0.0)): + raise ValueError( + "Rigid-body inertia must contain positive principal moments." + ) + if inertia.size == 9: + inertia_matrix = inertia.reshape(3, 3) + if not np.allclose(inertia_matrix, inertia_matrix.T, atol=1.0e-6): + raise ValueError("Rigid-body inertia matrix must be symmetric.") + if np.any(np.linalg.eigvalsh(inertia_matrix) <= 0.0): + raise ValueError("Rigid-body inertia matrix must be positive definite.") + if com_quaternion is not None: + quaternion_norm = float(np.linalg.norm(com_quaternion)) + if quaternion_norm <= 1.0e-8: + raise ValueError("Rigid-body com_quaternion cannot be zero.") + com_quaternion = com_quaternion / quaternion_norm + # DexSim descriptors use wxyz; EmbodiChain configuration uses xyzw. + com_quaternion = convert_quat(com_quaternion, to="wxyz") + + if body_type != "static": + mass = ( + float(mass_value) + if mass_value is not None and float(mass_value) > 0 + else None + ) + density = ( + float(density_value) + if mass is None and density_value is not None and float(density_value) > 0 + else None + ) + else: + # Both backends ignore mass properties on static actors. Omitting them + # also avoids a Newton build warning for the common default cfg. + mass = None + density = None + inertia = None + com_position = None + com_quaternion = None + + if physics.default_rigid_props: + default_values = {item.name: None for item in fields(DexsimPhysicsDesc)} + default_values.update(physics.default_rigid_props) + default_desc = DexsimPhysicsDesc(**default_values) + else: + default_desc = None + return RigidBodyPhysicsDesc( + actor_type=actor_type, + mass=mass, + density=density, + inertia=inertia, + com_position=com_position, + com_quaternion=com_quaternion, + dexsim=default_desc, + newton=None, + ) + + +def _rigid_array( + value: object | None, + *, + field_name: str, + allowed_sizes: tuple[int, ...], +) -> np.ndarray | None: + """Validate and copy a rigid-body mass-property array.""" + if value is None: + return None + result = np.asarray(value, dtype=np.float32).reshape(-1) + if result.size not in allowed_sizes or not np.all(np.isfinite(result)): + expected = " or ".join(str(size) for size in allowed_sizes) + raise ValueError( + f"Rigid-body {field_name} must contain {expected} finite values." + ) + return result.copy() + + +def _common_collision_envelope( + physics: _RigidPhysicsSpec, +) -> tuple[float | None, float | None]: + """Validate and return the portable contact/rest envelope.""" + + def optional_float(value: object | None, field_name: str) -> float | None: + if value is None: + return None + try: + result = float(value) + except (TypeError, ValueError) as exc: + raise TypeError(f"{field_name} must be a finite number.") from exc + if not math.isfinite(result): + raise ValueError(f"{field_name} must be finite.") + return result + + contact_offset = optional_float(physics.contact_offset, "contact_offset") + rest_offset = optional_float(physics.rest_offset, "rest_offset") + if contact_offset is not None and contact_offset < 0.0: + raise ValueError("contact_offset must be non-negative.") + if ( + contact_offset is not None + and rest_offset is not None + and contact_offset < rest_offset + ): + raise ValueError("contact_offset must be no smaller than rest_offset.") + return contact_offset, rest_offset + + +def _compile_default_collision( + physics: _RigidPhysicsSpec, +) -> DexsimCollisionDesc | None: + values = dict(physics.material_props) + contact_offset, rest_offset = _common_collision_envelope(physics) + if contact_offset is not None: + values["contact_offset"] = contact_offset + if rest_offset is not None: + values["rest_offset"] = rest_offset + values.update(physics.default_collision_props) + if not values: + return None + configured = {item.name: None for item in fields(DexsimCollisionDesc)} + configured.update(values) + return DexsimCollisionDesc(**configured) + + +def _compile_newton_collision( + physics: _RigidPhysicsSpec, + *, + mesh_collision: MeshCollisionCfg | None = None, + newton_solver_type: str | None = None, + author_shape_defaults: bool = False, +) -> NewtonCollisionDesc | None: + # Keep partial descriptors sparse for source overlays. Once a newly authored + # shape has a Newton override, fill the Spawn margin/gap defaults because a + # non-None descriptor suppresses DexSim's descriptor factory defaults. + values = {field.name: None for field in fields(NewtonCollisionDesc)} + contact_offset, rest_offset = _common_collision_envelope(physics) + native_margin = physics.newton_collision_props.get("margin") + native_gap = physics.newton_collision_props.get("gap") + if rest_offset is not None: + values["margin"] = rest_offset + if contact_offset is not None and native_gap is None: + effective_margin = native_margin if native_margin is not None else rest_offset + if effective_margin is None: + if newton_solver_type is not None: + raise ValueError( + "Newton requires rest_offset (or a native margin) when a " + "portable contact_offset is configured." + ) + else: + try: + gap = contact_offset - float(effective_margin) + except (TypeError, ValueError) as exc: + raise TypeError( + "Newton collision margin must be a finite number." + ) from exc + if not math.isfinite(gap): + raise ValueError("Newton collision margin must be finite.") + if gap < 0.0: + raise ValueError( + "Newton collision margin must be no larger than contact_offset." + ) + values["gap"] = gap + values.update(physics.newton_collision_props) + if mesh_collision is not None and mesh_collision.approximation == "sdf": + values["force_sdf"] = True + for field_name in ( + "is_hydroelastic", + "sdf_narrow_band_range", + "sdf_target_voxel_size", + "sdf_texture_format", + "sdf_padding", + ): + value = getattr(mesh_collision, field_name) + if value is not None: + values[field_name] = value + if mesh_collision.sdf_resolution is not None: + values["sdf_max_resolution"] = int(mesh_collision.sdf_resolution) + values.update(physics.newton_material_props) + dynamic_friction = physics.material_props.get("dynamic_friction") + if dynamic_friction is not None: + values["mu"] = float(dynamic_friction) + solver_contact_fields = NEWTON_CONTACT_SOLVER_FIELDS.get(newton_solver_type) + restitution = physics.material_props.get("restitution") + if restitution is not None and ( + solver_contact_fields is None or "restitution" in solver_contact_fields + ): + values["restitution"] = float(restitution) + if all(value is None for value in values.values()): + return None + if author_shape_defaults: + defaults = NewtonCollisionDesc() + if values["margin"] is None: + values["margin"] = defaults.margin + if values["gap"] is None: + values["gap"] = defaults.gap + return NewtonCollisionDesc(**values) + + +def _compile_geometry( + cfg: RigidObjectCfg, +) -> tuple[GeometryDesc, CollisionApproximation, int]: + shape = cfg.shape + if isinstance(shape, MeshCfg): + if _is_missing(shape.fpath) or not str(shape.fpath).strip(): + raise ValueError("MeshCfg.fpath must be a non-empty path.") + collision_cfg = shape.collision or MeshCollisionCfg() + approximation = { + "convex_hull": CollisionApproximation.CONVEX_HULL, + "convex_decomposition": CollisionApproximation.CONVEX_DECOMPOSITION, + "triangle_mesh": CollisionApproximation.NONE, + "sdf": CollisionApproximation.SDF, + }[collision_cfg.approximation] + max_hulls = collision_cfg.max_hulls or 1 + acd_method = collision_cfg.acd_method or "coacd" + + if collision_cfg.approximation == "triangle_mesh" and cfg.body_type != "static": + raise ValueError( + "triangle_mesh collision is supported only for static rigid objects." + ) + + if shape.compute_uv: + logger.log_warning( + "Mesh UV projection is not represented by GeometryDesc and was " + "not applied." + ) + if ( + collision_cfg.approximation == "convex_decomposition" + and acd_method != "coacd" + ): + raise ValueError( + "Spawn supports only acd_method='coacd' for convex_decomposition." + ) + if collision_cfg.sdf_resolution is not None: + logger.log_warning( + "CollisionApproximation.SDF is preserved and Newton receives " + "sdf_max_resolution, but the DexSim descriptor does not expose " + "its cooking resolution." + ) + return ( + GeometryDesc.mesh( + file_path=str(shape.fpath), segment_name=cfg.uid or "mesh" + ), + approximation, + max(1, max_hulls), + ) + + if isinstance(shape, CubeCfg): + size = tuple(float(value) for value in shape.size) + if len(size) != 3 or any(value <= 0 for value in size): + raise ValueError("CubeCfg.size must contain three positive values.") + return GeometryDesc.cube(size), CollisionApproximation.NONE, 1 + + if isinstance(shape, SphereCfg): + if shape.radius <= 0: + raise ValueError("SphereCfg.radius must be positive.") + return ( + GeometryDesc.sphere(float(shape.radius)), + CollisionApproximation.NONE, + 1, + ) + + raise NotImplementedError( + f"RigidObjectCfg shape {type(shape).__name__!r} is not supported by " + "the Spawn converter; supported shapes are MeshCfg, CubeCfg, and SphereCfg." + ) + + +def _compile_load_option(shape: object) -> DexsimLoadOption | None: + """Translate mesh import options without leaking EmbodiChain config types.""" + if not isinstance(shape, MeshCfg): + return None + source = shape.load_option + option = DexsimLoadOption() + option.rebuild_normals = bool(source.rebuild_normals) + option.rebuild_tangent = bool(source.rebuild_tangent) + option.rebuild_3rdnormal = bool(source.rebuild_3rdnormal) + option.rebuild_3rdtangent = bool(source.rebuild_3rdtangent) + option.smooth = float(source.smooth) + return option + + +def _compile_visual_material( + object_uid: str, + cfg: VisualMaterialCfg | None, +) -> tuple[str | None, tuple[str, MaterialDesc] | None]: + if cfg is None: + return None, None + key = str(cfg.uid or f"{object_uid}_material") + base_color = tuple(float(value) for value in cfg.base_color) + if len(base_color) != 4: + raise ValueError("VisualMaterialCfg.base_color must be RGBA.") + emissive_rgb = tuple( + float(value) * float(cfg.emissive_intensity) for value in cfg.emissive + ) + if len(emissive_rgb) != 3: + raise ValueError("VisualMaterialCfg.emissive must be RGB.") + desc = MaterialDesc( + name=key, + base_color=base_color, + base_color_map=cfg.base_color_texture, + normal_map=cfg.normal_texture, + emissive=(*emissive_rgb, 1.0), + roughness=float(cfg.roughness), + roughness_map=cfg.roughness_texture, + metallic=float(cfg.metallic), + metallic_map=cfg.metallic_texture, + ao_map=cfg.ao_texture, + ior=float(cfg.ior), + ) + return key, (key, desc) + + +def _pose_from_cfg(cfg: object) -> np.ndarray: + local_pose = getattr(cfg, "init_local_pose", None) + if local_pose is not None: + pose = np.asarray(local_pose, dtype=np.float32).reshape(4, 4).copy() + else: + position = _vector3(getattr(cfg, "init_pos"), field_name="init_pos") + rotation_deg = _vector3(getattr(cfg, "init_rot"), field_name="init_rot") + rx, ry, rz = np.deg2rad(rotation_deg) + cx, sx = math.cos(rx), math.sin(rx) + cy, sy = math.cos(ry), math.sin(ry) + cz, sz = math.cos(rz), math.sin(rz) + rot_x = np.array( + ((1.0, 0.0, 0.0), (0.0, cx, -sx), (0.0, sx, cx)), + dtype=np.float32, + ) + rot_y = np.array( + ((cy, 0.0, sy), (0.0, 1.0, 0.0), (-sy, 0.0, cy)), + dtype=np.float32, + ) + rot_z = np.array( + ((cz, -sz, 0.0), (sz, cz, 0.0), (0.0, 0.0, 1.0)), + dtype=np.float32, + ) + pose = np.eye(4, dtype=np.float32) + # Match EmbodiChain's shared matrix_from_euler(..., "XYZ") contract + # used by the legacy RigidObject reset path. + pose[:3, :3] = rot_x @ rot_y @ rot_z + pose[:3, 3] = position + + if not np.isfinite(pose).all(): + raise ValueError("init_local_pose must contain finite values.") + if not np.allclose(pose[3], (0.0, 0.0, 0.0, 1.0), atol=1e-6): + raise ValueError("init_local_pose must be a homogeneous 4x4 transform.") + return pose + + +def _vector3(value: object, *, field_name: str) -> np.ndarray: + result = np.asarray(value, dtype=np.float32).reshape(-1) + if result.size != 3 or not np.isfinite(result).all(): + raise ValueError(f"{field_name} must contain three finite values.") + if field_name == "body_scale" and np.any(result <= 0): + raise ValueError("body_scale values must be positive.") + return result.copy() + + +def _required_uid(value: str | None, label: str) -> str: + if value is None or not str(value).strip(): + raise ValueError(f"{label} uid must be specified before Spawn conversion.") + uid = str(value) + if "/" in uid: + raise ValueError(f"{label} uid cannot contain '/': {uid!r}.") + return uid + + +def _articulation_uid(value: str | None, path: str | None) -> str: + if value is not None and str(value).strip(): + return _required_uid(str(value), "Articulation") + if path is None or not str(path).strip(): + raise ValueError( + "Articulation uid is required when its source path is unresolved." + ) + inferred = os.path.splitext(os.path.basename(str(path)))[0] + return _required_uid(inferred, "Articulation") + + +def _is_usd_path(path: object) -> bool: + return str(path).lower().endswith((".usd", ".usda", ".usdc")) + + +def _is_missing(value: object) -> bool: + # ``@configclass`` deepcopy can create a distinct _MISSING_TYPE instance. + return value is MISSING or isinstance(value, type(MISSING)) diff --git a/embodichain/lab/sim/spawn/scene.py b/embodichain/lab/sim/spawn/scene.py new file mode 100644 index 000000000..9d0e19d6e --- /dev/null +++ b/embodichain/lab/sim/spawn/scene.py @@ -0,0 +1,322 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +"""Thin EmbodiChain coordination around DexSim Spawn.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Literal + +__all__ = ["SpawnScene"] + +_AssetKind = Literal[ + "rigid_object", + "rigid_object_group", + "articulation", + "soft_object", + "cloth_object", +] + + +@dataclass(slots=True) +class _AssetDeclaration: + kind: _AssetKind + descriptor: Any + facade: Any | None + source_configurator: Callable[[Any], None] | None = None + + +class SpawnScene: + """Map EmbodiChain asset declarations onto one DexSim Spawn scene. + + DexSim owns declaration materialization, stable handles, and topology + revisions. EmbodiChain resolves and configures source metadata before the + first backend build so Newton does not materialize an articulation twice. + """ + + def __init__( + self, + world: Any, + *, + num_envs: int, + spacing: tuple[float, float, float] = (0.0, 0.0, 0.0), + ) -> None: + from dexsim.spawn import SceneBuilder + + self.builder = SceneBuilder(world) + self.builder.replicate( + count=num_envs, + spacing=spacing, + name_format="arena_{i}", + collision_policy="isolated", + ) + self._assets: dict[str, _AssetDeclaration] = {} + + @property + def arena_names(self) -> tuple[str, ...]: + """Names of the replicated per-environment Arenas.""" + return tuple(self.builder.replicate_plan.env_names()) + + def __contains__(self, uid: str) -> bool: + return uid in self._assets + + def declare( + self, + kind: _AssetKind, + uid: str, + descriptor: Any, + *, + facade: Any | None = None, + configure_source: Callable[[Any], None] | None = None, + ) -> None: + """Add a descriptor and associate it with an EmbodiChain facade.""" + if uid in self._assets: + raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") + declaration = _AssetDeclaration( + kind=kind, + descriptor=descriptor, + facade=facade, + source_configurator=configure_source, + ) + + if kind == "rigid_object_group": + declaration.descriptor = tuple( + self.builder.add_object(member) for member in descriptor + ) + else: + if ( + kind == "articulation" + and configure_source is not None + and (self.builder.is_finalized or self.builder.result is not None) + and self._can_resolve_before_materialization() + ): + self._resolve_articulation_source(descriptor) + configure_source(descriptor) + declaration.source_configurator = None + add_name = { + "rigid_object": "add_object", + "articulation": "add_articulation", + "soft_object": "add_soft_object", + "cloth_object": "add_cloth_object", + }[kind] + declaration.descriptor = getattr(self.builder, add_name)(descriptor) + self._assets[uid] = declaration + self._configure_materialized_source(uid) + handles = self.handles(uid) + if facade is not None and handles: + facade.attach_spawn_handles(handles) + + def resolve_sources(self) -> None: + """Resolve and configure declarations before backend materialization.""" + if self.builder.is_finalized: + return + + builder_resolver = getattr(self.builder, "resolve_sources", None) + if builder_resolver is not None: + builder_resolver() + elif getattr(self.builder, "backend", None) == "newton": + for declaration in self._assets.values(): + if ( + declaration.kind == "articulation" + and declaration.source_configurator is not None + ): + self._resolve_articulation_source(declaration.descriptor) + else: + return + + for declaration in self._assets.values(): + configure = declaration.source_configurator + if configure is None: + continue + configure(declaration.descriptor) + declaration.source_configurator = None + + def track( + self, + kind: _AssetKind, + uid: str, + descriptor: Any, + *, + facade: Any | None = None, + ) -> None: + """Track a descriptor that was already added to ``SceneBuilder``.""" + if uid in self._assets: + raise ValueError(f"Spawn asset uid is already declared: {uid!r}.") + declaration = _AssetDeclaration(kind, descriptor, facade) + self._assets[uid] = declaration + handles = self.handles(uid) + if facade is not None and handles: + facade.attach_spawn_handles(handles) + + def remove(self, uid: str) -> None: + """Remove a declared asset from its DexSim owner.""" + declaration = self._assets[uid] + if declaration.kind in {"soft_object", "cloth_object"}: + raise NotImplementedError( + "DexSim Spawn does not yet expose pending removal for " + f"{declaration.kind.replace('_', ' ')}." + ) + if declaration.kind == "rigid_object_group": + for member in declaration.descriptor: + self.builder.remove_object(member.name) + else: + remove_name = { + "rigid_object": "remove_object", + "articulation": "remove_articulation", + }[declaration.kind] + removed = getattr(self.builder, remove_name)(declaration.descriptor.name) + if removed is None: + raise KeyError(f"Spawn asset is absent from SceneBuilder: {uid!r}.") + del self._assets[uid] + + def commit(self) -> Any: + """Finalize once or let ``SpawnResult`` consume pending changes.""" + if not self.builder.is_finalized: + self.resolve_sources() + result = self.builder.finalize() + else: + result = self.builder.result + assert result is not None + if self.builder.has_pending_changes or result.needs_rebuild: + result = result.rebuild(self.builder) + + for uid in self._assets: + self._configure_materialized_source(uid) + self.builder.result = result + return result + + def bind(self) -> None: + """Complete post-finalize runtime binding for declared facades. + + Native entity creation belongs to ``SceneBuilder`` and its backend + adapter. This method only attaches handles that were unavailable during + declaration, then lets each facade create its result-dependent + Batch/Data state through ``bind_spawn()``. Eager Default handles may + already be attached; deferred Newton handles are resolved here. + """ + result = self.builder.result + if result is None or not self.builder.is_finalized: + raise RuntimeError("Spawn scene must be materialized before binding.") + + for uid, declaration in self._assets.items(): + facade = declaration.facade + if facade is None or not facade.is_declared: + continue + if not facade._entities: + facade.attach_spawn_handles(self.handles(uid)) + facade.bind_spawn(result) + + def prepare_runtime_config(self, result: Any) -> None: + """Apply facade configuration required before backend initialization. + + Default Direct GPU simulation snapshots some native articulation + properties during initialization. Articulation facades therefore get + a narrow pre-bind hook after materialization but before the manager + initializes backend runtime buffers. + """ + if result is not self.builder.result or not self.builder.is_finalized: + raise RuntimeError("Spawn scene must be materialized before runtime setup.") + + for uid, declaration in self._assets.items(): + facade = declaration.facade + if facade is None or declaration.kind != "articulation": + continue + if not facade._entities: + facade.attach_spawn_handles(self.handles(uid)) + facade._prepare_spawn_runtime_config(result) + + def close(self) -> None: + """Release Spawn resources and facade references.""" + result = self.builder.result + if result is not None: + result.close() + self.builder.result = None + self._assets.clear() + + def handles(self, uid: str) -> tuple[Any, ...]: + """Return currently materialized handles for one logical asset.""" + result = self.builder.result + if result is None: + return () + declaration = self._assets[uid] + if declaration.kind == "rigid_object_group": + paths = tuple( + f"{arena}/{member.name}" + for arena in self.arena_names + for member in declaration.descriptor + ) + elif declaration.descriptor.per_env: + paths = tuple( + f"{arena}/{declaration.descriptor.name}" for arena in self.arena_names + ) + else: + paths = (declaration.descriptor.name,) + if any(path not in result.handles for path in paths): + return () + return tuple(result.handles[path] for path in paths) + + def _resolve_articulation_source(self, descriptor: Any) -> None: + """Resolve one descriptor through the available DexSim boundary.""" + builder_resolver = getattr( + self.builder, + "resolve_articulation_source", + None, + ) + if builder_resolver is not None: + builder_resolver(descriptor) + return + + from embodichain.lab.sim.spawn.source import resolve_articulation_source + + resolve_articulation_source(self.builder, descriptor) + + def _can_resolve_before_materialization(self) -> bool: + """Return whether exact source metadata is available before add.""" + return ( + getattr(self.builder, "resolve_articulation_source", None) is not None + or getattr(self.builder, "backend", None) == "newton" + ) + + def _configure_materialized_source(self, uid: str) -> None: + """Apply a pending source config to an eager Default articulation.""" + declaration = self._assets[uid] + configure = declaration.source_configurator + if configure is None or declaration.kind != "articulation": + return + + handles = self.handles(uid) + if not handles: + return + + result = self.builder.result + assert result is not None + if result.backend != "dexsim": + raise RuntimeError( + "Newton articulation source configuration must run before " + "SceneBuilder.finalize()." + ) + + prototype = declaration.descriptor + source = ( + prototype + if getattr(prototype, "links", None) or getattr(prototype, "joints", None) + else handles[0].articulation_desc + ) + configure(source) + for handle in handles: + handle.apply_dexsim_properties(source) + declaration.source_configurator = None diff --git a/embodichain/lab/sim/spawn/source.py b/embodichain/lab/sim/spawn/source.py new file mode 100644 index 000000000..9bbcd1e89 --- /dev/null +++ b/embodichain/lab/sim/spawn/source.py @@ -0,0 +1,116 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Resolve Newton articulation metadata before its first physics build.""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any + +import numpy as np +from dexsim.spawn import ArticulationDesc + +if TYPE_CHECKING: + from dexsim.spawn import SceneBuilder + +__all__ = ["resolve_articulation_source"] + + +def resolve_articulation_source( + builder: SceneBuilder, + desc: ArticulationDesc, +) -> ArticulationDesc: + """Populate exact URDF metadata without building a Newton model. + + DexSim 0.4.3 removed its public source-resolution phase while retaining + the same URDF-to-descriptor translator inside the Newton adapter. This + compatibility boundary invokes that translator with a disposable + render-only skeleton, allowing name-dependent EmbodiChain overlays to be + authored before :meth:`SceneBuilder.finalize`. + + Args: + builder: Scene builder that owns the target arena layout. + desc: Articulation descriptor to resolve in place. + + Returns: + The resolved descriptor. + """ + signature = _source_signature(desc) + previous = getattr(desc, "_embodichain_source_signature", None) + if previous == signature: + return desc + + if desc.urdf_path is None: + setattr(desc, "_embodichain_source_signature", signature) + return desc + + if previous is not None: + desc.links = [] + desc.joints = [] + desc.root_link_name = None + + arena = _source_arena(builder, desc) + temp_name = f"__embodichain_resolve__{desc.name.replace('/', '__')}__{id(desc)}" + skeleton = arena.create_skeleton("skeleton") + if skeleton is None: + raise RuntimeError(f"Failed to create a source resolver for {desc.name!r}.") + skeleton.set_name(temp_name) + skeleton.detach_parent() + try: + scale = np.asarray(desc.body_scale, dtype=np.float32).reshape(3) + load_result = skeleton.load_urdf(os.path.abspath(desc.urdf_path), scale) + if load_result != 0: + raise RuntimeError( + f"Skeleton.load_urdf({desc.urdf_path!r}) failed: {load_result}" + ) + + # DexSim currently exposes no public metadata-only resolver. Reuse the + # adapter's source translator so its retained descriptor semantics stay + # identical to the subsequent Newton build. + from dexsim.spawn.adapters.newton_articulation_adapter import ( + _translate_urdf_articulation, + ) + + _translate_urdf_articulation(skeleton, desc) + finally: + # Drop the wrapper before deleting its Arena-owned native object. + skeleton = None + arena.remove_skeleton(temp_name) + + setattr(desc, "_embodichain_source_signature", signature) + return desc + + +def _source_signature(desc: ArticulationDesc) -> tuple[object, ...]: + if desc.urdf_path is None: + return "explicit", id(desc) + return ( + "urdf", + os.path.abspath(desc.urdf_path), + tuple(float(value) for value in np.asarray(desc.body_scale).reshape(3)), + ) + + +def _source_arena(builder: SceneBuilder, desc: ArticulationDesc) -> Any: + if desc.per_env and builder.replicate_plan is not None: + arenas = builder.prepare_arenas() + if not arenas: + raise RuntimeError( + f"No replicated Arena is available to resolve {desc.name!r}." + ) + return arenas[0] + return builder.world.get_env() diff --git a/embodichain/lab/sim/spawn/usd.py b/embodichain/lab/sim/spawn/usd.py new file mode 100644 index 000000000..84d68bf49 --- /dev/null +++ b/embodichain/lab/sim/spawn/usd.py @@ -0,0 +1,245 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- +"""Compatibility translation for EmbodiChain's singleton USD APIs.""" + +from __future__ import annotations + +import os +from dataclasses import fields, replace +from typing import TypeVar + +from dexsim.spawn import ( + ArticulationDesc, + CollisionDesc, + MaterialDesc, + ObjectDesc, + RenderDesc, + RigidBodyPhysicsDesc, +) +from dexsim.types import ActorType + +from embodichain.lab.sim.cfg import ArticulationCfg, RigidObjectCfg +from embodichain.lab.sim.spawn.descriptors import ( + _compile_default_collision, + _compile_newton_collision, + _compile_rigid_physics, + _compile_visual_material, + _articulation_root_values, + _pose_from_cfg, + _required_uid, + _resolve_rigid_physics, + _validate_articulation_rigid_physics, + _vector3, +) + +__all__ = ["articulation_desc_from_usd", "rigid_desc_from_usd"] + +_PropertyCfgT = TypeVar("_PropertyCfgT") + + +def _overlay_optional_properties( + source: _PropertyCfgT | None, + configured: _PropertyCfgT | None, +) -> _PropertyCfgT | None: + """Overlay non-None dataclass fields without erasing source values.""" + if configured is None: + return source + if source is None: + return configured + for item in fields(configured): + value = getattr(configured, item.name) + if value is not None: + setattr(source, item.name, value) + return source + + +def _overlay_rigid_body_properties( + source: RigidBodyPhysicsDesc | None, + configured: RigidBodyPhysicsDesc, + *, + recompute_inertia: bool = False, +) -> RigidBodyPhysicsDesc: + """Merge a partial body config into properties parsed from USD.""" + if source is None: + return configured + source.actor_type = configured.actor_type + source.dexsim = _overlay_optional_properties(source.dexsim, configured.dexsim) + source.newton = _overlay_optional_properties(source.newton, configured.newton) + if configured.mass is not None: + source.mass = configured.mass + source.density = None + elif configured.density is not None: + source.mass = None + source.density = configured.density + if recompute_inertia: + source.inertia = None + for name in ("inertia", "com_position", "com_quaternion"): + value = getattr(configured, name) + if value is not None: + setattr(source, name, value) + return source + + +def _overlay_collision_properties( + source: CollisionDesc, + configured: CollisionDesc, +) -> None: + """Merge partial contact properties while retaining parsed geometry.""" + if configured.enable_collision is not None: + source.enable_collision = configured.enable_collision + source.dexsim = _overlay_optional_properties(source.dexsim, configured.dexsim) + source.newton = _overlay_optional_properties(source.newton, configured.newton) + + +def rigid_desc_from_usd( + cfg: RigidObjectCfg, + *, + per_env: bool = True, + newton_solver_type: str | None = None, +) -> tuple[ObjectDesc, dict[str, MaterialDesc]]: + """Select the sole rigid object in a USD stage.""" + uid = _required_uid(cfg.uid, "Rigid object") + path = getattr(cfg.shape, "fpath", None) + scene, desc = _parse_singleton(path, "mesh_objects", "rigid object") + + desc.name = uid + desc.pose = _pose_from_cfg(cfg) + desc.per_env = per_env + materials = _namespace_materials(desc.renders, scene.materials, uid) + + if cfg.resolve_asset_physics_mode() == "preserve": + if desc.physics is None: + raise ValueError(f"USD rigid object {path!r} has no physics.") + cfg.body_type = { + ActorType.DYNAMIC: "dynamic", + ActorType.KINEMATIC: "kinematic", + ActorType.STATIC: "static", + }[desc.physics.actor_type] + cfg.body_scale = tuple(float(value) for value in desc.body_scale) + return desc, materials + + physics = _resolve_rigid_physics( + cfg.attrs, + newton_solver_type=newton_solver_type, + ) + configured_body = _compile_rigid_physics(physics, cfg.body_type) + desc.physics = _overlay_rigid_body_properties( + desc.physics, + configured_body, + recompute_inertia=bool(physics.recompute_inertia), + ) + desc.body_scale = _vector3(cfg.body_scale, field_name="body_scale") + for collision in desc.collisions: + _overlay_collision_properties( + collision, + CollisionDesc( + enable_collision=physics.collision_enabled, + dexsim=_compile_default_collision(physics), + newton=_compile_newton_collision( + physics, + newton_solver_type=newton_solver_type, + ), + ), + ) + + material_ref, material_entry = _compile_visual_material( + uid, + cfg.shape.visual_material, + ) + if material_entry is not None: + materials = {material_entry[0]: material_entry[1]} + for render in desc.renders: + render.material = None + render.material_ref = material_ref + return desc, materials + + +def articulation_desc_from_usd( + cfg: ArticulationCfg, + *, + per_env: bool = True, + source_path: str | None = None, + newton_solver_type: str | None = None, +) -> tuple[ArticulationDesc, dict[str, MaterialDesc]]: + """Select the sole articulation in a USD stage.""" + preserve_asset_physics = cfg.resolve_asset_physics_mode() == "preserve" + if not preserve_asset_physics: + _validate_articulation_rigid_physics( + cfg, + newton_solver_type=newton_solver_type, + ) + path = source_path or cfg.fpath + scene, desc = _parse_singleton(path, "articulations", "articulation") + uid = _required_uid( + cfg.uid or os.path.splitext(os.path.basename(str(path)))[0], + "Articulation", + ) + cfg.uid = uid + desc.name = uid + desc.pose = _pose_from_cfg(cfg) + desc.per_env = per_env + renders = [visual for link in desc.links for visual in link.visuals] + materials = _namespace_materials(renders, scene.materials, uid) + + if preserve_asset_physics: + cfg.body_scale = tuple(float(value) for value in desc.body_scale) + else: + desc.body_scale = _vector3(cfg.body_scale, field_name="body_scale") + desc.fixed_base, desc.enable_self_collision = _articulation_root_values( + cfg, + fixed_base_default=bool(desc.fixed_base), + self_collision_default=desc.enable_self_collision, + ) + return desc, materials + + +def _parse_singleton(path: object, collection: str, label: str): + if path is None: + raise ValueError(f"A USD path is required for the {label}.") + + from dexsim.kit.usd import parse_usd + + scene = parse_usd(str(path)) + candidates = getattr(scene, collection) + if len(candidates) != 1: + found = [ + (item.name, None if item.usd is None else item.usd.prim_path) + for item in candidates + ] + raise ValueError( + f"Expected exactly one {label} in USD file {path!r}, found " + f"{len(candidates)}: {found}." + ) + return scene, candidates[0] + + +def _namespace_materials( + renders: list[RenderDesc], + materials: dict[str, MaterialDesc], + uid: str, +) -> dict[str, MaterialDesc]: + selected = {} + for render in renders: + if render.material_ref is None: + continue + source_ref = render.material_ref + material = materials[source_ref] + render.material_ref = f"{uid}::{source_ref}" + selected[render.material_ref] = replace( + material, + name=f"{uid}::{material.name}", + ) + return selected diff --git a/embodichain/lab/sim/utility/cfg_utils.py b/embodichain/lab/sim/utility/cfg_utils.py index 51ce7d028..62fe7b1d8 100644 --- a/embodichain/lab/sim/utility/cfg_utils.py +++ b/embodichain/lab/sim/utility/cfg_utils.py @@ -14,9 +14,38 @@ # limitations under the License. # ---------------------------------------------------------------------------- -from embodichain.lab.sim.cfg import RobotCfg +from typing import TypeVar + +from embodichain.lab.sim.cfg import ( + _raise_removed_articulation_cfg_fields, + JointDrivePropertiesCfg, + RigidBodyPhysicsCfg, + RobotCfg, +) +from embodichain.lab.sim.cfg.rigid import _rigid_body_physics_from_dict from embodichain.lab.sim.solvers import SolverCfg -from embodichain.utils import logger +from embodichain.utils import is_configclass, logger + +_ConfigT = TypeVar("_ConfigT") + + +def _merge_non_none_config(base: _ConfigT | None, override: _ConfigT) -> _ConfigT: + """Merge non-None configclass fields without discarding base defaults.""" + if base is None: + return override + for field_name in override.__dataclass_fields__: + value = getattr(override, field_name) + if value is not None: + base_value = getattr(base, field_name) + if ( + base_value is not None + and type(base_value) is type(value) + and is_configclass(base_value) + ): + _merge_non_none_config(base_value, value) + else: + setattr(base, field_name, value) + return base def merge_solver_cfg( @@ -83,6 +112,8 @@ def merge_robot_cfg(base_cfg: RobotCfg, override_cfg_dict: dict[str, any]) -> Ro RobotCfg: The merged robot configuration. """ + _raise_removed_articulation_cfg_fields(override_cfg_dict) + # Only parse keys the base RobotCfg recognizes, so subclass-only variant # fields (version, ...) set by _build_defaults don't trigger # spurious "Key not found in RobotCfg" warnings from the base from_dict. @@ -142,30 +173,47 @@ def merge_robot_cfg(base_cfg: RobotCfg, override_cfg_dict: dict[str, any]) -> Ro f"new solver entry, or ensure the part name " f"matches an existing solver." ) - elif key == "drive_pros": + elif key == "joint_drive_props": # merge joint drive properties - user_drive_pros_dict = override_cfg_dict.get("drive_pros") - if isinstance(user_drive_pros_dict, dict): - for prop, val in user_drive_pros_dict.items(): + user_joint_drive_props_dict = override_cfg_dict.get("joint_drive_props") + if isinstance(user_joint_drive_props_dict, dict): + if user_joint_drive_props_dict.get("backend") == "newton": + base_cfg.joint_drive_props = JointDrivePropertiesCfg.from_dict( + user_joint_drive_props_dict, + defaults=base_cfg.joint_drive_props, + ) + continue + for prop, val in user_joint_drive_props_dict.items(): + if prop == "backend": + continue # Get the current value in cfg (which has defaults) - default_val = getattr(base_cfg.drive_pros, prop, None) + default_val = getattr(base_cfg.joint_drive_props, prop, None) if isinstance(val, dict) and isinstance(default_val, dict): # Merge dictionaries default_val.update(val) else: # Overwrite if not both dicts - setattr(base_cfg.drive_pros, prop, val) + setattr(base_cfg.joint_drive_props, prop, val) else: logger.log_warning( - "drive_pros should be a dictionary. Skipping drive_pros merge." + "joint_drive_props should be a dictionary. Skipping joint_drive_props merge." ) elif key == "attrs": # merge physics attributes user_attrs_dict = override_cfg_dict.get("attrs") if isinstance(user_attrs_dict, dict): - for attr_key, attr_val in user_attrs_dict.items(): - setattr(base_cfg.attrs, attr_key, attr_val) + grouped_fields = set(RigidBodyPhysicsCfg.__dataclass_fields__) + parsed = _rigid_body_physics_from_dict(user_attrs_dict) + for field_name in grouped_fields: + override = getattr(parsed, field_name) + if override is None: + continue + base = getattr(base_cfg.attrs, field_name) + if base is not None and type(base) is type(override): + _merge_non_none_config(base, override) + else: + setattr(base_cfg.attrs, field_name, override) else: logger.log_warning( "attrs should be a dictionary. Skipping attrs merge." diff --git a/embodichain/lab/sim/utility/keyboard_utils.py b/embodichain/lab/sim/utility/keyboard_utils.py index d64eca180..a7623e4ab 100644 --- a/embodichain/lab/sim/utility/keyboard_utils.py +++ b/embodichain/lab/sim/utility/keyboard_utils.py @@ -220,8 +220,7 @@ def run_keyboard_control_for_camera( quaternion = rot.as_quat() log_info("Current Camera pose:") log_info(f"Translation: {translation}") - quat_wxyz = [quaternion[3], quaternion[0], quaternion[1], quaternion[2]] - log_info(f"Quaternion (w, x, y, z): {quat_wxyz}") + log_info(f"Quaternion (x, y, z, w): {quaternion.tolist()}") rotation_euler = rot.as_euler("xyz", degrees=True) log_info(f"Rotation (XYZ Euler, degrees): {rotation_euler}") diff --git a/embodichain/lab/sim/utility/sim_utils.py b/embodichain/lab/sim/utility/sim_utils.py index dc92b8563..a2bbefc04 100644 --- a/embodichain/lab/sim/utility/sim_utils.py +++ b/embodichain/lab/sim/utility/sim_utils.py @@ -14,20 +14,24 @@ # limitations under the License. # ---------------------------------------------------------------------------- +from __future__ import annotations + import os +import warnings as _warnings + import dexsim import open3d as o3d -from dataclasses import MISSING -from typing import List, Union +from typing import TYPE_CHECKING, List, Union from dexsim.types import ( + CloneStrategy, DriveType, ArticulationFlag, LoadOption, + ObjectCloneOptions, RigidBodyShape, SDFConfig, - PhysicalAttr, ) from dexsim.engine import Articulation from dexsim.environment import Env, Arena @@ -35,17 +39,30 @@ from embodichain.lab.sim.cfg import ( ArticulationCfg, + ArticulationRootPropertiesCfg, LinkPhysicsOverrideCfg, + RigidBodyPhysicsCfg, RigidObjectCfg, SoftObjectCfg, ClothObjectCfg, ) from embodichain.utils.string import resolve_matching_names -from embodichain.lab.sim.shapes import MeshCfg, CubeCfg, SphereCfg +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, MeshCollisionCfg, SphereCfg from embodichain.utils import logger from dexsim.kit.meshproc import get_mesh_auto_uv import numpy as np +if TYPE_CHECKING: + from dexsim.spawn import SpawnedArticulation + + +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 get_dexsim_arenas() -> List[dexsim.environment.Arena]: """Get all arenas in the default dexsim world. @@ -77,33 +94,9 @@ def get_dexsim_arena_num() -> int: def _resolve_mesh_collision_params( cfg: RigidObjectCfg, -) -> tuple[int, str, int]: - """Resolve legacy and shape-level mesh collision parameters.""" - - def is_missing(value) -> bool: - # deepcopy() can produce a distinct instance of dataclasses.MISSING. - return value is MISSING or isinstance(value, type(MISSING)) - - max_convex_hull_num = next( - value - for value in ( - cfg.max_convex_hull_num, - cfg.shape.max_convex_hull_num, - 1, - ) - if not is_missing(value) - ) - acd_method = next( - value - for value in (cfg.acd_method, cfg.shape.acd_method, "coacd") - if not is_missing(value) - ) - sdf_resolution = next( - value - for value in (cfg.sdf_resolution, cfg.shape.sdf_resolution, 0) - if not is_missing(value) - ) - return max_convex_hull_num, acd_method, sdf_resolution +) -> MeshCollisionCfg: + """Resolve mesh collision parameters from the shape configuration.""" + return cfg.shape.collision or MeshCollisionCfg() def get_dexsim_drive_type(drive_type: str) -> DriveType: @@ -160,67 +153,276 @@ def _apply_link_physics_overrides( group_cfg = link_to_group.get(name) if group_cfg is None: continue - physical_attr = group_cfg.attrs.merge_with(cfg.attrs) - replace_inertial = group_cfg.replace_inertial or ( - group_cfg.attrs.mass is not None + base_attr = cfg.attrs.to_dexsim_physical_attr() + physical_attr = group_cfg.attrs.to_dexsim_physical_attr(base=base_attr) + mass_props = group_cfg.attrs.mass_props + recompute_inertia = bool( + mass_props is not None and mass_props.recompute_inertia + ) + art.set_physical_attr( + physical_attr, + name, + is_replace_inertial=recompute_inertia, ) - 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 _warn_legacy_articulation_api(name: str) -> None: + _warnings.warn( + f"{name}() bypasses the Spawn ownership/configuration path and is " + "deprecated; declare the articulation through SimulationManager instead.", + DeprecationWarning, + stacklevel=3, + ) - Args: - arts (List[Articulation]): List of dexsim articulations to configure. - cfg (ArticulationCfg): Configuration object containing articulation settings. + +def _default_articulation_clone_options() -> ObjectCloneOptions: + options = ObjectCloneOptions() + options.render.material = CloneStrategy.DEEP_COPY + return options + + +def default_articulation_clone_options() -> ObjectCloneOptions: + """Return legacy articulation clone options. + + Deprecated: new scene code must use the Spawn declaration path. """ + _warn_legacy_articulation_api("default_articulation_clone_options") + return _default_articulation_clone_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 get_drive_type(drive_pros): - if isinstance(drive_pros, dict): - return drive_pros.get("drive_type", None) - return getattr(drive_pros, "drive_type", None) - drive_pros = getattr(cfg, "drive_pros", None) - drive_type = get_drive_type(drive_pros) if drive_pros is not None else None +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 + ) - if drive_type == "force": - drive_type = DriveType.FORCE - elif drive_type == "acceleration": - drive_type = DriveType.ACCELERATION - elif drive_type == "none": - drive_type = DriveType.NONE - else: - logger.log_error(f"Unknow drive type {drive_type}") - for i, art in enumerate(arts): - art.set_body_scale(cfg.body_scale) - 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 +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. + + Deprecated: use ``SimulationManager.add_articulation()`` or + ``SimulationManager.add_robot()``. + """ + _warn_legacy_articulation_api("spawn_articulation_entities") + 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 cfg.resolve_asset_physics_mode() == "overlay": + _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, ) - art.set_solver_iteration_counts( - min_position_iters=cfg.min_position_iters, - min_velocity_iters=cfg.min_velocity_iters, + 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. + + Deprecated: use ``SimulationManager.add_articulation()`` or + ``SimulationManager.add_robot()``. + """ + _warn_legacy_articulation_api("spawn_usd_articulation_entities") + 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 cfg.resolve_asset_physics_mode() == "overlay": + _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 - # 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) - if i == 0 and cfg.compute_uv: - render_body = art.get_render_body(name) - if render_body: - render_body.set_projective_uv() +def set_dexsim_articulation_cfg( + art: Articulation | SpawnedArticulation, + cfg: ArticulationCfg, +) -> None: + """Apply cfg through the deprecated raw DexSim articulation path. - # TODO: will crash when exit if not explicitly delete. - # This may due to the destruction of render body order when exiting. - del render_body + Args: + art: DexSim articulation (or Newton skeleton carrier) to configure. + cfg: EmbodiChain articulation configuration. + """ + _warn_legacy_articulation_api("set_dexsim_articulation_cfg") + _set_dexsim_articulation_cfg(art, cfg) + + +def _set_dexsim_articulation_cfg( + art: Articulation | SpawnedArticulation, + cfg: ArticulationCfg, +) -> None: + """Implement the retained legacy path for compatibility wrappers.""" + + is_newton_art = hasattr(art, "dexsim_meta_links") + if is_newton_art: + raise TypeError( + "The deprecated raw articulation configuration path is " + "Default-backend-only. Declare the asset through SimulationManager " + "and use grouped RigidBodyPhysicsCfg properties for Newton." + ) + lifecycle_state = getattr(getattr(art, "_mgr", None), "_lifecycle_state", None) + lifecycle_name = getattr(lifecycle_state, "name", "") + if lifecycle_name == "BUILDER" or not is_newton_art: + art.set_body_scale(cfg.body_scale) + + link_names = art.get_link_names() + physical_attr = cfg.attrs.to_dexsim_physical_attr() + art.set_physical_attr(physical_attr) + _apply_link_physics_overrides(art, cfg, link_names) + root_props = cfg.root_props + fixed_base = True if root_props.fixed_base is None else bool(root_props.fixed_base) + self_collision_enabled = ( + False + if root_props.self_collision_enabled is None + else bool(root_props.self_collision_enabled) + ) + art.set_articulation_flag(ArticulationFlag.FIX_BASE, fixed_base) + art.set_articulation_flag( + ArticulationFlag.DISABLE_SELF_COLLISION, not self_collision_enabled + ) + _apply_default_articulation_root_properties(art, root_props) + + 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 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 + + +def _apply_default_articulation_root_properties( + art: Articulation, + props: ArticulationRootPropertiesCfg, +) -> None: + """Apply explicitly configured Default-native articulation-root values.""" + if props.sleep_threshold is not None: + art.set_sleep_threshold(float(props.sleep_threshold)) + + position_iters = props.min_position_iters + velocity_iters = props.min_velocity_iters + if (position_iters is None) != (velocity_iters is None): + raise ValueError( + "Articulation-root min_position_iters and min_velocity_iters " + "must be configured together." + ) + if position_iters is not None: + assert velocity_iters is not None + art.set_solver_iteration_counts( + min_position_iters=int(position_iters), + min_velocity_iters=int(velocity_iters), + ) def is_rt_enabled() -> bool: @@ -284,124 +486,274 @@ 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. +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 - 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. - """ - obj_list = [] - body_type = cfg.to_dexsim_body_type() - if isinstance(cfg.shape, MeshCfg): +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 + + 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) - 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 - cfg: RigidObjectCfg - max_convex_hull_num, acd_method, sdf_resolution = ( - _resolve_mesh_collision_params(cfg) +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: + raise TypeError( + "The deprecated raw rigid-object initialization path is " + "Default-backend-only. Use SimulationManager with grouped " + "RigidBodyPhysicsCfg properties for Newton." ) - fpath = cfg.shape.fpath + obj.set_body_scale(*cfg.body_scale) + obj.add_rigidbody( + body_type, + shape_type, + cfg.attrs.to_dexsim_physical_attr(), + ) - 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 - else: - # non-usd file does not support this option, will be forced set False to avoid potential issues. - cfg.use_usd_properties = False - - for i, env in enumerate(env_list): - if max_convex_hull_num > 1: - obj = env.load_actor_with_acd( - fpath, - duplicate=True, - attach_scene=True, - option=option, - cache_path=cache_dir, - actor_type=body_type, - max_convex_hull_num=max_convex_hull_num, - method=acd_method, - ) - elif sdf_resolution > 0: - obj = env.load_actor( - fpath, duplicate=True, attach_scene=True, option=option - ) - sdf_cfg = SDFConfig() - sdf_cfg.resolution = 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) +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.""" + if is_newton_backend: + raise TypeError( + "The deprecated raw rigid-object initialization path is " + "Default-backend-only. Use SimulationManager with grouped " + "RigidBodyPhysicsCfg properties for Newton." + ) + option = _mesh_load_option_from_cfg(cfg) + fpath = cfg.shape.fpath + collision_cfg = _resolve_mesh_collision_params(cfg) + + if collision_cfg.approximation == "convex_decomposition": + obj = env.load_actor_with_acd( + fpath, + duplicate=True, + attach_scene=True, + option=option, + cache_path=cache_dir, + actor_type=body_type, + max_convex_hull_num=collision_cfg.max_hulls, + method=collision_cfg.acd_method or "coacd", + ) + elif collision_cfg.approximation == "sdf": + if collision_cfg.sdf_resolution is None: + raise ValueError( + "The deprecated raw Default path requires sdf_resolution for " + "MeshCollisionCfg(approximation='sdf')." + ) + if 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=collision_cfg.sdf_resolution) + obj.add_physical_body( + body_type, + RigidBodyShape.SDF, + config=sdf_cfg, + attr=cfg.attrs.to_dexsim_physical_attr(), + ) + else: + if collision_cfg.approximation == "triangle_mesh" and cfg.body_type != "static": + raise ValueError( + "triangle_mesh collision is supported only for static rigid objects." + ) + obj = env.load_actor(fpath, duplicate=True, attach_scene=True, option=option) + shape_type = ( + RigidBodyShape.MESH + if collision_cfg.approximation == "triangle_mesh" + else RigidBodyShape.CONVEX + ) + obj.add_rigidbody( + body_type, + shape_type, + cfg.attrs.to_dexsim_physical_attr(), + ) - if compute_uv: - vertices = obj.get_vertices() - triangles = obj.get_triangles() + _apply_mesh_uv_mapping(obj, cfg) + return obj - 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) - elif isinstance(cfg.shape, CubeCfg): - from embodichain.lab.sim.utility.sim_utils import create_cube +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." + ) - obj_list = create_cube(env_list, cfg.shape.size, uid=cfg.uid) - for obj in obj_list: - obj.add_rigidbody(body_type, RigidBodyShape.BOX) + 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 - 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 +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 clone_options is None: + clone_options = default_rigid_object_clone_options() + + body_type = cfg.to_dexsim_body_type() + is_newton_backend = _is_newton_backend_active() + if is_newton_backend: + raise TypeError( + "spawn_rigid_object_entities() is a deprecated " + "Default-backend-only initialization path. Use " + "SimulationManager.add_rigid_object() with grouped " + "RigidBodyPhysicsCfg properties for Newton." + ) + 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.asset_physics_mode = "overlay" + 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): + 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): + 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/embodichain/lab/visualization/protocol.py b/embodichain/lab/visualization/protocol.py index 3fd30b5e8..27b6238cb 100644 --- a/embodichain/lab/visualization/protocol.py +++ b/embodichain/lab/visualization/protocol.py @@ -129,8 +129,9 @@ def pose_to_position_wxyz(pose: object) -> tuple[np.ndarray, np.ndarray]: """Split pose arrays into positions and normalized wxyz quaternions. The accepted layouts are ``(..., 7)`` in EmbodiChain's - ``(x, y, z, qw, qx, qy, qz)`` convention or homogeneous ``(..., 4, 4)`` - matrices. This is the single conversion boundary used by scene exporters. + ``(x, y, z, qx, qy, qz, qw)`` convention or homogeneous ``(..., 4, 4)`` + matrices. Viser uses ``wxyz``, so this is the single conversion boundary + used by scene exporters. Args: pose: Pose or batch of poses. @@ -144,11 +145,12 @@ def pose_to_position_wxyz(pose: object) -> tuple[np.ndarray, np.ndarray]: array = _array(pose, np.float32) if array.ndim >= 1 and array.shape[-1] == 7: position = array[..., :3].copy() - wxyz = array[..., 3:7].copy() - norms = np.linalg.norm(wxyz, axis=-1, keepdims=True) + xyzw = array[..., 3:7].copy() + norms = np.linalg.norm(xyzw, axis=-1, keepdims=True) if np.any(norms <= np.finfo(np.float32).eps): raise ValueError("Pose contains a degenerate quaternion.") - return position, wxyz / norms + xyzw = xyzw / norms + return position, np.roll(xyzw, 1, axis=-1) if array.ndim >= 2 and array.shape[-2:] == (4, 4): position = array[..., :3, 3].copy() @@ -163,6 +165,22 @@ def pose_to_position_wxyz(pose: object) -> tuple[np.ndarray, np.ndarray]: ) +def _normalize_position_wxyz( + position: object, quaternion: object +) -> tuple[np.ndarray, np.ndarray]: + """Validate one protocol-native position and Viser ``wxyz`` quaternion.""" + position_array = _array(position, np.float32).copy() + quaternion_array = _array(quaternion, np.float32).copy() + if position_array.shape != (3,): + raise ValueError(f"position must have shape (3,), got {position_array.shape}.") + if quaternion_array.shape != (4,): + raise ValueError(f"wxyz must have shape (4,), got {quaternion_array.shape}.") + norm = np.linalg.norm(quaternion_array) + if norm <= np.finfo(np.float32).eps: + raise ValueError("Pose contains a degenerate quaternion.") + return position_array, quaternion_array / norm + + @dataclass(frozen=True) class MeshGeometry: """Backend-neutral triangle mesh stored in local coordinates.""" @@ -310,11 +328,7 @@ class GizmoState: visible: bool = True def __post_init__(self) -> None: - position, wxyz = pose_to_position_wxyz( - np.concatenate( - (_array(self.position, np.float32), _array(self.wxyz, np.float32)) - ) - ) + position, wxyz = _normalize_position_wxyz(self.position, self.wxyz) object.__setattr__(self, "position", position) object.__setattr__(self, "wxyz", wxyz) @@ -346,11 +360,7 @@ def __post_init__(self) -> None: raise ValueError("Gizmo command phase must be 'start', 'update', or 'end'.") if not self.client_id: raise ValueError("Gizmo command client_id must not be empty.") - position, wxyz = pose_to_position_wxyz( - np.concatenate( - (_array(self.position, np.float32), _array(self.wxyz, np.float32)) - ) - ) + position, wxyz = _normalize_position_wxyz(self.position, self.wxyz) object.__setattr__(self, "position", position) object.__setattr__(self, "wxyz", wxyz) @@ -534,11 +544,7 @@ class FrameOverlay: visible: bool = True def __post_init__(self) -> None: - position, wxyz = pose_to_position_wxyz( - np.concatenate( - (_array(self.position, np.float32), _array(self.wxyz, np.float32)) - ) - ) + position, wxyz = _normalize_position_wxyz(self.position, self.wxyz) object.__setattr__(self, "position", position) object.__setattr__(self, "wxyz", wxyz) @@ -554,11 +560,7 @@ class TargetOverlay: visible: bool = True def __post_init__(self) -> None: - position, wxyz = pose_to_position_wxyz( - np.concatenate( - (_array(self.position, np.float32), _array(self.wxyz, np.float32)) - ) - ) + position, wxyz = _normalize_position_wxyz(self.position, self.wxyz) object.__setattr__(self, "position", position) object.__setattr__(self, "wxyz", wxyz) diff --git a/embodichain/lab/visualization/scene_exporter.py b/embodichain/lab/visualization/scene_exporter.py index 31d12c027..61fac218d 100644 --- a/embodichain/lab/visualization/scene_exporter.py +++ b/embodichain/lab/visualization/scene_exporter.py @@ -367,18 +367,8 @@ def build_manifest(self) -> SceneManifest: self._append_deformable_objects( sources=sources, geometries=geometries, - uids=self._sim.get_soft_object_uid_list(), - getter=self._sim.get_soft_object, - kind="soft_object", - asset_prefix="soft", - ) - self._append_deformable_objects( - sources=sources, - geometries=geometries, - uids=self._sim.get_cloth_object_uid_list(), - getter=self._sim.get_cloth_object, - kind="cloth_object", - asset_prefix="cloth", + uids=self._sim.get_deformable_object_uid_list(), + getter=self._sim.get_deformable_object, ) self._append_cameras(camera_sources) self._append_gizmos(gizmo_sources) @@ -548,31 +538,29 @@ def _append_deformable_objects( geometries: dict[str, MeshGeometry], uids: list[str], getter: object, - kind: str, - asset_prefix: str, ) -> None: for uid in uids: asset = getter(uid) if asset is None: continue - if kind == "soft_object": - current_vertices = _to_numpy( - asset.get_current_collision_vertices(), - np.float32, - ) + if asset.deformable_type == "volume": + kind = "soft_object" + asset_prefix = "soft" + elif asset.deformable_type == "surface": + kind = "cloth_object" + asset_prefix = "cloth" else: - current_vertices = _to_numpy( - asset.get_current_vertex_position(), - np.float32, + raise ValueError( + f"Unsupported deformable_type {asset.deformable_type!r} " + f"for asset {uid!r}." ) + current_vertices = _to_numpy( + asset.get_surface_vertices(), + np.float32, + ) uid_component = safe_path_component(uid) selected_env_ids = list(self._env_ids) - if kind == "soft_object": - faces_by_env = asset.get_collision_surface_triangles( - env_ids=selected_env_ids, - ) - else: - faces_by_env = asset.get_triangles(env_ids=selected_env_ids) + faces_by_env = asset.get_surface_triangles(env_ids=selected_env_ids) for selected_index, env_id in enumerate(self._env_ids): vertices = current_vertices[env_id] - self._env_offsets[env_id] faces = faces_by_env[selected_index] @@ -849,10 +837,7 @@ def capture( if not source.node.dynamic_geometry: continue if source.asset_key not in dynamic_vertex_cache: - if source.asset_key[0] == "soft": - vertices = source.asset.get_current_collision_vertices() - else: - vertices = source.asset.get_current_vertex_position() + vertices = source.asset.get_surface_vertices() dynamic_vertex_cache[source.asset_key] = _to_numpy( vertices, np.float32, diff --git a/embodichain/learning/rl/train.py b/embodichain/learning/rl/train.py index 41c44a3f1..f22a4cd80 100644 --- a/embodichain/learning/rl/train.py +++ b/embodichain/learning/rl/train.py @@ -457,11 +457,11 @@ def train_from_config( gpu_index = device.index if gpu_index is None: gpu_index = torch.cuda.current_device() - gym_env_cfg.sim_cfg.sim_device = torch.device(f"cuda:{gpu_index}") + gym_env_cfg.sim_cfg.device = torch.device(f"cuda:{gpu_index}") if hasattr(gym_env_cfg.sim_cfg, "gpu_id"): gym_env_cfg.sim_cfg.gpu_id = gpu_index else: - gym_env_cfg.sim_cfg.sim_device = torch.device("cpu") + gym_env_cfg.sim_cfg.device = torch.device("cpu") gym_env_cfg.sim_cfg.headless = headless gym_env_cfg.sim_cfg.render_cfg = RenderCfg(renderer=renderer) gym_env_cfg.sim_cfg.gpu_id = gpu_id @@ -476,7 +476,7 @@ def train_from_config( ) if rank == 0: logger.log_info( - f"Loaded gym_config from {gym_config_path} (env_id={gym_config_data['id']}, num_envs={gym_env_cfg.num_envs}, headless={gym_env_cfg.sim_cfg.headless}, renderer={gym_env_cfg.sim_cfg.render_cfg.renderer}, sim_device={gym_env_cfg.sim_cfg.sim_device})" + f"Loaded gym_config from {gym_config_path} (env_id={gym_config_data['id']}, num_envs={gym_env_cfg.num_envs}, headless={gym_env_cfg.sim_cfg.headless}, renderer={gym_env_cfg.sim_cfg.render_cfg.renderer}, device={gym_env_cfg.sim_cfg.device})" ) env = build_env(gym_config_data["id"], base_env_cfg=gym_env_cfg) diff --git a/embodichain/utils/__init__.py b/embodichain/utils/__init__.py index fd680446c..3ba103010 100644 --- a/embodichain/utils/__init__.py +++ b/embodichain/utils/__init__.py @@ -19,6 +19,8 @@ The ``@configclass`` decorator, ``CfgNode`` configuration system, logging, math/tensor helpers, file/string/device/image utilities, non-maximum suppression, and high-performance ``warp`` kernels. """ +from __future__ import annotations + from .configclass import configclass, is_configclass from .config_paths import resolve_config_path diff --git a/embodichain/utils/configclass.py b/embodichain/utils/configclass.py index a2d0a5542..5813cfd49 100644 --- a/embodichain/utils/configclass.py +++ b/embodichain/utils/configclass.py @@ -154,6 +154,17 @@ def _combined(*args, **kwargs): return _combined +def _is_class_var_annotation(annotation: Any) -> bool: + """Return whether an eager or postponed annotation denotes ``ClassVar``.""" + if annotation is ClassVar or getattr(annotation, "__origin__", None) is ClassVar: + return True + if not isinstance(annotation, str): + return False + return annotation in {"ClassVar", "typing.ClassVar"} or annotation.startswith( + ("ClassVar[", "typing.ClassVar[") + ) + + def custom_post_init(obj): """Deepcopy all elements to avoid shared memory issues for mutable objects in dataclasses initialization. @@ -161,10 +172,13 @@ def custom_post_init(obj): proxy type i.e. a read only proxy for mapping objects. The error is thrown when using hierarchical data-classes for configuration. """ + annotations = obj.__class__.__dict__.get("__annotations__", {}) for key in dir(obj): # skip dunder members if key.startswith("__"): continue + if _is_class_var_annotation(annotations.get(key)): + continue # get data member value = getattr(obj, key) # check annotation @@ -538,8 +552,7 @@ class State: value = class_members.get(key, MISSING) # check if key belongs to ClassVar # in that case, we cannot use default_factory! - origin = getattr(ann[key], "__origin__", None) - if origin is ClassVar: + if _is_class_var_annotation(ann[key]): continue # check if f is MISSING # note: commented out for now since it causes issue with inheritance diff --git a/embodichain/utils/math.py b/embodichain/utils/math.py index 1e5842d6a..5b5fa6278 100644 --- a/embodichain/utils/math.py +++ b/embodichain/utils/math.py @@ -237,12 +237,12 @@ def quat_unique(q: torch.Tensor) -> torch.Tensor: rotation. This function ensures the real part of the quaternion is non-negative. Args: - q: The quaternion orientation in (w, x, y, z). Shape is (..., 4). + q: The quaternion orientation in (x, y, z, w). Shape is (..., 4). Returns: Standardized quaternions. Shape is (..., 4). """ - return torch.where(q[..., 0:1] < 0, -q, q) + return torch.where(q[..., 3:4] < 0, -q, q) @torch.jit.script @@ -250,7 +250,7 @@ def matrix_from_quat(quaternions: torch.Tensor) -> torch.Tensor: """Convert rotations given as quaternions to rotation matrices. Args: - quaternions: The quaternion orientation in (w, x, y, z). Shape is (..., 4). + quaternions: The quaternion orientation in (x, y, z, w). Shape is (..., 4). Returns: Rotation matrices. The shape is (..., 3, 3). @@ -258,7 +258,7 @@ def matrix_from_quat(quaternions: torch.Tensor) -> torch.Tensor: Reference: https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py#L41-L70 """ - r, i, j, k = torch.unbind(quaternions, -1) + i, j, k, r = torch.unbind(quaternions, -1) # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`. two_s = 2.0 / (quaternions * quaternions).sum(-1) @@ -282,14 +282,15 @@ def matrix_from_quat(quaternions: torch.Tensor) -> torch.Tensor: def convert_quat( quat: torch.Tensor | np.ndarray, to: Literal["xyzw", "wxyz"] = "xyzw" ) -> torch.Tensor | np.ndarray: - """Converts quaternion from one convention to another. + """Convert a quaternion between ``wxyz`` and ``xyzw`` conventions. The convention to convert TO is specified as an optional argument. If to == 'xyzw', then the input is in 'wxyz' format, and vice-versa. Args: quat: The quaternion of shape (..., 4). - to: Convention to convert quaternion to.. Defaults to "xyzw". + to: Convention to convert the quaternion to. The input is interpreted as + the opposite convention. Defaults to ``"xyzw"``. Returns: The converted quaternion in specified convention. @@ -332,14 +333,14 @@ def quat_conjugate(q: torch.Tensor) -> torch.Tensor: """Computes the conjugate of a quaternion. Args: - q: The quaternion orientation in (w, x, y, z). Shape is (..., 4). + q: The quaternion orientation in (x, y, z, w). Shape is (..., 4). Returns: - The conjugate quaternion in (w, x, y, z). Shape is (..., 4). + The conjugate quaternion in (x, y, z, w). Shape is (..., 4). """ shape = q.shape q = q.reshape(-1, 4) - return torch.cat((q[..., 0:1], -q[..., 1:]), dim=-1).view(shape) + return torch.cat((-q[..., :3], q[..., 3:4]), dim=-1).view(shape) @torch.jit.script @@ -347,11 +348,11 @@ def quat_inv(q: torch.Tensor, eps: float = 1e-9) -> torch.Tensor: """Computes the inverse of a quaternion. Args: - q: The quaternion orientation in (w, x, y, z). Shape is (N, 4). + q: The quaternion orientation in (x, y, z, w). Shape is (N, 4). eps: A small value to avoid division by zero. Defaults to 1e-9. Returns: - The inverse quaternion in (w, x, y, z). Shape is (N, 4). + The inverse quaternion in (x, y, z, w). Shape is (N, 4). """ return quat_conjugate(q) / q.pow(2).sum(dim=-1, keepdim=True).clamp(min=eps) @@ -371,7 +372,7 @@ def quat_from_euler_xyz( yaw: Rotation around z-axis (in radians). Shape is (N,). Returns: - The quaternion in (w, x, y, z). Shape is (N, 4). + The quaternion in (x, y, z, w). Shape is (N, 4). """ cy = torch.cos(yaw * 0.5) sy = torch.sin(yaw * 0.5) @@ -385,7 +386,7 @@ def quat_from_euler_xyz( qy = cy * cr * sp + sy * sr * cp qz = sy * cr * cp - cy * sr * sp - return torch.stack([qw, qx, qy, qz], dim=-1) + return torch.stack([qx, qy, qz, qw], dim=-1) @torch.jit.script @@ -407,7 +408,7 @@ def quat_from_matrix(matrix: torch.Tensor) -> torch.Tensor: matrix: The rotation matrices. Shape is (..., 3, 3). Returns: - The quaternion in (w, x, y, z). Shape is (..., 4). + The quaternion in (x, y, z, w). Shape is (..., 4). Reference: https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py#L102-L161 @@ -454,16 +455,17 @@ def quat_from_matrix(matrix: torch.Tensor) -> torch.Tensor: # if not for numerical problems, quat_candidates[i] should be same (up to a sign), # forall i; we pick the best-conditioned one (with the largest denominator) - return quat_candidates[ + quaternion_wxyz = quat_candidates[ torch.nn.functional.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, : ].reshape(batch_dim + (4,)) + return torch.cat([quaternion_wxyz[..., 1:], quaternion_wxyz[..., :1]], dim=-1) def xyz_quat_to_4x4_matrix(xyz_quat: torch.Tensor) -> torch.Tensor: - """Convert a 7D pose vector (x, y, z, qw, qx, qy, qz) to a 4x4 transformation matrix. + """Convert a 7D pose vector (x, y, z, qx, qy, qz, qw) to a 4x4 transformation matrix. Args: - xyz_quat: The pose vector in (x, y, z, qw, qx, qy, qz). Shape is (..., 7). + xyz_quat: The pose vector in (x, y, z, qx, qy, qz, qw). Shape is (..., 7). Returns: The transformation matrix. Shape is (..., 4, 4). @@ -492,7 +494,7 @@ def trans_matrix_to_xyz_quat(matrix: torch.Tensor) -> torch.Tensor: matrix: The pose transformation matrix in ((R, t), (0, 1)). Shape is (..., 4, 4). Returns: - The pose vector in (x, y, z, qw, qx, qy, qz). Shape is (..., 7). + The pose vector in (x, y, z, qx, qy, qz, qw). Shape is (..., 7). """ if matrix.shape[-2:] != (4, 4): raise ValueError(f"Invalid input shape {matrix.shape}, expected (..., 4, 4).") @@ -640,7 +642,7 @@ def euler_xyz_from_quat( The euler angles are assumed in XYZ extrinsic convention. Args: - quat: The quaternion orientation in (w, x, y, z). Shape is (N, 4). + quat: The quaternion orientation in (x, y, z, w). Shape is (N, 4). wrap_to_2pi (bool): Whether to wrap output Euler angles into [0, 2π). If False, angles are returned in the default range (−π, π]. Defaults to False. @@ -651,7 +653,7 @@ def euler_xyz_from_quat( Reference: https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles """ - q_w, q_x, q_y, q_z = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3] + q_x, q_y, q_z, q_w = quat[:, 0], quat[:, 1], quat[:, 2], quat[:, 3] # roll (x-axis rotation) sin_roll = 2.0 * (q_w * q_x + q_y * q_z) cos_roll = 1 - 2 * (q_x * q_x + q_y * q_y) @@ -680,7 +682,7 @@ def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tenso """Convert rotations given as quaternions to axis/angle. Args: - quat: The quaternion orientation in (w, x, y, z). Shape is (..., 4). + quat: The quaternion orientation in (x, y, z, w). Shape is (..., 4). eps: The tolerance for Taylor approximation. Defaults to 1.0e-6. Returns: @@ -690,21 +692,21 @@ def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tenso Reference: https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py#L526-L554 """ - # Modified to take in quat as [q_w, q_x, q_y, q_z] - # Quaternion is [q_w, q_x, q_y, q_z] = [cos(theta/2), n_x * sin(theta/2), n_y * sin(theta/2), n_z * sin(theta/2)] + # Modified to take in quat as [q_x, q_y, q_z, q_w] + # Quaternion is [q_x, q_y, q_z, q_w] = [n_x * sin(theta/2), n_y * sin(theta/2), n_z * sin(theta/2), cos(theta/2)] # Axis-angle is [a_x, a_y, a_z] = [theta * n_x, theta * n_y, theta * n_z] # Thus, axis-angle is [q_x, q_y, q_z] / (sin(theta/2) / theta) # When theta = 0, (sin(theta/2) / theta) is undefined # However, as theta --> 0, we can use the Taylor approximation 1/2 - theta^2 / 48 - quat = quat * (1.0 - 2.0 * (quat[..., 0:1] < 0.0)) - mag = torch.linalg.norm(quat[..., 1:], dim=-1) - half_angle = torch.atan2(mag, quat[..., 0]) + quat = quat * (1.0 - 2.0 * (quat[..., 3:4] < 0.0)) + mag = torch.linalg.norm(quat[..., :3], dim=-1) + half_angle = torch.atan2(mag, quat[..., 3]) angle = 2.0 * half_angle # check whether to apply Taylor approximation sin_half_angles_over_angles = torch.where( angle.abs() > eps, torch.sin(half_angle) / angle, 0.5 - angle * angle / 48 ) - return quat[..., 1:4] / sin_half_angles_over_angles.unsqueeze(-1) + return quat[..., :3] / sin_half_angles_over_angles.unsqueeze(-1) @torch.jit.script @@ -716,12 +718,12 @@ def quat_from_angle_axis(angle: torch.Tensor, axis: torch.Tensor) -> torch.Tenso axis: The axis of rotation. Shape is (N, 3). Returns: - The quaternion in (w, x, y, z). Shape is (N, 4). + The quaternion in (x, y, z, w). Shape is (N, 4). """ theta = (angle / 2).unsqueeze(-1) xyz = normalize(axis) * theta.sin() w = theta.cos() - return normalize(torch.cat([w, xyz], dim=-1)) + return normalize(torch.cat([xyz, w], dim=-1)) @torch.jit.script @@ -729,11 +731,11 @@ def quat_mul(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: """Multiply two quaternions together. Args: - q1: The first quaternion in (w, x, y, z). Shape is (..., 4). - q2: The second quaternion in (w, x, y, z). Shape is (..., 4). + q1: The first quaternion in (x, y, z, w). Shape is (..., 4). + q2: The second quaternion in (x, y, z, w). Shape is (..., 4). Returns: - The product of the two quaternions in (w, x, y, z). Shape is (..., 4). + The product of the two quaternions in (x, y, z, w). Shape is (..., 4). Raises: ValueError: Input shapes of ``q1`` and ``q2`` are not matching. @@ -747,8 +749,8 @@ def quat_mul(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: q1 = q1.reshape(-1, 4) q2 = q2.reshape(-1, 4) # extract components from quaternions - w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3] - w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3] + x1, y1, z1, w1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3] + x2, y2, z2, w2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3] # perform multiplication ww = (z1 + x1) * (x2 + y2) yy = (w1 - y1) * (w2 + z2) @@ -760,7 +762,7 @@ def quat_mul(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: y = qq - yy + (w1 - x1) * (y2 + z2) z = qq - zz + (z1 + y1) * (w2 - x2) - return torch.stack([w, x, y, z], dim=-1).view(shape) + return torch.stack([x, y, z, w], dim=-1).view(shape) @torch.jit.script @@ -768,21 +770,21 @@ def yaw_quat(quat: torch.Tensor) -> torch.Tensor: """Extract the yaw component of a quaternion. Args: - quat: The orientation in (w, x, y, z). Shape is (..., 4) + quat: The orientation in (x, y, z, w). Shape is (..., 4) Returns: A quaternion with only yaw component. """ shape = quat.shape quat_yaw = quat.view(-1, 4) - qw = quat_yaw[:, 0] - qx = quat_yaw[:, 1] - qy = quat_yaw[:, 2] - qz = quat_yaw[:, 3] + qx = quat_yaw[:, 0] + qy = quat_yaw[:, 1] + qz = quat_yaw[:, 2] + qw = quat_yaw[:, 3] yaw = torch.atan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz)) quat_yaw = torch.zeros_like(quat_yaw) - quat_yaw[:, 3] = torch.sin(yaw / 2) - quat_yaw[:, 0] = torch.cos(yaw / 2) + quat_yaw[:, 2] = torch.sin(yaw / 2) + quat_yaw[:, 3] = torch.cos(yaw / 2) quat_yaw = normalize(quat_yaw) return quat_yaw.view(shape) @@ -792,8 +794,8 @@ def quat_box_minus(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: """The box-minus operator (quaternion difference) between two quaternions. Args: - q1: The first quaternion in (w, x, y, z). Shape is (N, 4). - q2: The second quaternion in (w, x, y, z). Shape is (N, 4). + q1: The first quaternion in (x, y, z, w). Shape is (N, 4). + q2: The second quaternion in (x, y, z, w). Shape is (N, 4). Returns: The difference between the two quaternions. Shape is (N, 3). @@ -812,7 +814,7 @@ def quat_box_plus( """The box-plus operator (quaternion update) to apply an increment to a quaternion. Args: - q: The initial quaternion in (w, x, y, z). Shape is (N, 4). + q: The initial quaternion in (x, y, z, w). Shape is (N, 4). delta: The axis-angle perturbation. Shape is (N, 3). eps: A small value to avoid division by zero. Defaults to 1e-6. @@ -837,7 +839,7 @@ def quat_apply(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: """Apply a quaternion rotation to a vector. Args: - quat: The quaternion in (w, x, y, z). Shape is (..., 4). + quat: The quaternion in (x, y, z, w). Shape is (..., 4). vec: The vector in (x, y, z). Shape is (..., 3). Returns: @@ -849,9 +851,9 @@ def quat_apply(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: quat = quat.reshape(-1, 4) vec = vec.reshape(-1, 3) # extract components from quaternions - xyz = quat[:, 1:] + xyz = quat[:, :3] t = xyz.cross(vec, dim=-1) * 2 - return (vec + quat[:, 0:1] * t + xyz.cross(t, dim=-1)).view(shape) + return (vec + quat[:, 3:4] * t + xyz.cross(t, dim=-1)).view(shape) @torch.jit.script @@ -859,7 +861,7 @@ def quat_apply_inverse(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: """Apply an inverse quaternion rotation to a vector. Args: - quat: The quaternion in (w, x, y, z). Shape is (..., 4). + quat: The quaternion in (x, y, z, w). Shape is (..., 4). vec: The vector in (x, y, z). Shape is (..., 3). Returns: @@ -871,9 +873,9 @@ def quat_apply_inverse(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: quat = quat.reshape(-1, 4) vec = vec.reshape(-1, 3) # extract components from quaternions - xyz = quat[:, 1:] + xyz = quat[:, :3] t = xyz.cross(vec, dim=-1) * 2 - return (vec - quat[:, 0:1] * t + xyz.cross(t, dim=-1)).view(shape) + return (vec - quat[:, 3:4] * t + xyz.cross(t, dim=-1)).view(shape) @torch.jit.script @@ -881,7 +883,7 @@ def quat_apply_yaw(quat: torch.Tensor, vec: torch.Tensor) -> torch.Tensor: """Rotate a vector only around the yaw-direction. Args: - quat: The orientation in (w, x, y, z). Shape is (N, 4). + quat: The orientation in (x, y, z, w). Shape is (N, 4). vec: The vector in (x, y, z). Shape is (N, 3). Returns: @@ -896,8 +898,8 @@ def quat_error_magnitude(q1: torch.Tensor, q2: torch.Tensor) -> torch.Tensor: """Computes the rotation difference between two quaternions. Args: - q1: The first quaternion in (w, x, y, z). Shape is (..., 4). - q2: The second quaternion in (w, x, y, z). Shape is (..., 4). + q1: The first quaternion in (x, y, z, w). Shape is (..., 4). + q2: The second quaternion in (x, y, z, w). Shape is (..., 4). Returns: Angular error between input quaternions in radians. @@ -952,7 +954,7 @@ def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool: Args: pos: The cartesian position. Shape is (N, 3). - rot: The quaternion in (w, x, y, z). Shape is (N, 4). + rot: The quaternion in (x, y, z, w). Shape is (N, 4). Returns: True if all the input poses result in identity transform. Otherwise, False. @@ -960,7 +962,7 @@ def is_identity_pose(pos: torch.tensor, rot: torch.tensor) -> bool: # create identity transformations pos_identity = torch.zeros_like(pos) rot_identity = torch.zeros_like(rot) - rot_identity[..., 0] = 1 + rot_identity[..., 3] = 1 # compare input to identity return torch.allclose(pos, pos_identity) and torch.allclose(rot, rot_identity) @@ -979,10 +981,10 @@ def combine_frame_transforms( Args: t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3). - q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (x, y, z, w). Shape is (N, 4). t12: Position of frame 2 w.r.t. frame 1. Shape is (N, 3). Defaults to None, in which case the position is assumed to be zero. - q12: Quaternion orientation of frame 2 w.r.t. frame 1 in (w, x, y, z). Shape is (N, 4). + q12: Quaternion orientation of frame 2 w.r.t. frame 1 in (x, y, z, w). Shape is (N, 4). Defaults to None, in which case the orientation is assumed to be identity. Returns: @@ -1028,7 +1030,7 @@ def rigid_body_twist_transform( v0: Linear velocity of 0 in frame 0. Shape is (N, 3). w0: Angular velocity of 0 in frame 0. Shape is (N, 3). t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3). - q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (x, y, z, w). Shape is (N, 4). Returns: A tuple containing: @@ -1054,10 +1056,10 @@ def subtract_frame_transforms( Args: t01: Position of frame 1 w.r.t. frame 0. Shape is (N, 3). - q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + q01: Quaternion orientation of frame 1 w.r.t. frame 0 in (x, y, z, w). Shape is (N, 4). t02: Position of frame 2 w.r.t. frame 0. Shape is (N, 3). Defaults to None, in which case the position is assumed to be zero. - q02: Quaternion orientation of frame 2 w.r.t. frame 0 in (w, x, y, z). Shape is (N, 4). + q02: Quaternion orientation of frame 2 w.r.t. frame 0 in (x, y, z, w). Shape is (N, 4). Defaults to None, in which case the orientation is assumed to be identity. Returns: @@ -1090,9 +1092,9 @@ def compute_pose_error( Args: t01: Position of source frame. Shape is (N, 3). - q01: Quaternion orientation of source frame in (w, x, y, z). Shape is (N, 4). + q01: Quaternion orientation of source frame in (x, y, z, w). Shape is (N, 4). t02: Position of target frame. Shape is (N, 3). - q02: Quaternion orientation of target frame in (w, x, y, z). Shape is (N, 4). + q02: Quaternion orientation of target frame in (x, y, z, w). Shape is (N, 4). rot_error_type: The rotation error type to return: "quat", "axis_angle". Defaults to "axis_angle". @@ -1111,7 +1113,7 @@ def compute_pose_error( # Compute quaternion error (i.e., difference quaternion) # Reference: https://personal.utdallas.edu/~sxb027100/dock/quaternion.html # q_current_norm = q_current * q_current_conj - source_quat_norm = quat_mul(q01, quat_conjugate(q01))[:, 0] + source_quat_norm = quat_mul(q01, quat_conjugate(q01))[:, 3] # q_current_inv = q_current_conj / q_current_norm source_quat_inv = quat_conjugate(q01) / source_quat_norm.unsqueeze(-1) # q_error = q_target * q_current_inv @@ -1148,7 +1150,7 @@ def apply_delta_pose( Args: source_pos: Position of source frame. Shape is (N, 3). - source_rot: Quaternion orientation of source frame in (w, x, y, z). Shape is (N, 4).. + source_rot: Quaternion orientation of source frame in (x, y, z, w). Shape is (N, 4).. delta_pose: Position and orientation displacements. Shape is (N, 6). eps: The tolerance to consider orientation displacement as zero. Defaults to 1.0e-6. @@ -1167,7 +1169,7 @@ def apply_delta_pose( angle = torch.linalg.vector_norm(rot_actions, dim=1) axis = rot_actions / angle.unsqueeze(-1) # change from axis-angle to quat convention - identity_quat = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device).repeat( + identity_quat = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device).repeat( num_poses, 1 ) rot_delta_quat = torch.where( @@ -1205,7 +1207,7 @@ def transform_points( points: Points to transform. Shape is (N, P, 3) or (P, 3). pos: Position of the target frame. Shape is (N, 3) or (3,). Defaults to None, in which case the position is assumed to be zero. - quat: Quaternion orientation of the target frame in (w, x, y, z). Shape is (N, 4) or (4,). + quat: Quaternion orientation of the target frame in (x, y, z, w). Shape is (N, 4) or (4,). Defaults to None, in which case the orientation is assumed to be identity. Returns: @@ -1567,10 +1569,10 @@ def default_orientation(num: int, device: str) -> torch.Tensor: device: Device to create tensor on. Returns: - Identity quaternion in (w, x, y, z). Shape is (num, 4). + Identity quaternion in (x, y, z, w). Shape is (num, 4). """ quat = torch.zeros((num, 4), dtype=torch.float32, device=device) - quat[..., 0] = 1.0 + quat[..., 3] = 1.0 return quat @@ -1584,7 +1586,7 @@ def random_orientation(num: int, device: str) -> torch.Tensor: device: Device to create tensor on. Returns: - Sampled quaternion in (w, x, y, z). Shape is (num, 4). + Sampled quaternion in (x, y, z, w). Shape is (num, 4). Reference: https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.random.html @@ -1604,7 +1606,7 @@ def random_yaw_orientation(num: int, device: str) -> torch.Tensor: device: Device to create tensor on. Returns: - Sampled quaternion in (w, x, y, z). Shape is (num, 4). + Sampled quaternion in (x, y, z, w). Shape is (num, 4). """ roll = torch.zeros(num, dtype=torch.float32, device=device) pitch = torch.zeros(num, dtype=torch.float32, device=device) @@ -1802,12 +1804,12 @@ def convert_camera_frame_orientation_convention( - :obj:`"world"` - forward axis: +X - up axis +Z - Offset is applied in the World Frame convention Args: - orientation: Quaternion of form `(w, x, y, z)` with shape (..., 4) in source convention. + orientation: Quaternion of form `(x, y, z, w)` with shape (..., 4) in source convention. origin: Convention to convert from. Defaults to "opengl". target: Convention to convert to. Defaults to "ros". Returns: - Quaternion of form `(w, x, y, z)` with shape (..., 4) in target convention + Quaternion of form `(x, y, z, w)` with shape (..., 4) in target convention """ if target == origin: return orientation.clone() @@ -2013,12 +2015,12 @@ def quat_slerp(q1: torch.Tensor, q2: torch.Tensor, tau: float) -> torch.Tensor: This function does not support batch processing. Args: - q1: First quaternion in (w, x, y, z) format. - q2: Second quaternion in (w, x, y, z) format. + q1: First quaternion in (x, y, z, w) format. + q2: Second quaternion in (x, y, z, w) format. tau: Interpolation coefficient between 0 (q1) and 1 (q2). Returns: - Interpolated quaternion in (w, x, y, z) format. + Interpolated quaternion in (x, y, z, w) format. """ assert isinstance(q1, torch.Tensor), "Input must be a torch tensor" assert isinstance(q2, torch.Tensor), "Input must be a torch tensor" diff --git a/embodichain/utils/nms.py b/embodichain/utils/nms.py index ca1047405..5ece184bb 100644 --- a/embodichain/utils/nms.py +++ b/embodichain/utils/nms.py @@ -139,8 +139,7 @@ def _poses_to_components(poses: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso # dtype and autograd relationship. poses_f32 = poses.detach().to(dtype=torch.float32).contiguous() positions = poses_f32[:, :3, 3].contiguous() - quaternions_wxyz = quat_from_matrix(poses_f32[:, :3, :3]) - quaternions = torch.cat([quaternions_wxyz[:, 1:], quaternions_wxyz[:, :1]], dim=-1) + quaternions = quat_from_matrix(poses_f32[:, :3, :3]) quaternions = quaternions / torch.linalg.vector_norm( quaternions, dim=-1, keepdim=True ).clamp_min(torch.finfo(quaternions.dtype).eps) diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.json b/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.json index 085d94a3d..4499aff7a 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.json +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.json @@ -46,7 +46,7 @@ "init_pos": [0.0, 0.0, 0.5], "init_rot": [0.0, 0.0, 0.0], "init_qpos": [-0.2, 0.07], - "drive_pros": { + "joint_drive_props": { "stiffness": { "slider_to_cart": 1e1, "cart_to_pole":1e-2 diff --git a/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.yaml b/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.yaml index e5d50843b..90df9bc82 100644 --- a/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.yaml +++ b/embodichain_tasks/configs/tasks/classic_control/cart_pole/env.yaml @@ -43,7 +43,7 @@ robot: init_qpos: - -0.2 - 0.07 - drive_pros: + joint_drive_props: stiffness: slider_to_cart: 10.0 cart_to_pole: 0.01 diff --git a/embodichain_tasks/configs/tasks/manipulation/hand_over/env.json b/embodichain_tasks/configs/tasks/manipulation/hand_over/env.json index 8e921ee6f..2b66a02c4 100644 --- a/embodichain_tasks/configs/tasks/manipulation/hand_over/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/hand_over/env.json @@ -163,11 +163,11 @@ { "kind": "configured_pose", "final_position": [0.0, -0.2, 0.6], - "final_quaternion_wxyz": [ - 0.7071067812, + "final_quaternion_xyzw": [ 0.7071067812, 0.0, - 0.0 + 0.0, + 0.7071067812 ] } ], @@ -255,7 +255,7 @@ "left_hand": ["left_gripper_finger1_joint_1"], "right_hand": ["right_gripper_finger1_joint_1"] }, - "drive_pros": { + "joint_drive_props": { "stiffness": { "left_joint[0-9]": 10000.0, "right_joint[0-9]": 10000.0, @@ -288,8 +288,10 @@ "(left|right)_gripper_finger[12]_link_1" ], "attrs": { - "dynamic_friction": 2.0, - "static_friction": 2.0 + "material_props": { + "dynamic_friction": 2.0, + "static_friction": 2.0 + } } } }, @@ -361,10 +363,14 @@ "size": [0.8, 1.2, 0.02] }, "attrs": { - "mass": 10.0, - "dynamic_friction": 0.9, - "static_friction": 0.95, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "dynamic_friction": 0.9, + "static_friction": 0.95, + "restitution": 0.01 + } }, "body_type": "static", "init_pos": [0.0, 0.0, 0.49], @@ -377,22 +383,33 @@ "shape": { "shape_type": "Mesh", "fpath": "SodaCan/simple_cola_can.obj", - "compute_uv": false + "compute_uv": false, + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 16 + } }, "attrs": { - "mass": 0.33, - "dynamic_friction": 0.97, - "static_friction": 0.99, - "angular_damping": 1.0, - "linear_damping": 0.5, - "contact_offset": 0.001, - "rest_offset": 0.0, - "restitution": 0.01, - "min_position_iters": 32, - "min_velocity_iters": 8, - "max_depenetration_velocity": 2.0 + "mass_props": { + "mass": 0.33 + }, + "rigid_props": { + "angular_damping": 1.0, + "linear_damping": 0.5, + "min_position_iters": 32, + "min_velocity_iters": 8, + "max_depenetration_velocity": 2.0 + }, + "collision_props": { + "contact_offset": 0.001, + "rest_offset": 0.0 + }, + "material_props": { + "dynamic_friction": 0.97, + "static_friction": 0.99, + "restitution": 0.01 + } }, - "max_convex_hull_num": 16, "init_pos": [0.0, 0.02, 0.62], "init_rot": [90.0, 0.0, 0.0], "body_scale": [0.56, 0.56, 0.56] diff --git a/embodichain_tasks/configs/tasks/manipulation/hand_over/expert/program.yaml b/embodichain_tasks/configs/tasks/manipulation/hand_over/expert/program.yaml index f73bd51d4..0650c448b 100644 --- a/embodichain_tasks/configs/tasks/manipulation/hand_over/expert/program.yaml +++ b/embodichain_tasks/configs/tasks/manipulation/hand_over/expert/program.yaml @@ -10,7 +10,7 @@ targets: kind: cyclic_pose values: - position: [0.0, -0.2, 0.6] - quaternion_wxyz: [0.7071067812, 0.7071067812, 0.0, 0.0] + quaternion_xyzw: [0.7071067812, 0.0, 0.0, 0.7071067812] program: kind: segment diff --git a/embodichain_tasks/configs/tasks/manipulation/open_drawer/env.json b/embodichain_tasks/configs/tasks/manipulation/open_drawer/env.json index 35a2d9dc6..c4a6fd4e8 100644 --- a/embodichain_tasks/configs/tasks/manipulation/open_drawer/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/open_drawer/env.json @@ -137,7 +137,7 @@ "control_parts": { "hand": ["gripper_finger1_joint_1"] }, - "drive_pros": { + "joint_drive_props": { "stiffness": { "gripper_finger1_joint_1": 1000.0 }, @@ -181,13 +181,18 @@ "init_pos": [-1.1, 0.0, 0.0], "init_rot": [0.0, 0.0, 90.0], "init_qpos": [0.0], - "fix_base": true, - "drive_pros": { + "root_props": { + "fixed_base": true + }, + "asset_physics_mode": "overlay", + "joint_drive_props": { "drive_type": "none" }, "attrs": { - "static_friction": 1.0, - "dynamic_friction": 1.0 + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0 + } } } ] diff --git a/embodichain_tasks/configs/tasks/manipulation/push_cube/env.json b/embodichain_tasks/configs/tasks/manipulation/push_cube/env.json index 1399faf0a..7bb2639f2 100644 --- a/embodichain_tasks/configs/tasks/manipulation/push_cube/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/push_cube/env.json @@ -135,7 +135,7 @@ "init_pos": [0.0, 0.0, 0.0], "init_rot": [0.0, 0.0, 0.0], "init_qpos": [0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.04, 0.04], - "drive_pros": { + "joint_drive_props": { "drive_type": "force", "stiffness": 100000.0, "damping": 1000.0, @@ -166,17 +166,25 @@ "body_type": "dynamic", "init_pos": [-0.6, -0.4, 0.05], "attrs": { - "mass": 2.0, - "static_friction": 1.0, - "dynamic_friction": 0.8, - "linear_damping": 2.0, - "angular_damping": 2.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "restitution": 0.1, - "max_depenetration_velocity": 10.0, - "max_linear_velocity": 1.0, - "max_angular_velocity": 1.0 + "mass_props": { + "mass": 2.0 + }, + "rigid_props": { + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_depenetration_velocity": 10.0, + "max_linear_velocity": 1.0, + "max_angular_velocity": 1.0 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 0.8, + "restitution": 0.1 + } } } ], diff --git a/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/env.json b/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/env.json index d171818f6..22f406c50 100644 --- a/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/env.json @@ -120,7 +120,7 @@ "control_parts": { "hand": ["gripper_finger1_joint_1"] }, - "drive_pros": { + "joint_drive_props": { "stiffness": { "gripper_finger1_joint_1": 1000.0 }, @@ -163,14 +163,19 @@ "size": [0.05, 0.05, 0.05] }, "body_type": "dynamic", - "max_convex_hull_num": 16, "init_pos": [-0.42, -0.08, 0.025], "attrs": { - "mass": 0.05, - "dynamic_friction": 0.97, - "static_friction": 0.99, - "linear_damping": 0.2, - "angular_damping": 0.2 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "linear_damping": 0.2, + "angular_damping": 0.2 + }, + "material_props": { + "dynamic_friction": 0.97, + "static_friction": 0.99 + } } } ], diff --git a/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/expert/program.yaml b/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/expert/program.yaml index e26ea24a8..55e076dbf 100644 --- a/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/expert/program.yaml +++ b/embodichain_tasks/configs/tasks/manipulation/repeated_pick_place/expert/program.yaml @@ -8,9 +8,9 @@ targets: kind: cyclic_pose values: - position: [-0.40, 0.48, 0.10] - quaternion_wxyz: [1.0, 0.0, 0.0, 0.0] + quaternion_xyzw: [0.0, 0.0, 0.0, 1.0] - position: [-0.42, -0.08, 0.10] - quaternion_wxyz: [1.0, 0.0, 0.0, 0.0] + quaternion_xyzw: [0.0, 0.0, 0.0, 1.0] program: kind: repeat count: 3 diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_rgb/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_rgb/env.json index ba76af4ca..10675a799 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_rgb/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_rgb/env.json @@ -153,7 +153,7 @@ "extrinsics": { "parent": "right_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } }, { @@ -166,7 +166,7 @@ "extrinsics": { "parent": "left_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } } ], @@ -190,10 +190,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -209,19 +213,26 @@ "size": [0.05, 0.05, 0.05] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.015, 0.86], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_2", @@ -230,19 +241,26 @@ "size": [0.05, 0.05, 0.05] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.015, 0.86], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_3", @@ -251,19 +269,26 @@ "size": [0.05, 0.05, 0.05] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.015, 0.86], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] } ] } diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_size/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_size/env.json index 3f803066d..6c6c6f034 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_size/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/blocks_ranking_size/env.json @@ -140,7 +140,7 @@ "extrinsics": { "parent": "right_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } }, { @@ -153,7 +153,7 @@ "extrinsics": { "parent": "left_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } } ], @@ -177,10 +177,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -196,19 +200,26 @@ "size": [0.063, 0.063, 0.063] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.015, 0.86], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_2", @@ -217,19 +228,26 @@ "size": [0.051, 0.051, 0.051] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.015, 0.86], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_3", @@ -238,20 +256,26 @@ "size": [0.039, 0.039, 0.039] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.015, 0.86], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] } ] } - diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/match_object_container/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/match_object_container/env.json index a127b47f4..30fd9a3c9 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/match_object_container/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/match_object_container/env.json @@ -147,7 +147,7 @@ "extrinsics": { "parent": "right_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } }, { @@ -160,7 +160,7 @@ "extrinsics": { "parent": "left_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } } ], @@ -184,10 +184,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -203,20 +207,27 @@ "size": [0.04, 0.04, 0.04] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.565, -0.075, 0.86], "init_rot": [0, 0, 0], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_sphere_1", @@ -225,74 +236,103 @@ "radius": 0.025 }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.635, -0.075, 0.86], "init_rot": [0, 0, 0], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"container_cube", "shape": { "shape_type": "Mesh", - "fpath": "ContainerMetal/container_metal.obj" + "fpath": "ContainerMetal/container_metal.obj", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "body_type": "dynamic", "attrs" : { - "mass": 0.5, - "static_friction": 1.0, - "dynamic_friction": 1.0, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 2.0, - "linear_damping": 2.0, - "angular_damping": 2.0, - "max_linear_velocity": 5.0, - "max_angular_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.5 + }, + "rigid_props": { + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0 + } }, "init_pos": [0.875, -0.25, 0.86], "init_rot": [0, 0, 0], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 8 + "body_scale":[1, 1, 1] }, { "uid":"container_sphere", "shape": { "shape_type": "Mesh", - "fpath": "ContainerMetal/container_metal.obj" + "fpath": "ContainerMetal/container_metal.obj", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "body_type": "dynamic", "attrs" : { - "mass": 0.5, - "static_friction": 1.0, - "dynamic_friction": 1.0, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 2.0, - "linear_damping": 2.0, - "angular_damping": 2.0, - "max_linear_velocity": 5.0, - "max_angular_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.5 + }, + "rigid_props": { + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0 + } }, "init_pos": [0.875, 0.25, 0.86], "init_rot": [0, 0, 0], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 8 + "body_scale":[1, 1, 1] }, { "uid":"block_cube_2", @@ -301,20 +341,27 @@ "size": [0.04, 0.04, 0.04] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.565, 0.075, 0.86], "init_rot": [0, 0, 0], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_sphere_2", @@ -323,21 +370,27 @@ "radius": 0.025 }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.635, 0.075, 0.86], "init_rot": [0, 0, 0], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] } ] } - diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/place_object_drawer/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/place_object_drawer/env.json index 185e4617d..1b6423c6a 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/place_object_drawer/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/place_object_drawer/env.json @@ -77,7 +77,7 @@ "extrinsics": { "parent": "right_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } }, { @@ -90,7 +90,7 @@ "extrinsics": { "parent": "left_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } } ], @@ -114,10 +114,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -130,27 +134,38 @@ "uid":"object", "shape": { "shape_type": "Mesh", - "fpath": "ToyDuck/toy_duck.glb" + "fpath": "ToyDuck/toy_duck.glb", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "attrs" : { - "mass": 0.01, - "static_friction": 1.0, - "dynamic_friction": 1.0, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 2.0, - "linear_damping": 2.0, - "angular_damping": 2.0, - "max_linear_velocity": 5.0, - "max_angular_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.01 + }, + "rigid_props": { + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0 + } }, "init_pos": [0.725, -0.1, 0.86], "init_rot": [0, 0, 0], - "body_scale":[0.2, 0.2, 0.2], - "max_convex_hull_num": 8 + "body_scale":[0.2, 0.2, 0.2] } ], "articulation": [ @@ -162,4 +177,3 @@ } ] } - diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/env.json index 4b06308d6..e66badf6d 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/env.json @@ -229,10 +229,14 @@ "compute_uv": true }, "attrs": { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -246,40 +250,62 @@ "shape": { "shape_type": "Mesh", "fpath": "PaperCup/paper_cup.ply", - "compute_uv": true + "compute_uv": true, + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "attrs": { - "mass": 0.01, - "contact_offset": 0.003, - "rest_offset": 0.001, - "restitution": 0.01, - "max_depenetration_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.01 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "restitution": 0.01 + } }, "init_pos": [0.75, 0.1, 0.9], - "body_scale": [0.75, 0.75, 1.0], - "max_convex_hull_num": 8 + "body_scale": [0.75, 0.75, 1.0] }, { "uid": "bottle", "shape": { "shape_type": "Mesh", "fpath": "ScannedBottle/kashijia_processed.ply", - "compute_uv": true + "compute_uv": true, + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "attrs": { - "mass": 0.01, - "contact_offset": 0.003, - "rest_offset": 0.001, - "restitution": 0.01, - "max_depenetration_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.01 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "restitution": 0.01 + } }, "init_pos": [0.75, -0.1, 0.932], - "body_scale": [1, 1, 1], - "max_convex_hull_num": 8 + "body_scale": [1, 1, 1] } ], "rigid_object_group": [], diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/expert/program.yaml b/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/expert/program.yaml index a50a07737..c0d1ecb88 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/expert/program.yaml +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/pour_water/expert/program.yaml @@ -8,7 +8,7 @@ targets: kind: cyclic_pose values: - position: [0.75, -0.1, 0.962] - quaternion_wxyz: [1.0, 0.0, 0.0, 0.0] + quaternion_xyzw: [0.0, 0.0, 0.0, 1.0] program: kind: segment name: pour_and_return_bottle diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/scoop_ice/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/scoop_ice/env.json index 675566b67..c7cadfff5 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/scoop_ice/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/scoop_ice/env.json @@ -114,7 +114,7 @@ "extrinsics": { "parent": "right_ee", "pos": [0.09, 0.05, 0.04], - "quat": [0.36497168, -0.11507513, 0.88111957, 0.27781593] + "quat": [-0.11507513, 0.88111957, 0.27781593, 0.36497168] } }, { @@ -127,7 +127,7 @@ "extrinsics": { "parent": "left_ee", "pos": [0.09, -0.05, 0.04], - "quat": [0.27781593, 0.88111957, -0.11507513, 0.36497168] + "quat": [0.88111957, -0.11507513, 0.36497168, 0.27781593] } } ], @@ -141,10 +141,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 1.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.05 + "mass_props": { + "mass": 1.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.05 + } }, "body_type": "kinematic", "init_pos": [0.80, 0, 0.54], @@ -156,34 +160,52 @@ "uid": "scoop", "shape": { "shape_type": "Mesh", - "fpath": "ScoopIceNewEnv/scoop.ply" + "fpath": "ScoopIceNewEnv/scoop.ply", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "attrs" : { - "mass": 0.5, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.5 + }, + "rigid_props": { + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.0 + } }, - "max_convex_hull_num": 8, "init_pos": [0, 10, 10] }, { "uid": "paper_cup", "shape": { "shape_type": "Mesh", - "fpath": "PaperCup/paper_cup.ply" + "fpath": "PaperCup/paper_cup.ply", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 16 + } }, "attrs" : { - "mass": 0.5, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.5 + }, + "rigid_props": { + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.0 + } }, - "max_convex_hull_num": 16, "init_pos": [0, 10, 10] } ], @@ -196,15 +218,23 @@ "rigid_objects": { "obj": { "attrs" : { - "mass": 0.004, - "contact_offset": 0.001, - "rest_offset": 0, - "dynamic_friction": 0.05, - "static_friction": 0.1, - "restitution": 0.00, - "min_position_iters": 32, - "min_velocity_iters": 8, - "max_depenetration_velocity": 1.0 + "mass_props": { + "mass": 0.004 + }, + "rigid_props": { + "min_position_iters": 32, + "min_velocity_iters": 8, + "max_depenetration_velocity": 1.0 + }, + "collision_props": { + "contact_offset": 0.001, + "rest_offset": 0 + }, + "material_props": { + "dynamic_friction": 0.05, + "static_friction": 0.1, + "restitution": 0.0 + } }, "shape": { "shape_type": "Mesh" @@ -222,12 +252,18 @@ "init_pos": [0.635, -0.04, 0.94], "init_rot": [0, 0, -80], "attrs": { - "mass": 1.0, - "dynamic_friction": 0.05, - "static_friction": 0.1, - "max_depenetration_velocity": 1.0 + "mass_props": { + "mass": 1.0 + }, + "rigid_props": { + "max_depenetration_velocity": 1.0 + }, + "material_props": { + "dynamic_friction": 0.05, + "static_friction": 0.1 + } }, - "drive_pros": { + "joint_drive_props": { "stiffness": 1.0, "damping": 0.1, "max_effort": 100.0 diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/stack_blocks_two/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/stack_blocks_two/env.json index c58ed08ca..bb488af7f 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/stack_blocks_two/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/stack_blocks_two/env.json @@ -96,7 +96,7 @@ "extrinsics": { "parent": "right_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } }, { @@ -109,7 +109,7 @@ "extrinsics": { "parent": "left_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } } ], @@ -133,10 +133,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -152,19 +156,26 @@ "size": [0.05, 0.05, 0.05] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.75, -0.1, 0.9], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] }, { "uid":"block_2", @@ -173,19 +184,26 @@ "size": [0.05, 0.05, 0.05] }, "attrs" : { - "mass": 0.05, - "static_friction": 0.5, - "dynamic_friction": 0.5, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 1e1, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.05 + }, + "rigid_props": { + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 0.5, + "dynamic_friction": 0.5, + "restitution": 0.0 + } }, "init_pos": [0.75, 0.1, 0.9], - "body_scale":[1, 1, 1], - "max_convex_hull_num": 1 + "body_scale":[1, 1, 1] } ] } diff --git a/embodichain_tasks/configs/tasks/manipulation/tableware/stack_cups/env.json b/embodichain_tasks/configs/tasks/manipulation/tableware/stack_cups/env.json index 09daa1494..f148cbe8b 100644 --- a/embodichain_tasks/configs/tasks/manipulation/tableware/stack_cups/env.json +++ b/embodichain_tasks/configs/tasks/manipulation/tableware/stack_cups/env.json @@ -95,7 +95,7 @@ "extrinsics": { "parent": "right_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } }, { @@ -108,7 +108,7 @@ "extrinsics": { "parent": "left_link6", "pos": [-0.08, 0.0, 0.04], - "quat": [0.15304635, 0.69034543, -0.69034543, -0.15304635] + "quat": [0.69034543, -0.69034543, -0.15304635, 0.15304635] } } ], @@ -132,10 +132,14 @@ "fpath": "CircleTableSimple/circle_table_simple.ply" }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", @@ -148,55 +152,75 @@ "uid":"cup_1", "shape": { "shape_type": "Mesh", - "fpath": "PaperCup/paper_cup.ply" + "fpath": "PaperCup/paper_cup.ply", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "attrs" : { - "mass": 0.01, - "static_friction": 1.0, - "dynamic_friction": 1.0, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 2.0, - "linear_damping": 2.0, - "angular_damping": 2.0, - "max_linear_velocity": 5.0, - "max_angular_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.01 + }, + "rigid_props": { + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0 + } }, "init_pos": [0.70, -0.1, 0.86], "init_rot": [0, 0, 0], - "body_scale":[0.75, 0.75, 1.0], - "max_convex_hull_num": 8 + "body_scale":[0.75, 0.75, 1.0] }, { "uid":"cup_2", "shape": { "shape_type": "Mesh", - "fpath": "PaperCup/paper_cup.ply" + "fpath": "PaperCup/paper_cup.ply", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 8 + } }, "attrs" : { - "mass": 0.01, - "static_friction": 1.0, - "dynamic_friction": 1.0, - "restitution": 0.0, - "contact_offset": 0.003, - "rest_offset": 0.001, - "max_depenetration_velocity": 2.0, - "linear_damping": 2.0, - "angular_damping": 2.0, - "max_linear_velocity": 5.0, - "max_angular_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8 + "mass_props": { + "mass": 0.01 + }, + "rigid_props": { + "max_depenetration_velocity": 2.0, + "linear_damping": 2.0, + "angular_damping": 2.0, + "max_linear_velocity": 5.0, + "max_angular_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8 + }, + "collision_props": { + "contact_offset": 0.003, + "rest_offset": 0.001 + }, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0, + "restitution": 0.0 + } }, "init_pos": [0.80, -0.1, 0.86], "init_rot": [0, 0, 0], - "body_scale":[0.75, 0.75, 1.0], - "max_convex_hull_num": 8 + "body_scale":[0.75, 0.75, 1.0] } ] } - - diff --git a/embodichain_tasks/configs/tasks/special/simple_task/env_ur10.json b/embodichain_tasks/configs/tasks/special/simple_task/env_ur10.json index 41faec8ae..c2b71d8e6 100644 --- a/embodichain_tasks/configs/tasks/special/simple_task/env_ur10.json +++ b/embodichain_tasks/configs/tasks/special/simple_task/env_ur10.json @@ -87,10 +87,14 @@ "compute_uv": true }, "attrs" : { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", diff --git a/embodichain_tasks/configs/tasks/special/stay_still_save/env_async_ur10.json b/embodichain_tasks/configs/tasks/special/stay_still_save/env_async_ur10.json index 16329668c..3baa75c94 100644 --- a/embodichain_tasks/configs/tasks/special/stay_still_save/env_async_ur10.json +++ b/embodichain_tasks/configs/tasks/special/stay_still_save/env_async_ur10.json @@ -67,10 +67,14 @@ "compute_uv": true }, "attrs": { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", diff --git a/embodichain_tasks/configs/tasks/special/stay_still_save/env_ur10.json b/embodichain_tasks/configs/tasks/special/stay_still_save/env_ur10.json index b6121fc08..de9ade65f 100644 --- a/embodichain_tasks/configs/tasks/special/stay_still_save/env_ur10.json +++ b/embodichain_tasks/configs/tasks/special/stay_still_save/env_ur10.json @@ -63,10 +63,14 @@ "compute_uv": true }, "attrs": { - "mass": 10.0, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "restitution": 0.01 + "mass_props": { + "mass": 10.0 + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01 + } }, "body_scale": [1, 1, 1], "body_type": "kinematic", diff --git a/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py b/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py new file mode 100644 index 000000000..3b32eb9d4 --- /dev/null +++ b/embodichain_tasks/embodichain_tasks/special/franka_reach_apg.py @@ -0,0 +1,521 @@ +# ---------------------------------------------------------------------------- +# 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 this task explicitly selects +the kinematics route and runs 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 ( + ArticulationRootPropertiesCfg, + 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 + + This task explicitly uses the ``kinematics`` route because the + semi_implicit dynamics solver 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 + FK-only workaround without changing the default route for other + differentiable environments. + """ + + metadata = {"render_modes": ["human"], "default_num_envs": 4} + differentiable_step_mode = "kinematics" + + 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), + root_props=ArticulationRootPropertiesCfg(fixed_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.""" + runtime = self.sim.differentiable_runtime + model = runtime.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) + # Every taped forward replaces these with private primal buffers before + # the bridge opens its tape. They must never alias manager live state. + self._current_joint_q_snapshot: wp.array | None = None + 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). + """ + model = self.sim.differentiable_runtime.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([0.0, 0.0, 0.0, 1.0], device=d).expand( + n, -1 + ) + + # -- DifferentiableEmbodiedEnv contract ------------------------------ # + + def _build_sim_state_dict(self, action: torch.Tensor) -> dict: + """Detach FK primal buffers before the parent opens a Warp tape.""" + runtime = self.sim.differentiable_runtime + self._current_joint_q_snapshot = wp.clone(runtime.current_state.joint_q) + self._fk_state = runtime.model.state() + return super()._build_sim_state_dict(action) + + def _make_kinematic_step_fn(self) -> Callable[[], Any]: + """Explicit FK hook: compute body_q from new_joint_q via ``eval_fk``. + + The semi_implicit solver does not propagate grad through + ``joint_target_pos`` to ``body_q`` (the grad path is zero), so + this kinematics-mode task runs 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.differentiable_runtime.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 + explicit kinematic hook then consumes this array via ``newton.eval_fk``. + """ + n = self.sim.num_envs + total = n * FRANKA_NUM_ARM_JOINTS + if self._current_joint_q_snapshot is None: + raise RuntimeError( + "Franka kinematics requires a detached joint_q snapshot " + "before opening its Warp tape." + ) + # 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, + self._current_joint_q_snapshot, + 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 + both Spawn live states for non-terminal envs so the next step starts + from the new configuration. The tape reads a per-forward detached + snapshot, so this continuation cannot overwrite its primal input. + """ + 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(): + runtime = self.sim.differentiable_runtime + current_q = wp.to_torch(runtime.current_state.joint_q).view( + self.sim.num_envs, -1 + ) + cur = current_q[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) + next_q = torch.clamp(cur + delta, lo, hi) + for state in runtime.live_states: + joint_q = wp.to_torch(state.joint_q).view(self.sim.num_envs, -1) + joint_q[live, :FRANKA_NUM_ARM_JOINTS] = next_q + newton.eval_fk( + runtime.model, + state.joint_q, + state.joint_qd, + state, + ) + 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. No-grad terminal steps use it for + auto-reset; grad-tracked terminal steps expose the IDs in + ``info`` for an explicit reset after backward. + + 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) + runtime = self.sim.differentiable_runtime + for state in runtime.live_states: + joint_q = wp.to_torch(state.joint_q).view(self.sim.num_envs, -1) + joint_q[env_ids] = 0.0 + newton.eval_fk( + runtime.model, + state.joint_q, + state.joint_qd, + state, + ) + obs = self._initial_obs() + return obs, {} + + def _initial_obs(self) -> torch.Tensor: + """Compute the initial observation from the live Spawn state.""" + with torch.no_grad(): + state = self.sim.differentiable_runtime.current_state + 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/examples/sim/demo/grasp_cup_to_caffe.py b/examples/sim/demo/grasp_cup_to_caffe.py index dde7cbaae..bac008db3 100644 --- a/examples/sim/demo/grasp_cup_to_caffe.py +++ b/examples/sim/demo/grasp_cup_to_caffe.py @@ -26,21 +26,22 @@ import torch from tqdm import tqdm from typing import Union -from scipy.spatial.transform import Rotation as R from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, MarkerCfg, JointDrivePropertiesCfg, RigidObjectCfg, - RigidBodyAttributesCfg, + MassPropertiesCfg, + RigidBodyPhysicsCfg, ArticulationCfg, ) from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.data import get_data_path from embodichain.utils import logger from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -58,6 +59,12 @@ def parse_arguments(): description="Create and simulate a robot in SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--seed", + type=int, + default=0, + help="Seed for scene XY perturbations; use a negative value for random runs.", + ) return parser.parse_args() @@ -71,10 +78,22 @@ def initialize_simulation(args) -> SimulationManager: Returns: SimulationManager: Configured simulation manager instance. """ + physics_cfg = physics_cfg_for_backend(args.physics) + if args.physics == "newton": + # This contact-heavy URDF scene needs Newton's collision pipeline; + # MuJoCo's native contact path is not reliable for these convex meshes. + physics_cfg.solver_cfg = { + "solver_type": "mujoco_warp", + "use_mujoco_contacts": False, + "nconmax": 16384, + "njmax": 65536, + } + config = SimulationManagerCfg( headless=True, - sim_device=args.device, + device=args.device, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg, physics_dt=1.0 / 100.0, num_envs=args.num_envs, arena_space=2.5, @@ -178,11 +197,14 @@ def create_table(sim: SimulationManager) -> RigidObject: uid="table", shape=MeshCfg( fpath=get_data_path("MultiW1Data/table_a.obj"), + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=8, + ), ), - attrs=RigidBodyAttributesCfg( - mass=0.5, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=0.5), ), - max_convex_hull_num=8, body_type="kinematic", init_pos=[1.1, -0.5, 0.08], init_rot=[0.0, 0.0, 0.0], @@ -206,11 +228,15 @@ def create_caffe(sim: SimulationManager) -> Robot: fpath=get_data_path("MultiW1Data/cafe/cafe.urdf"), init_pos=[1.05, -0.5, 0.79], init_rot=[0, 0, -30], - attrs=RigidBodyAttributesCfg( - mass=1.0, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), ), - drive_pros=JointDrivePropertiesCfg( - stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + stiffness=1.0, + damping=0.1, + max_effort=100.0, ), ) container = sim.add_articulation(cfg=container_cfg) @@ -232,10 +258,9 @@ def create_cup(sim: SimulationManager) -> RigidObject: shape=MeshCfg( fpath=get_data_path("MultiW1Data/paper_cup_2.obj"), ), - attrs=RigidBodyAttributesCfg( - mass=0.3, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=0.3), ), - max_convex_hull_num=1, body_type="dynamic", init_pos=[0.86, -0.76, 0.841], init_rot=[0.0, 0.0, 0.0], @@ -274,7 +299,11 @@ def create_trajectory( cup_position = cup.get_local_pose(to_matrix=True)[:, :3, 3] # grasp cup waypoint generation - rest_right_qpos = robot.get_qpos()[:, right_arm_ids] # [num_envs, dof] + # Build the task trajectory from the authored hold target. The measured + # pose after the first physics step includes backend-specific gravity and + # constraint settling, which can send the redundant arm IK to a different + # solution before the task even starts. + rest_right_qpos = robot.get_qpos(target=True)[:, right_arm_ids] right_arm_xpos = robot.compute_fk( qpos=rest_right_qpos, name="right_arm", to_matrix=True ) @@ -426,18 +455,20 @@ def main(): caffe = create_caffe(sim) cup = create_cup(sim) - sim.update(step=1) + sim.prepare() - # apply random perturbation + # Apply initialization-time poses before Newton captures its CUDA graph. + # Seed here so backend initialization cannot consume a different random + # prefix and make Default/Newton comparisons use different scenes. + if args.seed >= 0: + np.random.seed(args.seed) apply_random_xy_perturbation(cup, max_perturbation=0.05) apply_random_xy_perturbation(caffe, max_perturbation=0.05) + sim.update(step=1) 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 7184ea6de..c6a3d6033 100644 --- a/examples/sim/demo/pick_up_cloth.py +++ b/examples/sim/demo/pick_up_cloth.py @@ -38,8 +38,9 @@ from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, LightCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, @@ -71,7 +72,7 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]): {"component_type": "hand", "urdf_path": gripper_urdf_path}, ] }, - "drive_pros": { + "joint_drive_props": { "stiffness": {"FINGER[1-2]": 1e2}, "damping": {"FINGER[1-2]": 1e1}, "max_effort": {"FINGER[1-2]": 1e3}, @@ -112,13 +113,16 @@ def create_padding_box(sim: SimulationManager): shape=CubeCfg( size=[0.02, 0.07, 0.05], ), - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.01, - dynamic_friction=0.00, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "rigid_props": {"min_position_iters": 32, "min_velocity_iters": 8}, + "material_props": { + "static_friction": 0.01, + "dynamic_friction": 0.00, + "restitution": 0.01, + }, + } ), body_type="kinematic", init_pos=[0.5, 0.0, 0.026], @@ -257,10 +261,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), visualization=visualization_cfg_from_args(args), ) @@ -270,7 +275,7 @@ def main(): robot = create_robot(sim) cloth = create_cloth(sim) padding_box = create_padding_box(sim) - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() sim.update(step=10) # Let the cloth settle before interaction diff --git a/examples/sim/demo/press_softbody.py b/examples/sim/demo/press_softbody.py index 84d89b580..017235276 100644 --- a/examples/sim/demo/press_softbody.py +++ b/examples/sim/demo/press_softbody.py @@ -37,6 +37,7 @@ from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, SoftObjectCfg, SoftbodyVoxelAttributesCfg, @@ -73,8 +74,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, visualization=visualization_cfg_from_args(args), @@ -188,7 +190,7 @@ def main(): robot = create_robot(sim) soft_cow = create_soft_cow(sim) - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() diff --git a/examples/sim/demo/scoop_ice.py b/examples/sim/demo/scoop_ice.py index 29fde4877..94da16379 100644 --- a/examples/sim/demo/scoop_ice.py +++ b/examples/sim/demo/scoop_ice.py @@ -33,8 +33,9 @@ from embodichain.lab.sim.objects import Robot, RigidObject, RigidObjectGroup from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ArticulationCfg, RigidObjectGroupCfg, JointDrivePropertiesCfg, @@ -42,7 +43,7 @@ ) from embodichain.lab.sim.material import VisualMaterialCfg from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg, CubeCfg +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, MeshCollisionCfg from embodichain.data import get_data_path from embodichain.utils import logger from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -62,6 +63,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, visualization=visualization_cfg_from_args(args), ) @@ -145,7 +147,7 @@ def create_robot(sim): "LEFT_HAND_PINKY", ], }, - "drive_pros": { + "joint_drive_props": { "stiffness": {"LEFT_[A-Z|_]+[0-9]?": 1e2}, "damping": {"LEFT_[A-Z|_]+[0-9]?": 1e1}, "max_effort": {"LEFT_[A-Z|_]+[0-9]?": 1e3}, @@ -182,16 +184,22 @@ def create_scoop(sim: SimulationManager): uid="scoop", shape=MeshCfg( fpath=get_data_path("ScoopIceNewEnv/scoop.ply"), + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=12, + ), ), - attrs=RigidBodyAttributesCfg( - mass=0.5, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 0.5}, + "rigid_props": {"min_position_iters": 32, "min_velocity_iters": 8}, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01, + }, + } ), - max_convex_hull_num=12, body_type="dynamic", init_pos=[0.6, 0.0, 0.09], init_rot=[0.0, 0.0, 0.0], @@ -207,13 +215,16 @@ def create_heave_ice(sim: SimulationManager): shape=MeshCfg( fpath=get_data_path("ScoopIceNewEnv/ice_mesh_small/ice_000.obj"), ), - attrs=RigidBodyAttributesCfg( - mass=0.5, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 0.5}, + "rigid_props": {"min_position_iters": 32, "min_velocity_iters": 8}, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01, + }, + } ), body_type="dynamic", init_pos=[10, 10, 0.08], @@ -229,13 +240,16 @@ def create_padding_box(sim: SimulationManager): shape=CubeCfg( size=[0.1, 0.16, 0.05], ), - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "rigid_props": {"min_position_iters": 32, "min_velocity_iters": 8}, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01, + }, + } ), body_type="kinematic", init_pos=[0.6, 0.15, 0.025], @@ -251,15 +265,18 @@ def create_container(sim: SimulationManager): fpath=get_data_path("ScoopIceNewEnv/IceContainer/ice_container.urdf"), init_pos=[0.7, -0.4, 0.21], init_rot=[0, 0, -90], - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "rigid_props": {"min_position_iters": 32, "min_velocity_iters": 8}, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01, + }, + } ), - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" ), ) @@ -277,15 +294,21 @@ def create_ice_cubes(sim: SimulationManager): "rigid_objects": { "obj": { "attrs": { - "mass": 0.003, - "contact_offset": 0.001, - "rest_offset": 0, - "dynamic_friction": 0.05, - "static_friction": 0.1, - "restitution": 0.01, - "min_position_iters": 32, - "min_velocity_iters": 4, - "max_depenetration_velocity": 1.0, + "mass_props": {"mass": 0.003}, + "rigid_props": { + "min_position_iters": 32, + "min_velocity_iters": 4, + "max_depenetration_velocity": 1.0, + }, + "collision_props": { + "contact_offset": 0.001, + "rest_offset": 0, + }, + "material_props": { + "dynamic_friction": 0.05, + "static_friction": 0.1, + "restitution": 0.01, + }, }, "shape": {"shape_type": "Mesh"}, "init_pos": [20.0, 0, 1.0], @@ -307,6 +330,7 @@ def create_ice_cubes(sim: SimulationManager): material_type="BSDF", ) ) + sim.prepare() ice_cubes.set_visual_material(mat=ice_mat) return ice_cubes diff --git a/examples/sim/gizmo/gizmo_camera.py b/examples/sim/gizmo/gizmo_camera.py index 8ff4e842a..a855ec8a0 100644 --- a/examples/sim/gizmo/gizmo_camera.py +++ b/examples/sim/gizmo/gizmo_camera.py @@ -31,7 +31,12 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.sensors import Camera, CameraCfg -from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + RigidObjectCfg, + RigidBodyPhysicsCfg, + 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 @@ -53,8 +58,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), visualization=visualization_cfg_from_args(args), ) @@ -68,11 +74,15 @@ def main(): uid=f"cube_{i}", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="dynamic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.3, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.3, + }, + } ), init_pos=[0.5 + i * 0.3, 0.0, 0.5], ) @@ -97,6 +107,7 @@ def main(): # Add camera to simulation camera = sim.add_sensor(sensor_cfg=camera_cfg) + sim.prepare() # Wait for initialization time.sleep(0.2) diff --git a/examples/sim/gizmo/gizmo_object.py b/examples/sim/gizmo/gizmo_object.py index 690c7a6de..f61cc713f 100644 --- a/examples/sim/gizmo/gizmo_object.py +++ b/examples/sim/gizmo/gizmo_object.py @@ -25,8 +25,12 @@ import time from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import ( + RigidBodyPhysicsCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg 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 @@ -49,10 +53,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), visualization=visualization_cfg_from_args(args), ) @@ -65,11 +70,15 @@ def main(): uid="cube1", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="kinematic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.1, + }, + } ), init_pos=[0.0, 0.0, 1.0], ) @@ -79,15 +88,20 @@ def main(): uid="cube2", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="kinematic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.1, + }, + } ), init_pos=[0.3, 0.0, 1.0], ) ) + sim.prepare() native_window_opened = False if not args.headless: @@ -123,9 +137,6 @@ def main(): def run_simulation(sim: SimulationManager): """Run the simulation loop.""" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 gizmo_enabled = True try: diff --git a/examples/sim/gizmo/gizmo_robot.py b/examples/sim/gizmo/gizmo_robot.py index cfccfd566..190d07ec5 100644 --- a/examples/sim/gizmo/gizmo_robot.py +++ b/examples/sim/gizmo/gizmo_robot.py @@ -29,6 +29,7 @@ from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, @@ -55,8 +56,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), visualization=visualization_cfg_from_args(args), ) @@ -93,7 +95,7 @@ def main(): num_samples=30, ) }, - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness={"JOINT[0-9]": 1e4, "FINGER[1-2]": 1e2}, damping={"JOINT[0-9]": 1e3, "FINGER[1-2]": 1e1}, max_effort={"JOINT[0-9]": 1e5, "FINGER[1-2]": 1e3}, @@ -102,6 +104,7 @@ def main(): init_qpos=[0.0, -np.pi / 2, -np.pi / 2, np.pi / 2, -np.pi / 2, 0.0, 0.0, 0.0], ) robot = sim.add_robot(cfg=robot_cfg) + sim.prepare() # Set initial joint positions initial_qpos = torch.tensor( diff --git a/examples/sim/gizmo/gizmo_scene.py b/examples/sim/gizmo/gizmo_scene.py index 083405c8e..6cd0d1b47 100644 --- a/examples/sim/gizmo/gizmo_scene.py +++ b/examples/sim/gizmo/gizmo_scene.py @@ -34,11 +34,12 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim.shapes import CubeCfg @@ -63,8 +64,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), visualization=visualization_cfg_from_args(args), ) @@ -101,7 +103,7 @@ def main(): dt=0.1, ), }, - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness={"LEFT_J[1-7]": 1e4, "RIGHT_J[1-7]": 1e4}, damping={"LEFT_J[1-7]": 1e3, "RIGHT_J[1-7]": 1e3}, ), @@ -120,22 +122,20 @@ def main(): device="cpu", ) - left_joint_ids = robot.get_joint_ids("left_arm") - right_joint_ids = robot.get_joint_ids("right_arm") - - robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) - robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) - # Create a rigid object (cube) positioned to the side of the robot cube_cfg = RigidObjectCfg( uid="interactive_cube", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="kinematic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.1, + }, + } ), init_pos=[1.0, 0.0, 0.5], # Position to the side of the robot ) @@ -157,6 +157,12 @@ def main(): ), ) camera = sim.add_sensor(sensor_cfg=camera_cfg) + sim.prepare() + + left_joint_ids = robot.get_joint_ids("left_arm") + right_joint_ids = robot.get_joint_ids("right_arm") + robot.set_qpos(qpos=left_arm_qpos, joint_ids=left_joint_ids) + robot.set_qpos(qpos=right_arm_qpos, joint_ids=right_joint_ids) native_window_opened = False if not args.headless: diff --git a/examples/sim/gizmo/gizmo_w1.py b/examples/sim/gizmo/gizmo_w1.py index 554859c39..b0f4ef55c 100644 --- a/examples/sim/gizmo/gizmo_w1.py +++ b/examples/sim/gizmo/gizmo_w1.py @@ -28,6 +28,7 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, URDFCfg, JointDrivePropertiesCfg, @@ -55,8 +56,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), visualization=visualization_cfg_from_args(args), ) @@ -128,6 +130,7 @@ def main(): 0.0000e00, ] robot = sim.add_robot(cfg=cfg) + sim.prepare() # Set initial joint positions for both arms # Left arm: 8 joints (WAIST + 7 LEFT_J), Right arm: 8 joints (WAIST + 7 RIGHT_J) diff --git a/examples/sim/planners/curobo_planner.py b/examples/sim/planners/curobo_planner.py index 908f8619a..c2b193932 100644 --- a/examples/sim/planners/curobo_planner.py +++ b/examples/sim/planners/curobo_planner.py @@ -29,6 +29,7 @@ python examples/sim/planners/curobo_planner.py --headless python examples/sim/planners/curobo_planner.py --headless --num_envs 4 python examples/sim/planners/curobo_planner.py --headless --device cuda:1 + python examples/sim/planners/curobo_planner.py --headless --physics newton Requirements: an NVIDIA CUDA device and the CUDA-matched cuRobo V2 source package installed in the active environment. Installation instructions: @@ -65,7 +66,11 @@ MotionPolicy, ) from embodichain.data import get_data_path -from embodichain.lab.sim.cfg import RenderCfg, RigidBodyAttributesCfg +from embodichain.lab.sim.cfg import ( + RenderCfg, + RigidBodyPhysicsCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.objects import RigidObjectCfg, Robot, RigidObject from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator from embodichain.lab.sim.planners.curobo.curobo_planner import ( @@ -243,6 +248,7 @@ def _build_scene( arena_space: float = 2.0, gpu_id: int = 0, visualization: VisualizationCfg | None = None, + physics: str = "default", ) -> tuple[SimulationManager, Robot, RigidObject, torch.Tensor, str]: """Create the batched robot scene with an identical cuboid in each arena.""" sim = SimulationManager( @@ -253,6 +259,7 @@ def _build_scene( arena_space=arena_space, gpu_id=gpu_id, render_cfg=RenderCfg(renderer=renderer), + physics_cfg=physics_cfg_for_backend(physics), visualization=visualization or VisualizationCfg(), ) ) @@ -321,7 +328,7 @@ def _build_scene( "LEFT_HAND_PINKY", ], }, - "drive_pros": { + "joint_drive_props": { "stiffness": {"LEFT_[A-Z|_]+[0-9]?": 1e2}, "damping": {"LEFT_[A-Z|_]+[0-9]?": 1e1}, "max_effort": {"LEFT_[A-Z|_]+[0-9]?": 1e3}, @@ -457,11 +464,6 @@ def _build_scene( if robot is None: raise RuntimeError(f"Failed to add robot '{robot_type}' to the cuRobo demo.") target_xpos = _resolve_batched_target(target_xpos, robot.num_instances) - if robot_type == "w1": - # Keep the W1-specific IK diagnostic batched so it remains useful when - # checking solver and cuRobo reachability across multiple environments. - is_success, ik_qpos = robot.compute_ik(pose=target_xpos, name=control_part) - print(f"robot compute ik success: {is_success}, ik_qpos: {ik_qpos}") # This object is also exported into the cuRobo collision world below via # CuroboWorldCfg.rigid_objects, so the simulator and planner share geometry @@ -470,12 +472,21 @@ def _build_scene( cfg=RigidObjectCfg( uid="demo_block", shape=CubeCfg(size=demo_block_size), - attrs=RigidBodyAttributesCfg(), + # The grouped form is backend-neutral; the deprecated flat attrs + # configuration cannot be spawned by Newton. + attrs=RigidBodyPhysicsCfg(), body_type="kinematic", init_pos=demo_block_position, init_rot=(0.0, 0.0, 0.0), ) ) + sim.prepare() + + if robot_type == "w1": + # Keep the W1-specific IK diagnostic batched so it remains useful when + # checking solver and cuRobo reachability across multiple environments. + is_success, ik_qpos = robot.compute_ik(pose=target_xpos, name=control_part) + print(f"robot compute ik success: {is_success}, ik_qpos: {ik_qpos}") return sim, robot, demo_block, target_xpos, control_part @@ -697,9 +708,8 @@ def main() -> None: args.arena_space, effective_gpu_id, visualization_cfg_from_args(args), + physics=args.physics, ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() obstacles = [demo_block] obstacle_poses = _perturb_obstacles( diff --git a/examples/sim/planners/neural_planner.py b/examples/sim/planners/neural_planner.py index 9ce5a18db..d234f001f 100644 --- a/examples/sim/planners/neural_planner.py +++ b/examples/sim/planners/neural_planner.py @@ -34,7 +34,7 @@ from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import MarkerCfg, RenderCfg +from embodichain.lab.sim.cfg import MarkerCfg, RenderCfg, physics_cfg_for_backend from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg from embodichain.lab.sim.planners import ( @@ -207,11 +207,12 @@ def main() -> None: sim = SimulationManager( SimulationManagerCfg( headless=args.headless, - sim_device=sim_device, + device=sim_device, num_envs=args.num_envs, arena_space=args.arena_space, gpu_id=effective_gpu_id, render_cfg=RenderCfg(renderer=args.renderer), + physics_cfg=physics_cfg_for_backend(args.physics), visualization=visualization_cfg_from_args(args), ) ) @@ -220,8 +221,7 @@ def main() -> None: arm_name = "arm" device = robot.device - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() diff --git a/examples/sim/robot/dexforce_w1.py b/examples/sim/robot/dexforce_w1.py index 9a4e78383..37a71c5d8 100644 --- a/examples/sim/robot/dexforce_w1.py +++ b/examples/sim/robot/dexforce_w1.py @@ -60,7 +60,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: }, ] }, - "drive_pros": { + "joint_drive_props": { "max_effort": { "left_eef": 10.0, "right_eef": 10.0, @@ -70,6 +70,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: ) robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.update(step=1) print("DexforceW1 with a user defined end-effector added to the simulation.") diff --git a/examples/sim/scene/scene_demo.py b/examples/sim/scene/scene_demo.py index aeb39a3a5..a260a2b6b 100644 --- a/examples/sim/scene/scene_demo.py +++ b/examples/sim/scene/scene_demo.py @@ -29,7 +29,8 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( RenderCfg, - RigidBodyAttributesCfg, + physics_cfg_for_backend, + RigidBodyPhysicsCfg, LightCfg, RobotCfg, URDFCfg, @@ -77,9 +78,6 @@ def resolve_asset_path(scene_name: str) -> str: def run_simulation(sim: SimulationManager): """Run the simulation loop.""" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - try: while True: time.sleep(0.01) @@ -119,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, visualization=visualization_cfg_from_args(args), @@ -142,11 +141,15 @@ def main(): cfg = LightCfg(uid=uid, intensity=intensity, radius=600, init_pos=[x, y, z]) lights.append(sim.add_light(cfg)) - physics_attrs = RigidBodyAttributesCfg( - mass=10, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + physics_attrs = RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 10}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.1, + }, + } ) try: @@ -179,6 +182,8 @@ def main(): logger.log_info(f"Failed to load scene asset: {e}") return + sim.prepare() + logger.log_info(f"Scene '{args.scene}' setup complete!") logger.log_info(f"Running simulation with {args.num_envs} environment(s)") logger.log_info("Press Ctrl+C to stop the simulation") diff --git a/examples/sim/sensors/batch_camera.py b/examples/sim/sensors/batch_camera.py index e52b180fc..0af567c7b 100644 --- a/examples/sim/sensors/batch_camera.py +++ b/examples/sim/sensors/batch_camera.py @@ -21,8 +21,13 @@ import matplotlib.pyplot as plt from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import ( + RenderCfg, + physics_cfg_for_backend, + RigidObjectCfg, + LightCfg, +) from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import RenderCfg, RigidObjectCfg, LightCfg from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import RigidObject, Light from embodichain.lab.sim.sensors import ( @@ -38,10 +43,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), visualization=visualization_cfg_from_args(args), ) sim = SimulationManager(config) @@ -54,8 +60,7 @@ def main(args): ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() @@ -117,6 +122,8 @@ def main(args): else: plt.show() + sim.destroy() + if __name__ == "__main__": import argparse diff --git a/examples/sim/sensors/create_contact_sensor.py b/examples/sim/sensors/create_contact_sensor.py index e2332d9d1..e48c59c1e 100644 --- a/examples/sim/sensors/create_contact_sensor.py +++ b/examples/sim/sensors/create_contact_sensor.py @@ -28,8 +28,12 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( + DefaultRigidBodyPropertiesCfg, + MassPropertiesCfg, RenderCfg, - RigidBodyAttributesCfg, + physics_cfg_for_backend, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.sim.sensors import ( ContactSensorCfg, @@ -60,12 +64,14 @@ def create_cube( uid=uid, shape=CubeCfg(size=cube_size), body_type="dynamic", - attrs=RigidBodyAttributesCfg( - mass=0.1, - dynamic_friction=0.9, - static_friction=0.95, - restitution=0.01, - sleep_threshold=0.0, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=0.1), + rigid_props=DefaultRigidBodyPropertiesCfg(sleep_threshold=0.0), + material_props=RigidBodyMaterialCfg( + dynamic_friction=0.9, + static_friction=0.95, + restitution=0.01, + ), ), init_pos=position, ) @@ -151,10 +157,10 @@ def create_robot( }, "init_pos": position, "init_qpos": [0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.0, 0.0], - "drive_pros": { - "stiffness": {"JOINT[1-6]": 1e4, "FINGER[1-2]_JOINT": 1e2}, - "damping": {"JOINT[1-6]": 1e3, "FINGER[1-2]_JOINT": 1e1}, - "max_effort": {"JOINT[1-6]": 1e5, "FINGER[1-2]_JOINT": 1e3}, + "joint_drive_props": { + "stiffness": {"Joint[1-6]": 1e4, "finger[1-2]_joint": 1e2}, + "damping": {"Joint[1-6]": 1e3, "finger[1-2]_joint": 1e1}, + "max_effort": {"Joint[1-6]": 1e5, "finger[1-2]_joint": 1e3}, }, "solver_cfg": { "arm": { @@ -169,7 +175,7 @@ def create_robot( ], } }, - "control_parts": {"arm": ["JOINT[1-6]"], "hand": ["FINGER[1-2]_JOINT"]}, + "control_parts": {"arm": ["Joint[1-6]"], "hand": ["finger[1-2]_joint"]}, } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(robot_cfg_dict)) return robot @@ -192,10 +198,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), visualization=visualization_cfg_from_args(args), ) @@ -207,6 +214,7 @@ def main(): cube1 = create_cube(sim, "cube1", position=[0.0, 0.0, 0.06]) cube2 = create_cube(sim, "cube2", position=[0.0, 0.0, 0.09]) robot = create_robot(sim, "UR10_PGI", position=[0.5, 0.0, 0.0]) + sim.prepare() print("[INFO]: Scene setup complete!") print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") @@ -228,10 +236,6 @@ def run_simulation(sim: SimulationManager): sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 # contact filter config contact_filter_cfg = ContactSensorCfg() @@ -242,6 +246,10 @@ def run_simulation(sim: SimulationManager): contact_filter_cfg.articulation_cfg_list = [contact_filter_art_cfg] contact_filter_cfg.filter_need_both_actor = True + if sim.is_newton_backend: + run_newton_contact_query(sim, contact_filter_cfg) + return + contact_sensor = sim.add_sensor(sensor_cfg=contact_filter_cfg) try: @@ -285,5 +293,125 @@ def run_simulation(sim: SimulationManager): print("[INFO]: Simulation terminated successfully") +def run_newton_contact_query( + sim: SimulationManager, contact_filter_cfg: ContactSensorCfg +) -> None: + """Run Newton's raw contact query for the configured collision shapes. + + The generic :class:`ContactSensor` currently consumes Default-backend + ``PhysicsScene`` buffers. Newton's Spawn runtime instead owns the contact + buffers directly, so this example queries those buffers without claiming + that the generic sensor API is backend-neutral yet. + + Args: + sim: Prepared simulation manager using the Newton backend. + contact_filter_cfg: Rigid objects and articulation links to monitor. + """ + import warp as wp + from dexsim.engine.newton_physics.backend_registry import get_newton_backend + + result = sim.spawn_result + if result is None: + raise RuntimeError("Newton contact queries require a prepared Spawn scene.") + backend = get_newton_backend(result.world) + if backend is None: + raise RuntimeError("Newton Spawn runtime is unavailable for contact queries.") + if not callable(getattr(backend.solver, "update_contacts", None)): + raise RuntimeError( + "The active Newton solver does not expose contact-query support." + ) + + filter_shape_ids = _newton_filter_shape_ids(sim, contact_filter_cfg) + step_count = 0 + accumulated_cost_time = 0.0 + + try: + while True: + sim.update(step=1) + start_time = time.time() + backend.solver.update_contacts(backend.contacts, backend.state_0) + total_contacts = int( + wp.to_torch(backend.contacts.rigid_contact_count).reshape(-1)[0].item() + ) + matched_contacts = 0 + if total_contacts > 0: + shape0 = wp.to_torch(backend.contacts.rigid_contact_shape0)[ + :total_contacts + ] + shape1 = wp.to_torch(backend.contacts.rigid_contact_shape1)[ + :total_contacts + ] + shape0_matches = torch.isin(shape0, filter_shape_ids) + shape1_matches = torch.isin(shape1, filter_shape_ids) + if contact_filter_cfg.filter_need_both_actor: + matched_contacts = int( + torch.logical_and(shape0_matches, shape1_matches).sum().item() + ) + else: + matched_contacts = int( + torch.logical_or(shape0_matches, shape1_matches).sum().item() + ) + accumulated_cost_time += time.time() - start_time + step_count += 1 + + if step_count % 100 == 0: + average_cost_time = accumulated_cost_time / 100.0 + print( + "[INFO]: Fetch Newton contact cost time: " + f"{average_cost_time * 1000:.2f} ms, " + f"contacts: {matched_contacts}, num_envs: {sim.num_envs}" + ) + accumulated_cost_time = 0.0 + except KeyboardInterrupt: + print("\n[INFO]: Stopping simulation...") + finally: + sim.destroy() + print("[INFO]: Simulation terminated successfully") + + +def _newton_filter_shape_ids( + sim: SimulationManager, contact_filter_cfg: ContactSensorCfg +) -> torch.Tensor: + """Resolve a contact filter configuration to Newton Spawn shape IDs.""" + shape_ids: list[int] = [] + for rigid_uid in contact_filter_cfg.rigid_uid_list: + rigid_object = sim.get_rigid_object(rigid_uid) + if rigid_object is None: + continue + for entity in rigid_object._entities: + physics_body = entity.physics_body + if physics_body is not None: + shape_ids.extend(int(shape_id) for shape_id in physics_body.shape_ids) + + for articulation_cfg in contact_filter_cfg.articulation_cfg_list: + articulation = sim.get_robot(articulation_cfg.articulation_uid) + if articulation is None: + articulation = sim.get_articulation(articulation_cfg.articulation_uid) + if articulation is None: + continue + for entity in articulation._entities: + physics_articulation = entity.physics_articulation + if physics_articulation is None: + continue + link_names = ( + set(articulation_cfg.link_name_list) + if articulation_cfg.link_name_list + else {link.name for link in physics_articulation.links} + ) + for link in physics_articulation.links: + if link.name in link_names: + shape_ids.extend(int(shape_id) for shape_id in link.shape_ids) + + if not shape_ids: + raise ValueError( + "The Newton contact filter did not resolve to any collision shapes." + ) + return torch.tensor( + sorted(set(shape_ids)), + dtype=torch.int32, + device=sim.device, + ) + + if __name__ == "__main__": main() diff --git a/examples/sim/solvers/differential_solver.py b/examples/sim/solvers/differential_solver.py index cb9d18c6f..111cd4c53 100644 --- a/examples/sim/solvers/differential_solver.py +++ b/examples/sim/solvers/differential_solver.py @@ -43,11 +43,11 @@ def main( 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, + device=device, arena_space=1.5, num_envs=num_envs, visualization=visualization or VisualizationCfg(), @@ -82,6 +82,7 @@ def main( } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() # Prepare initial joint positions for all environments rad = torch.deg2rad(torch.tensor(45.0)) diff --git a/examples/sim/solvers/neural_ik_solver.py b/examples/sim/solvers/neural_ik_solver.py index 5d2b9a83f..2fdd6ae43 100644 --- a/examples/sim/solvers/neural_ik_solver.py +++ b/examples/sim/solvers/neural_ik_solver.py @@ -94,12 +94,12 @@ def main() -> None: np.set_printoptions(precision=5, suppress=True) torch.set_printoptions(precision=5, sci_mode=False) - sim_device = _resolve_device(args.device) + device = _resolve_device(args.device) num_envs = args.num_envs config = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=device, num_envs=num_envs, arena_space=2.0, visualization=visualization_cfg_from_args(args), @@ -128,6 +128,7 @@ def main() -> None: ) robot: Robot = sim.add_robot(cfg=cfg) + sim.prepare() sim.open_window() @@ -182,7 +183,7 @@ def main() -> None: ik_success_flags: list[torch.Tensor] = [] print( - f"\nRunning {num_steps} batch IK steps: num_envs={num_envs}, device='{sim_device}' ..." + f"\nRunning {num_steps} batch IK steps: num_envs={num_envs}, device='{device}' ..." ) ik_compute_begin = time.time() for step in range(num_steps): diff --git a/examples/sim/solvers/opw_solver.py b/examples/sim/solvers/opw_solver.py index 1583caa59..56ae124eb 100644 --- a/examples/sim/solvers/opw_solver.py +++ b/examples/sim/solvers/opw_solver.py @@ -40,10 +40,10 @@ def main(visualization: VisualizationCfg | None = None) -> None: torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" + device = "cpu" config = SimulationManagerCfg( headless=False, - sim_device=sim_device, + device=device, visualization=visualization or VisualizationCfg(), ) sim = SimulationManager(config) @@ -89,6 +89,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: # Add robot to simulation robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + sim.prepare() # Left arm control arm_name = "left_arm" diff --git a/examples/sim/solvers/pink_solver.py b/examples/sim/solvers/pink_solver.py index fcf90fc82..9d0e71b4e 100644 --- a/examples/sim/solvers/pink_solver.py +++ b/examples/sim/solvers/pink_solver.py @@ -40,10 +40,10 @@ def main(visualization: VisualizationCfg | None = None) -> None: torch.set_printoptions(precision=5, sci_mode=False) # Set up simulation with specified device (CPU or CUDA) - sim_device = "cpu" + device = "cpu" config = SimulationManagerCfg( headless=False, - sim_device=sim_device, + device=device, visualization=visualization or VisualizationCfg(), ) sim = SimulationManager(config) @@ -76,6 +76,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() # Define a sample target pose as a 1x4x4 homogeneous matrix rad = torch.deg2rad(torch.tensor(45.0)) diff --git a/examples/sim/solvers/pinocchio_solver.py b/examples/sim/solvers/pinocchio_solver.py index adf290525..bfc3610a9 100644 --- a/examples/sim/solvers/pinocchio_solver.py +++ b/examples/sim/solvers/pinocchio_solver.py @@ -41,10 +41,10 @@ def main(visualization: VisualizationCfg | None = None) -> None: torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" + device = "cpu" config = SimulationManagerCfg( headless=False, - sim_device=sim_device, + device=device, visualization=visualization or VisualizationCfg(), ) sim = SimulationManager(config) @@ -76,6 +76,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: } robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() arm_name = "left_arm" # Set initial joint positions for left arm qpos_seed = torch.tensor( diff --git a/examples/sim/solvers/pytorch_solver.py b/examples/sim/solvers/pytorch_solver.py index 89af61495..46749573a 100644 --- a/examples/sim/solvers/pytorch_solver.py +++ b/examples/sim/solvers/pytorch_solver.py @@ -41,11 +41,11 @@ def main(visualization: VisualizationCfg | None = None) -> None: 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, + device=device, arena_space=2.0, num_envs=num_envs, visualization=visualization or VisualizationCfg(), @@ -82,6 +82,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: # Add robot to simulation robot: Robot = sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + sim.prepare() # Prepare initial joint positions for all environments arm_name = "left_arm" diff --git a/examples/sim/solvers/srs_solver.py b/examples/sim/solvers/srs_solver.py index 25bbe4d5e..76693f96c 100644 --- a/examples/sim/solvers/srs_solver.py +++ b/examples/sim/solvers/srs_solver.py @@ -39,11 +39,11 @@ def main(visualization: VisualizationCfg | None = None) -> None: torch.set_printoptions(precision=5, sci_mode=False) # Initialize simulation - sim_device = "cpu" + device = "cpu" sim = SimulationManager( SimulationManagerCfg( headless=False, - sim_device=sim_device, + device=device, width=2200, height=1200, visualization=visualization or VisualizationCfg(), @@ -53,6 +53,7 @@ def main(visualization: VisualizationCfg | None = None) -> None: sim.set_manual_update(False) robot: Robot = sim.add_robot(cfg=DexforceW1Cfg.from_dict({"uid": "dexforce_w1"})) + sim.prepare() arm_name = "left_arm" # Set initial joint positions for left arm qpos_fk_list = [ diff --git a/examples/sim/workspace/analyze_cartesian_workspace.py b/examples/sim/workspace/analyze_cartesian_workspace.py index fb9160067..d514c71e2 100644 --- a/examples/sim/workspace/analyze_cartesian_workspace.py +++ b/examples/sim/workspace/analyze_cartesian_workspace.py @@ -101,6 +101,7 @@ def main() -> None: } ) robot = sim.add_robot(cfg=cfg) + sim.prepare() print("DexforceW1 robot added to the simulation.") left_qpos = torch.tensor( diff --git a/examples/sim/workspace/analyze_joint_workspace.py b/examples/sim/workspace/analyze_joint_workspace.py index 3695bdb79..ba96f7ca2 100644 --- a/examples/sim/workspace/analyze_joint_workspace.py +++ b/examples/sim/workspace/analyze_joint_workspace.py @@ -98,6 +98,7 @@ def main() -> None: } ) robot = sim_manager.add_robot(cfg=cfg) + sim_manager.prepare() print("DexforceW1 robot added to the simulation.") analyzer = WorkspaceAnalyzer( diff --git a/examples/sim/workspace/analyze_plane_workspace.py b/examples/sim/workspace/analyze_plane_workspace.py index 95e381e1e..7fbcccb24 100644 --- a/examples/sim/workspace/analyze_plane_workspace.py +++ b/examples/sim/workspace/analyze_plane_workspace.py @@ -101,6 +101,7 @@ def main() -> None: } ) robot = sim.add_robot(cfg=cfg) + sim.prepare() print("DexforceW1 robot added to the simulation.") left_qpos = torch.tensor( diff --git a/scripts/benchmark/atomic_action/common.py b/scripts/benchmark/atomic_action/common.py index a7b7e6db9..bfd726ca3 100644 --- a/scripts/benchmark/atomic_action/common.py +++ b/scripts/benchmark/atomic_action/common.py @@ -29,7 +29,7 @@ from dataclasses import dataclass from datetime import datetime from pathlib import Path -from typing import Callable +from typing import Callable, Literal try: import psutil @@ -82,7 +82,7 @@ class MeshObjectPreset: mesh_path: str = "" shape_type: str = "mesh" cube_size: tuple[float, float, float] | None = None - use_usd_properties: bool = False + asset_physics_mode: Literal["preserve", "overlay"] = "overlay" dynamic_friction: float = 0.97 static_friction: float = 0.99 restitution: float = 0.0 @@ -95,7 +95,10 @@ class MeshObjectPreset: min_velocity_iters: int = 1 max_linear_velocity: float = 100.0 max_angular_velocity: float = 100.0 - max_convex_hull_num: int = 16 + collision_approximation: Literal["convex_hull", "convex_decomposition"] = ( + "convex_decomposition" + ) + max_hulls: int | None = 16 enable_ccd: bool = False @@ -123,7 +126,7 @@ class MeshObjectPreset: body_scale=(0.8, 0.8, 0.8), mass=0.05, initial_z=0.05, - use_usd_properties=False, + asset_physics_mode="overlay", ), "coffee_cup": MeshObjectPreset( object_type="coffee_cup", @@ -134,7 +137,7 @@ class MeshObjectPreset: body_scale=(4.0, 4.0, 4.0), mass=0.01, initial_z=0.01, - use_usd_properties=False, + asset_physics_mode="overlay", ), "cube": MeshObjectPreset( object_type="cube", @@ -146,7 +149,7 @@ class MeshObjectPreset: body_scale=(1.0, 1.0, 1.0), mass=0.05, initial_z=0.05, - use_usd_properties=False, + asset_physics_mode="overlay", dynamic_friction=0.5, static_friction=0.5, contact_offset=0.003, @@ -154,7 +157,8 @@ class MeshObjectPreset: max_depenetration_velocity=10.0, min_position_iters=32, min_velocity_iters=8, - max_convex_hull_num=1, + collision_approximation="convex_hull", + max_hulls=None, ), "paper_cup": MeshObjectPreset( object_type="paper_cup", @@ -165,7 +169,7 @@ class MeshObjectPreset: body_scale=(0.75, 0.75, 1.0), mass=0.01, initial_z=0.05, - use_usd_properties=False, + asset_physics_mode="overlay", dynamic_friction=1.0, static_friction=1.0, contact_offset=0.003, @@ -177,7 +181,7 @@ class MeshObjectPreset: min_velocity_iters=8, max_linear_velocity=5.0, max_angular_velocity=10.0, - max_convex_hull_num=8, + max_hulls=8, ), "scanned_bottle": MeshObjectPreset( object_type="scanned_bottle", @@ -188,7 +192,7 @@ class MeshObjectPreset: body_scale=(1.0, 1.0, 1.0), mass=0.05, initial_z=0.05, - use_usd_properties=False, + asset_physics_mode="overlay", ), } COVERAGE_MESH_OBJECT_TYPES = ("sugar_box", "cube", "paper_cup") @@ -518,11 +522,17 @@ def create_benchmark_object( ): """Create one benchmark object at a selected initial position.""" from embodichain.data import get_data_path - from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg - from embodichain.lab.sim.shapes import CubeCfg, MeshCfg + from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg, RigidObjectCfg + from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, MeshCollisionCfg if preset.shape_type == "mesh": - shape = MeshCfg(fpath=get_data_path(preset.mesh_path)) + shape = MeshCfg( + fpath=get_data_path(preset.mesh_path), + collision=MeshCollisionCfg( + approximation=preset.collision_approximation, + max_hulls=preset.max_hulls, + ), + ) elif preset.shape_type == "cube": if preset.cube_size is None: raise ValueError(f"Cube preset {preset.object_type!r} misses cube_size.") @@ -535,27 +545,34 @@ def create_benchmark_object( cfg = RigidObjectCfg( uid=f"benchmark_{preset.label}_{position_case.name}_{uid_suffix}", shape=shape, - attrs=RigidBodyAttributesCfg( - mass=preset.mass, - dynamic_friction=preset.dynamic_friction, - static_friction=preset.static_friction, - restitution=preset.restitution, - contact_offset=preset.contact_offset, - rest_offset=preset.rest_offset, - linear_damping=preset.linear_damping, - angular_damping=preset.angular_damping, - max_depenetration_velocity=preset.max_depenetration_velocity, - min_position_iters=preset.min_position_iters, - min_velocity_iters=preset.min_velocity_iters, - max_linear_velocity=preset.max_linear_velocity, - max_angular_velocity=preset.max_angular_velocity, - enable_ccd=preset.enable_ccd, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": preset.mass}, + "rigid_props": { + "linear_damping": preset.linear_damping, + "angular_damping": preset.angular_damping, + "max_depenetration_velocity": preset.max_depenetration_velocity, + "min_position_iters": preset.min_position_iters, + "min_velocity_iters": preset.min_velocity_iters, + "max_linear_velocity": preset.max_linear_velocity, + "max_angular_velocity": preset.max_angular_velocity, + "enable_ccd": preset.enable_ccd, + }, + "collision_props": { + "contact_offset": preset.contact_offset, + "rest_offset": preset.rest_offset, + }, + "material_props": { + "dynamic_friction": preset.dynamic_friction, + "static_friction": preset.static_friction, + "restitution": preset.restitution, + }, + } ), - max_convex_hull_num=preset.max_convex_hull_num, init_pos=[position_case.xy[0], position_case.xy[1], preset.initial_z], init_rot=preset.init_rot, body_scale=preset.body_scale, - use_usd_properties=preset.use_usd_properties, + asset_physics_mode=preset.asset_physics_mode, ) obj = sim.add_rigid_object(cfg=cfg) sim.update(step=10) diff --git a/scripts/benchmark/rl/runtime.py b/scripts/benchmark/rl/runtime.py index 216d98d7c..20ea70b55 100644 --- a/scripts/benchmark/rl/runtime.py +++ b/scripts/benchmark/rl/runtime.py @@ -111,7 +111,7 @@ def _build_env_cfg( gym_env_cfg.seed = getattr(gym_env_cfg, "seed", None) gym_env_cfg.sim_cfg.headless = headless gym_env_cfg.sim_cfg.gpu_id = gpu_id - gym_env_cfg.sim_cfg.sim_device = device + gym_env_cfg.sim_cfg.device = device return gym_config_data, gym_env_cfg diff --git a/scripts/tutorials/atomic_action/assemble.py b/scripts/tutorials/atomic_action/assemble.py index c664c90fe..ca0451825 100644 --- a/scripts/tutorials/atomic_action/assemble.py +++ b/scripts/tutorials/atomic_action/assemble.py @@ -48,7 +48,7 @@ MotionPolicy, SceneEntityPose, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.data import get_data_path from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.shapes import CubeCfg, MeshCfg @@ -66,6 +66,7 @@ create_curobo_motion_generator, create_parallel_jaw_grasp_pose_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -170,8 +171,11 @@ def create_assemble_object(sim: SimulationManager) -> RigidObject: return sim.add_rigid_object( cfg=RigidObjectCfg( uid="assemble_object", - shape=MeshCfg(fpath=OBJECT_MESH_PATH, compute_uv=False), - attrs=RigidBodyAttributesCfg( + shape=MeshCfg( + fpath=OBJECT_MESH_PATH, + compute_uv=False, + ), + attrs=create_tutorial_rigid_body_physics( mass=0.01, dynamic_friction=0.97, static_friction=0.99, @@ -183,8 +187,8 @@ def create_assemble_object(sim: SimulationManager) -> RigidObject: min_position_iters=32, min_velocity_iters=8, max_depenetration_velocity=2.0, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=1, init_pos=[ OBJECT_A_XY[0], OBJECT_A_XY[1], @@ -202,7 +206,7 @@ def create_base_object(sim: SimulationManager) -> RigidObject: cfg=RigidObjectCfg( uid="base_object", shape=CubeCfg(size=[CUBE_SIZE, CUBE_SIZE, CUBE_SIZE]), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=1.0, dynamic_friction=0.9, static_friction=0.95, @@ -222,15 +226,20 @@ def create_base_object(sim: SimulationManager) -> RigidObject: def compute_can_half_height(can: RigidObject) -> float: """Return half the soda-can extent along world Z when laid on its side.""" vertices = can.get_vertices(env_ids=[0], scale=True)[0].to(torch.float32) - rotated = vertices @ _CAN_INIT_ROTATION.T + rotation = _CAN_INIT_ROTATION.to(device=vertices.device, dtype=vertices.dtype) + rotated = vertices @ rotation.T extent_z = float(rotated[:, 2].max().item() - rotated[:, 2].min().item()) return 0.5 * extent_z -def make_assemble_to_base_pose(dz: float) -> torch.Tensor: +def make_assemble_to_base_pose( + dz: float, + *, + device: torch.device | str | None = None, +) -> torch.Tensor: """Build the can pose relative to the cube: above it, same orientation.""" - pose = torch.eye(4, dtype=torch.float32) - pose[:3, :3] = _CAN_INIT_ROTATION + pose = torch.eye(4, dtype=torch.float32, device=device) + pose[:3, :3] = _CAN_INIT_ROTATION.to(device=pose.device, dtype=pose.dtype) pose[2, 3] = dz return pose @@ -244,6 +253,7 @@ def run_assemble_demo( create_support_surface(sim) can = create_assemble_object(sim) cube = create_base_object(sim) + sim.prepare() settle_object(sim, can, step=0) clone_local_pose_from_first_env(can) @@ -257,16 +267,20 @@ def run_assemble_demo( can, label="soda_can", ) - motion_gen = create_curobo_motion_generator(robot) + motion_gen = create_curobo_motion_generator( + robot, + use_cuda_graph=args.physics != "newton", + ) left_open, left_close = get_hand_open_close_qpos( robot, hand_control_part="left_hand", close_qpos=HAND_CLOSE_QPOS ) + cube_pose = cube.get_local_pose(to_matrix=True) can_half_z = compute_can_half_height(can) assemble_to_base = make_assemble_to_base_pose( - 0.5 * CUBE_SIZE + can_half_z + ASSEMBLE_MARGIN + 0.5 * CUBE_SIZE + can_half_z + ASSEMBLE_MARGIN, + device=cube_pose.device, ) - cube_pose = cube.get_local_pose(to_matrix=True) assemble_object_target_pose = cube_pose[0] @ assemble_to_base num_envs = robot.get_qpos().shape[0] diff --git a/scripts/tutorials/atomic_action/axis_align.py b/scripts/tutorials/atomic_action/axis_align.py index 8685f3f6c..523372905 100644 --- a/scripts/tutorials/atomic_action/axis_align.py +++ b/scripts/tutorials/atomic_action/axis_align.py @@ -37,7 +37,7 @@ MotionPolicy, ObjectSemantics, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger @@ -48,6 +48,7 @@ create_parallel_jaw_grasp_pose_generator, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -105,12 +106,12 @@ def create_align_object( cfg=RigidObjectCfg( uid="cube", shape=CubeCfg(size=list(OBJECT_SIZE)), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.05, dynamic_friction=0.97, static_friction=0.99, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=16, init_pos=init_pos, ) ) diff --git a/scripts/tutorials/atomic_action/control_dt.py b/scripts/tutorials/atomic_action/control_dt.py index 617567cdc..b1b33eb91 100644 --- a/scripts/tutorials/atomic_action/control_dt.py +++ b/scripts/tutorials/atomic_action/control_dt.py @@ -64,6 +64,7 @@ def main() -> None: args = parse_arguments() sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot) + sim.prepare() engine = AtomicActionEngine(motion_generator=create_toppra_motion_generator(robot)) initial_qpos = robot.get_qpos().clone() diff --git a/scripts/tutorials/atomic_action/coordinated_pickment.py b/scripts/tutorials/atomic_action/coordinated_pickment.py index b3ea01d1f..3d64dd602 100644 --- a/scripts/tutorials/atomic_action/coordinated_pickment.py +++ b/scripts/tutorials/atomic_action/coordinated_pickment.py @@ -43,12 +43,9 @@ CoordinatedPickmentOptions, MotionPolicy, ) -from embodichain.lab.sim.cfg import ( - RigidBodyAttributesCfg, - RigidObjectCfg, -) +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.utils import logger from embodichain.utils.math import matrix_from_euler from scripts.tutorials.atomic_action.scenario_utils import ( @@ -69,6 +66,7 @@ create_parallel_jaw_grasp_pose_generator, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, format_tensor, @@ -222,9 +220,14 @@ def create_pickment_object( cfg=RigidObjectCfg( uid=preset.label, shape=MeshCfg( - fpath=resolve_cached_data_path(preset.mesh_path), compute_uv=False + fpath=resolve_cached_data_path(preset.mesh_path), + compute_uv=False, + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=16, + ), ), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.01, dynamic_friction=0.97, static_friction=0.99, @@ -236,13 +239,14 @@ def create_pickment_object( min_position_iters=32, min_velocity_iters=8, max_depenetration_velocity=2.0, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=16, init_pos=[preset.init_xy[0], preset.init_xy[1], SUPPORT_SURFACE_Z], init_rot=list(preset.init_rot), body_scale=preset.body_scale, ) ) + sim.prepare() obj.cfg.init_pos = compute_supported_init_pos(obj, preset) obj.reset() return obj diff --git a/scripts/tutorials/atomic_action/coordinated_placement.py b/scripts/tutorials/atomic_action/coordinated_placement.py index e78034e5b..eab62b7c1 100644 --- a/scripts/tutorials/atomic_action/coordinated_placement.py +++ b/scripts/tutorials/atomic_action/coordinated_placement.py @@ -49,12 +49,9 @@ MotionPolicy, TaskState, ) -from embodichain.lab.sim.cfg import ( - RigidBodyAttributesCfg, - RigidObjectCfg, -) +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.scenario_utils import ( add_dual_tutorial_robot, @@ -76,6 +73,7 @@ clone_local_pose_from_first_env, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, format_tensor, @@ -231,7 +229,7 @@ def create_table(sim: SimulationManager) -> RigidObject: cfg=RigidObjectCfg( uid="table", shape=MeshCfg(fpath=resolve_cached_data_path(TABLE_MESH_PATH)), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=10.0, dynamic_friction=0.9, static_friction=0.95, @@ -250,9 +248,14 @@ def create_bread(sim: SimulationManager) -> RigidObject: cfg=RigidObjectCfg( uid="bread", shape=MeshCfg( - fpath=resolve_cached_data_path(BREAD_MESH_PATH), compute_uv=False + fpath=resolve_cached_data_path(BREAD_MESH_PATH), + compute_uv=False, + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=8, + ), ), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.01, contact_offset=0.003, rest_offset=0.001, @@ -260,9 +263,9 @@ def create_bread(sim: SimulationManager) -> RigidObject: min_position_iters=32, min_velocity_iters=8, max_depenetration_velocity=10.0, + newton_contact=sim.is_newton_backend, ), body_scale=(1.75, 1.75, 1.75), - max_convex_hull_num=8, init_pos=list(BREAD_INIT_POS), init_rot=list(BREAD_INIT_ROT), ) @@ -275,9 +278,14 @@ def create_pan(sim: SimulationManager) -> RigidObject: cfg=RigidObjectCfg( uid="pan", shape=MeshCfg( - fpath=resolve_cached_data_path(PAN_MESH_PATH), compute_uv=False + fpath=resolve_cached_data_path(PAN_MESH_PATH), + compute_uv=False, + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=16, + ), ), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.01, dynamic_friction=0.97, static_friction=0.99, @@ -289,9 +297,9 @@ def create_pan(sim: SimulationManager) -> RigidObject: min_position_iters=32, min_velocity_iters=8, max_depenetration_velocity=2.0, + newton_contact=sim.is_newton_backend, ), body_scale=(1.75, 1.75, 1.75), - max_convex_hull_num=16, init_pos=list(PAN_INIT_POS), init_rot=list(PAN_INIT_ROT), ) @@ -535,6 +543,7 @@ def run_coordinated_placement_demo( create_table(sim) bread = create_bread(sim) pan = create_pan(sim) + sim.prepare() settle_object(sim, bread, step=0) settle_object(sim, pan, step=0) bread_pose_batch = clone_local_pose_from_first_env(bread) diff --git a/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py b/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py index c7a716a09..d7edd8022 100644 --- a/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py +++ b/scripts/tutorials/atomic_action/dynamic_obstacle_recovery.py @@ -49,7 +49,6 @@ TimedCommandSequence, TrackingPolicy, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg, Robot from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator from embodichain.lab.sim.planners.curobo.curobo_planner import ( @@ -418,12 +417,12 @@ def main() -> None: roughness=0.35, ), ), - attrs=RigidBodyAttributesCfg(), body_type="kinematic", init_pos=list(OBSTACLE_START_POSITION), init_rot=[0.0, 0.0, 0.0], ) ) + sim.prepare() # Initialize GPU physics before planning or recording so the first visible # frame and the initial planning context share the same settled state. sim.update(step=10) @@ -431,6 +430,8 @@ def main() -> None: MotionGenCfg( planner_cfg=CuroboPlannerCfg( robot_uid=robot.uid, + # Newton physics captures CUDA graphs on the same device. + use_cuda_graph=args.physics != "newton", # The coarse default voxel fit under-covers the hand and # fingertips. Keep the denser morphit fit, but no extra radius # padding: 5 mm makes this tutorial's initial pose infeasible. diff --git a/scripts/tutorials/atomic_action/hand_over.py b/scripts/tutorials/atomic_action/hand_over.py index fc110e525..495b649ff 100644 --- a/scripts/tutorials/atomic_action/hand_over.py +++ b/scripts/tutorials/atomic_action/hand_over.py @@ -36,10 +36,10 @@ HandOverOptions, MotionPolicy, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.data import get_data_path from embodichain.lab.sim.objects import RigidObject, Robot -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.scenario_utils import ( add_dual_tutorial_robot, @@ -52,6 +52,7 @@ create_parallel_jaw_grasp_pose_generator, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, get_hand_open_close_qpos, clone_local_pose_from_first_env, @@ -142,19 +143,28 @@ def create_support_surface(sim: SimulationManager) -> RigidObject: ) -def create_handover_object(sim: SimulationManager, args) -> RigidObject: +def create_handover_object( + sim: SimulationManager, + args: argparse.Namespace | None = None, +) -> RigidObject: """Create the mode-specific mesh object on the support surface.""" + is_horizontal = bool(getattr(args, "is_horizontal", False)) mesh_path = ( - HORIZONTAL_OBJECT_MESH_PATH if args.is_horizontal else VERTICAL_OBJECT_MESH_PATH - ) - body_scale = ( - HORIZONTAL_OBJECT_SCALE if args.is_horizontal else VERTICAL_OBJECT_SCALE + HORIZONTAL_OBJECT_MESH_PATH if is_horizontal else VERTICAL_OBJECT_MESH_PATH ) + body_scale = HORIZONTAL_OBJECT_SCALE if is_horizontal else VERTICAL_OBJECT_SCALE return sim.add_rigid_object( cfg=RigidObjectCfg( uid="handover_object", - shape=MeshCfg(fpath=mesh_path, compute_uv=False), - attrs=RigidBodyAttributesCfg( + shape=MeshCfg( + fpath=mesh_path, + compute_uv=False, + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=16, + ), + ), + attrs=create_tutorial_rigid_body_physics( mass=0.01, dynamic_friction=0.97, static_friction=0.99, @@ -166,12 +176,10 @@ def create_handover_object(sim: SimulationManager, args) -> RigidObject: min_position_iters=32, min_velocity_iters=8, max_depenetration_velocity=2.0, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=16, init_pos=[OBJECT_INIT_XY[0], OBJECT_INIT_XY[1], SUPPORT_SURFACE_Z + 0.12], - init_rot=( - OBJECT_ROT_VERTICAL if not args.is_horizontal else OBJECT_ROT_HORIZONTAL - ), + init_rot=(OBJECT_ROT_HORIZONTAL if is_horizontal else OBJECT_ROT_VERTICAL), body_scale=body_scale, ) ) @@ -185,6 +193,7 @@ def run_handover_demo( """Plan and optionally execute one unified pick-up and handover.""" create_support_surface(sim) obj = create_handover_object(sim, args) + sim.prepare() settle_object(sim, obj, step=0) clone_local_pose_from_first_env(obj) obj.clear_dynamics() diff --git a/scripts/tutorials/atomic_action/move_end_effector.py b/scripts/tutorials/atomic_action/move_end_effector.py index 9993916e1..be696a34b 100644 --- a/scripts/tutorials/atomic_action/move_end_effector.py +++ b/scripts/tutorials/atomic_action/move_end_effector.py @@ -65,6 +65,7 @@ def main() -> None: args = parse_arguments() sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot) + sim.prepare() motion_gen = create_curobo_motion_generator(robot) engine = AtomicActionEngine(motion_generator=motion_gen) diff --git a/scripts/tutorials/atomic_action/move_held_object.py b/scripts/tutorials/atomic_action/move_held_object.py index efc553ac1..4548b81e4 100644 --- a/scripts/tutorials/atomic_action/move_held_object.py +++ b/scripts/tutorials/atomic_action/move_held_object.py @@ -38,9 +38,9 @@ MotionPolicy, PickUpOptions, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_tutorial_robot, @@ -50,6 +50,7 @@ create_curobo_motion_generator, create_parallel_jaw_grasp_pose_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -62,6 +63,7 @@ OBJECT_MESH_PATH = "PaperCup/paper_cup.ply" OBJECT_XY = (-0.42, -0.08) +OBJECT_INITIAL_Z = 0.05 MOVE_SAMPLE_INTERVAL = 60 PICK_SAMPLE_INTERVAL = 120 MOVE_HELD_OBJECT_SAMPLE_INTERVAL = 120 @@ -83,17 +85,24 @@ def create_pick_object(sim) -> RigidObject: obj = sim.add_rigid_object( cfg=RigidObjectCfg( uid="paper_cup", - shape=MeshCfg(fpath=get_data_path(OBJECT_MESH_PATH)), - attrs=RigidBodyAttributesCfg( + shape=MeshCfg( + fpath=get_data_path(OBJECT_MESH_PATH), + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=16, + ), + ), + attrs=create_tutorial_rigid_body_physics( mass=0.01, dynamic_friction=0.97, static_friction=0.99, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=16, - init_pos=[*OBJECT_XY, 0.0], + init_pos=[*OBJECT_XY, OBJECT_INITIAL_Z], body_scale=(0.75, 0.75, 1.0), ) ) + sim.prepare() sim.update(step=10) clone_local_pose_from_first_env(obj) obj.clear_dynamics() @@ -118,6 +127,7 @@ def main() -> None: sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot) obj = create_pick_object(sim) + sim.prepare() motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) diff --git a/scripts/tutorials/atomic_action/move_joints.py b/scripts/tutorials/atomic_action/move_joints.py index 0a35a5b9f..4a7a2a5e2 100644 --- a/scripts/tutorials/atomic_action/move_joints.py +++ b/scripts/tutorials/atomic_action/move_joints.py @@ -64,6 +64,7 @@ def main() -> None: args = parse_arguments() sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot) + sim.prepare() motion_gen = create_curobo_motion_generator(robot) home = robot.get_qpos(name="arm")[0].clone() diff --git a/scripts/tutorials/atomic_action/moving_target_recovery.py b/scripts/tutorials/atomic_action/moving_target_recovery.py index 3e007477a..8155e2670 100644 --- a/scripts/tutorials/atomic_action/moving_target_recovery.py +++ b/scripts/tutorials/atomic_action/moving_target_recovery.py @@ -54,7 +54,7 @@ TaskState, TrackingPolicy, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger @@ -62,6 +62,7 @@ add_tutorial_robot, create_curobo_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -93,7 +94,7 @@ class _MovingTargetScene: - """Publish a versioned target pose and physically push it exactly once.""" + """Publish a versioned target pose and move it exactly once.""" def __init__( self, @@ -136,7 +137,11 @@ def push( force_duration: float, force_magnitude: float, ) -> torch.Tensor: - """Push the visible target with a short force pulse. + """Move the visible target with backend-appropriate behavior. + + The Default backend demonstrates a physical force pulse. Newton uses a + deterministic pose update because its tutorial target is kinematic, + avoiding an unbounded impulse while the runner is replanning. Args: clock: Simulation adapter used to advance physics. @@ -145,7 +150,7 @@ def push( force_magnitude: Magnitude of the applied force in newtons. Returns: - Batched target pose after the physical motion. + Batched target pose after the move. """ if self.moved: return self.target.get_local_pose(to_matrix=True) @@ -171,7 +176,14 @@ def push( raise ValueError("destination must differ from the current planar pose.") force = force_magnitude * planar_offset / planar_distance.unsqueeze(-1) - self.target.set_body_type("dynamic") + if clock.simulation.is_newton_backend: + moved_pose = start_pose.clone() + moved_pose[:, :3, 3] = self.destination + self.target.set_local_pose(moved_pose) + self.version += 1 + self.moved = True + return moved_pose + self.target.clear_dynamics() step_count = max(1, math.ceil(duration / clock.physics_dt)) force_step_count = min( @@ -191,7 +203,7 @@ def push( def _create_moving_target(sim: SimulationManager) -> RigidObject: - """Create the bright cube, held kinematic until the physical push.""" + """Create the bright cube used for target-motion recovery.""" return sim.add_rigid_object( cfg=RigidObjectCfg( uid=TARGET_ENTITY_ID, @@ -204,14 +216,14 @@ def _create_moving_target(sim: SimulationManager) -> RigidObject: roughness=0.3, ), ), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.05, dynamic_friction=0.97, static_friction=0.99, enable_ccd=True, + newton_contact=sim.is_newton_backend, ), - body_type="kinematic", - max_convex_hull_num=16, + body_type="kinematic" if sim.is_newton_backend else "dynamic", init_pos=INITIAL_TARGET_POSITION, ) ) @@ -245,6 +257,7 @@ def main() -> None: sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot) target = _create_moving_target(sim) + sim.prepare() sim.update(step=10) target_scene = _MovingTargetScene(target, MOVED_TARGET_POSITION) sim_runtime = SimulationExecutionAdapter( @@ -253,16 +266,37 @@ def main() -> None: control_dt=2.0 * sim.sim_config.physics_dt, scene_supplier=target_scene.snapshot, ) - motion_gen = create_curobo_motion_generator(robot) + motion_gen = create_curobo_motion_generator( + robot, + use_cuda_graph=args.physics != "newton", + ) hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, target, hand_open) if args.no_target_motion: - target.set_body_type("dynamic") target.clear_dynamics() target_to_grasp = make_top_down_eef_pose( torch.zeros(3, dtype=torch.float32, device=sim.device) ) + + def attach_target_to_end_effector() -> None: + """Apply the logical grasp pose when Newton cannot retain contacts.""" + eef_pose = robot.compute_fk( + qpos=robot.get_qpos(name="arm"), + name="arm", + to_matrix=True, + ) + target_to_eef = ( + torch.linalg.inv(target_to_grasp) + .unsqueeze(0) + .expand( + eef_pose.shape[0], + -1, + -1, + ) + ) + target.set_local_pose(torch.bmm(eef_pose, target_to_eef)) + initial_target_pose = target.get_local_pose(to_matrix=True) draw_axis_marker( sim, @@ -345,9 +379,13 @@ def on_step(step: RunnerStep) -> None: and not target_scene.moved and step.command_count >= MOVE_AFTER_COMMAND ): + motion_description = ( + "Moving the blue target kinematically" + if sim.is_newton_backend + else f"Applying a {TARGET_PUSH_FORCE:.2f} N force pulse to the blue target" + ) logger.log_warning( - f"Applying a {TARGET_PUSH_FORCE:.2f} N force pulse to the blue " - "target while the robot holds its current command." + f"{motion_description} while the robot holds its current command." ) moved_pose = target_scene.push( sim_runtime, @@ -366,7 +404,7 @@ def on_step(step: RunnerStep) -> None: dim=1, ) logger.log_warning( - "The force pulse moved the blue target after " + "The target moved after " f"{step.command_count} accepted commands by " f"{displacement.detach().cpu().tolist()} m; the original goal " "axis remains visible." @@ -401,14 +439,20 @@ def on_step(step: RunnerStep) -> None: and not pickup_dynamics_cleared and step.command_count - plan_start_command >= clear_after_pick_command ): + if sim.is_newton_backend: + attach_target_to_end_effector() target.clear_dynamics() pickup_dynamics_cleared = True + elif pickup_dynamics_cleared and sim.is_newton_backend: + attach_target_to_end_effector() def verify_pickup_effect( _context: PlanningContext, request: EffectVerificationRequest, ) -> EffectVerificationResult: """Verify that the cube rose with, and remains near, the end effector.""" + if sim.is_newton_backend: + attach_target_to_end_effector() cube_position = target.get_local_pose(to_matrix=True)[:, :3, 3] eef_position = robot.compute_fk( qpos=robot.get_qpos(name="arm"), diff --git a/scripts/tutorials/atomic_action/open_door.py b/scripts/tutorials/atomic_action/open_door.py index cab6cfebb..c9c7d88cc 100644 --- a/scripts/tutorials/atomic_action/open_door.py +++ b/scripts/tutorials/atomic_action/open_door.py @@ -43,12 +43,13 @@ from embodichain.lab.sim.cfg import ( ArticulationCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.sim.objects import Articulation from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, + configure_newton_link_contacts, create_parallel_jaw_grasp_pose_generator, create_toppra_motion_generator, create_tutorial_argument_parser, @@ -91,20 +92,24 @@ def parse_arguments() -> argparse.Namespace: def create_microwave(sim: SimulationManager) -> Articulation: """Create the fixed-base microwave with an unactuated door hinge.""" - microwave = sim.add_articulation( - cfg=ArticulationCfg( - uid="microwave", - fpath=get_data_path(MICROWAVE_ASSET), - init_pos=MICROWAVE_POSITION, - init_rot=MICROWAVE_ORIENTATION, - drive_pros=JointDrivePropertiesCfg(drive_type="none"), - attrs=RigidBodyAttributesCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - fix_base=True, - ) + microwave_cfg = ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + asset_physics_mode="overlay", + init_pos=MICROWAVE_POSITION, + init_rot=MICROWAVE_ORIENTATION, + joint_drive_props=JointDrivePropertiesCfg(drive_type="none"), + attrs=RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": 1.0, "dynamic_friction": 1.0}} + ), + ) + configure_newton_link_contacts( + sim, + microwave_cfg, + group_name="newton_handle_contacts", + link_names_expr=[HANDLE_LINK_NAME], ) + microwave = sim.add_articulation(cfg=microwave_cfg) sim.update(step=10) return microwave diff --git a/scripts/tutorials/atomic_action/pickup.py b/scripts/tutorials/atomic_action/pickup.py index d67a154dd..8075e5ed3 100644 --- a/scripts/tutorials/atomic_action/pickup.py +++ b/scripts/tutorials/atomic_action/pickup.py @@ -35,7 +35,7 @@ PickUpOptions, MotionPolicy, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger @@ -46,6 +46,7 @@ create_curobo_motion_generator, create_parallel_jaw_grasp_pose_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -87,15 +88,16 @@ def create_pick_object(sim) -> RigidObject: cfg=RigidObjectCfg( uid="cube", shape=CubeCfg(size=list(OBJECT_SIZE)), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.05, dynamic_friction=0.97, static_friction=0.99, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=16, init_pos=[*OBJECT_XY, OBJECT_SIZE[2]], ) ) + sim.prepare() sim.update(step=10) clone_local_pose_from_first_env(obj) obj.clear_dynamics() @@ -127,9 +129,13 @@ def main() -> None: sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot, tcp_z=0.15) obj = create_pick_object(sim) + sim.prepare() hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) - motion_gen = create_curobo_motion_generator(robot) + motion_gen = create_curobo_motion_generator( + robot, + use_cuda_graph=args.physics != "newton", + ) engine = create_simulation_atomic_action_engine( motion_generator=motion_gen, diff --git a/scripts/tutorials/atomic_action/place.py b/scripts/tutorials/atomic_action/place.py index f4e0b22ed..7d6e24dfc 100644 --- a/scripts/tutorials/atomic_action/place.py +++ b/scripts/tutorials/atomic_action/place.py @@ -37,7 +37,7 @@ PlaceOptions, MotionPolicy, ) -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import RigidObjectCfg from embodichain.lab.sim.objects import RigidObject from embodichain.lab.sim.shapes import CubeCfg from embodichain.utils import logger @@ -50,6 +50,7 @@ create_curobo_motion_generator, create_parallel_jaw_grasp_pose_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -84,16 +85,17 @@ def create_pick_object(sim) -> RigidObject: cfg=RigidObjectCfg( uid="cube", shape=CubeCfg(size=list(OBJECT_SIZE)), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=0.05, dynamic_friction=0.97, static_friction=0.99, enable_ccd=True, + newton_contact=sim.is_newton_backend, ), - max_convex_hull_num=16, init_pos=[*OBJECT_XY, 0.5 * OBJECT_SIZE[2]], ) ) + sim.prepare() sim.update(step=10) clone_local_pose_from_first_env(obj) obj.clear_dynamics() @@ -125,6 +127,7 @@ def main() -> None: sim = create_tutorial_simulation(args) robot = add_tutorial_robot(sim, args.robot, tcp_z=0.15) obj = create_pick_object(sim) + sim.prepare() motion_gen = create_curobo_motion_generator(robot) hand_open, hand_close = get_hand_open_close_qpos(robot) initialize_pre_pick_robot_pose(robot, obj, hand_open) diff --git a/scripts/tutorials/atomic_action/press.py b/scripts/tutorials/atomic_action/press.py index 8e0400776..55a51a916 100644 --- a/scripts/tutorials/atomic_action/press.py +++ b/scripts/tutorials/atomic_action/press.py @@ -48,8 +48,10 @@ from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, + configure_newton_link_contacts, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, get_hand_open_close_qpos, prepare_tutorial_scene, @@ -94,19 +96,27 @@ def parse_arguments() -> argparse.Namespace: def create_microwave(sim) -> Articulation: """Create the fixed-base microwave articulation used by the demo.""" - microwave = sim.add_articulation( - cfg=ArticulationCfg( - uid="microwave", - fpath=get_data_path(MICROWAVE_ASSET), - init_pos=MICROWAVE_POSITION, - init_qpos=(0, 0, 0, 0), - init_rot=MICROWAVE_ORIENTATION, - drive_pros=JointDrivePropertiesCfg( - stiffness=1e-3, damping=1e2, max_effort=1e-2 - ), - fix_base=True, - ) + microwave_cfg = ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + asset_physics_mode="overlay", + init_pos=MICROWAVE_POSITION, + init_qpos=(0, 0, 0, 0), + init_rot=MICROWAVE_ORIENTATION, + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + stiffness=1e-3, + damping=1e2, + max_effort=1e-2, + ), ) + configure_newton_link_contacts( + sim, + microwave_cfg, + group_name="newton_button_contacts", + link_names_expr=[BUTTON_LINK_NAME], + ) + microwave = sim.add_articulation(cfg=microwave_cfg) sim.update(step=10) return microwave @@ -117,6 +127,9 @@ def create_rigid_button(sim) -> RigidObject: cfg=RigidObjectCfg( uid="rigid_button", shape=CubeCfg(size=list(RIGID_BUTTON_SIZE)), + attrs=create_tutorial_rigid_body_physics( + newton_contact=sim.is_newton_backend, + ), body_type="static", init_pos=RIGID_BUTTON_POSITION, ) @@ -182,6 +195,7 @@ def main() -> None: sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] ) target = create_rigid_button(sim) if args.rigid_object else create_microwave(sim) + sim.prepare() hand_open, hand_close = get_hand_open_close_qpos(robot, close_qpos=0.040) motion_gen = create_toppra_motion_generator(robot) semantics, target_pose = create_button_semantics(target) diff --git a/scripts/tutorials/atomic_action/scenario_utils.py b/scripts/tutorials/atomic_action/scenario_utils.py index 055947a0c..d74e48c42 100644 --- a/scripts/tutorials/atomic_action/scenario_utils.py +++ b/scripts/tutorials/atomic_action/scenario_utils.py @@ -31,7 +31,6 @@ from embodichain.lab.sim import SimulationManager from embodichain.lab.sim.atomic_actions import Affordance, ObjectSemantics from embodichain.lab.sim.cfg import ( - RigidBodyAttributesCfg, RigidObjectCfg, RobotCfg, ) @@ -43,6 +42,8 @@ from scripts.tutorials.atomic_action.tutorial_utils import ( ROBOTIQ_2F_140_TCP, TutorialRobot, + configure_newton_gripper_contacts, + create_tutorial_rigid_body_physics, create_tutorial_robot_cfg, ) @@ -159,7 +160,7 @@ def create_dual_tutorial_robot_cfg( ("damping", hand_damping), ("max_effort", hand_max_effort), ): - getattr(base_cfg.drive_pros, property_name)[hand_joint_pattern] = value + getattr(base_cfg.joint_drive_props, property_name)[hand_joint_pattern] = value arm_facing_rotation = make_yaw_transform( (0.0, 0.0, 0.0), @@ -256,24 +257,24 @@ def add_dual_tutorial_robot( Returns: The added dual-arm robot instance. """ - return sim.add_robot( - cfg=create_dual_tutorial_robot_cfg( - robot_type=robot_type, - uid=uid, - urdf_name=urdf_name, - tcp_z=tcp_z, - solver=solver, - ur_ik_nearest_weight=ur_ik_nearest_weight, - pytorch_num_samples=pytorch_num_samples, - init_pos=init_pos, - init_rot=init_rot, - left_arm_home=left_arm_home, - right_arm_home=right_arm_home, - hand_stiffness=hand_stiffness, - hand_damping=hand_damping, - hand_max_effort=hand_max_effort, - ) + robot_cfg = create_dual_tutorial_robot_cfg( + robot_type=robot_type, + uid=uid, + urdf_name=urdf_name, + tcp_z=tcp_z, + solver=solver, + ur_ik_nearest_weight=ur_ik_nearest_weight, + pytorch_num_samples=pytorch_num_samples, + init_pos=init_pos, + init_rot=init_rot, + left_arm_home=left_arm_home, + right_arm_home=right_arm_home, + hand_stiffness=hand_stiffness, + hand_damping=hand_damping, + hand_max_effort=hand_max_effort, ) + configure_newton_gripper_contacts(sim, robot_cfg) + return sim.add_robot(cfg=robot_cfg) def add_support_surface( @@ -287,7 +288,7 @@ def add_support_surface( cfg=RigidObjectCfg( uid="support_surface", shape=CubeCfg(size=list(size)), - attrs=RigidBodyAttributesCfg( + attrs=create_tutorial_rigid_body_physics( mass=10.0, dynamic_friction=0.9, static_friction=0.95, @@ -302,8 +303,6 @@ def add_support_surface( def settle_object(sim: SimulationManager, obj: RigidObject, step: int = 5) -> None: """Reset, settle, and freeze an object before tutorial planning.""" - if sim.device.type == "cuda": - sim.init_gpu_physics() obj.reset() if step > 0: sim.update(step=step) diff --git a/scripts/tutorials/atomic_action/slide.py b/scripts/tutorials/atomic_action/slide.py index 9780d77a9..d069e9e76 100644 --- a/scripts/tutorials/atomic_action/slide.py +++ b/scripts/tutorials/atomic_action/slide.py @@ -44,15 +44,16 @@ from embodichain.lab.sim.cfg import ( ArticulationCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, ) from embodichain.lab.sim.objects import Articulation from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, + configure_newton_link_contacts, create_parallel_jaw_grasp_pose_generator, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, draw_axis_marker, get_hand_open_close_qpos, @@ -87,21 +88,26 @@ def create_drawer( sim: SimulationManager, ) -> Articulation: """Create the fixed-base drawer in its closed initial state.""" - drawer = sim.add_articulation( - cfg=ArticulationCfg( - uid="drawer", - fpath=get_data_path(DRAWER_ASSET), - init_pos=DRAWER_POSITION, - init_rot=DRAWER_ORIENTATION, - init_qpos=(0.0,), - drive_pros=JointDrivePropertiesCfg(drive_type="none"), - attrs=RigidBodyAttributesCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - fix_base=True, - ) + drawer_cfg = ArticulationCfg( + uid="drawer", + fpath=get_data_path(DRAWER_ASSET), + asset_physics_mode="overlay", + init_pos=DRAWER_POSITION, + init_rot=DRAWER_ORIENTATION, + init_qpos=(0.0,), + joint_drive_props=JointDrivePropertiesCfg(drive_type="none"), + attrs=create_tutorial_rigid_body_physics( + static_friction=1.0, + dynamic_friction=1.0, + ), + ) + configure_newton_link_contacts( + sim, + drawer_cfg, + group_name="newton_handle_contacts", + link_names_expr=[HANDLE_LINK_NAME], ) + drawer = sim.add_articulation(cfg=drawer_cfg) sim.update(step=10) return drawer @@ -185,6 +191,7 @@ def main() -> None: sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0], tcp_z=0.15 ) drawer = create_drawer(sim) + sim.prepare() hand_open, hand_close = get_hand_open_close_qpos(robot) motion_gen = create_toppra_motion_generator(robot) semantics = create_drawer_semantics(drawer) diff --git a/scripts/tutorials/atomic_action/tutorial_utils.py b/scripts/tutorials/atomic_action/tutorial_utils.py index 7171c7d21..351ba9706 100644 --- a/scripts/tutorials/atomic_action/tutorial_utils.py +++ b/scripts/tutorials/atomic_action/tutorial_utils.py @@ -36,7 +36,23 @@ ObjectSemantics, TimedTrajectory, ) -from embodichain.lab.sim.cfg import LightCfg, MarkerCfg, RenderCfg, RobotCfg +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + CollisionPropertiesCfg, + DefaultRigidBodyPropertiesCfg, + LightCfg, + LinkPhysicsOverrideCfg, + MassPropertiesCfg, + MarkerCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + PhysicsBackendCfg, + RenderCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RobotCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.planners import ( CuroboPlannerCfg, @@ -88,9 +104,15 @@ palm_depth=0.096, ) DEFAULT_GRIPPER_CLOSE_QPOS = 0.024 +NEWTON_GRASP_CONTACT_STIFFNESS = 4.0e4 +NEWTON_GRASP_CONTACT_DAMPING = 4.0e2 DEFAULT_TUTORIAL_LIGHT_POS = (1.0, 0.0, 3.0) _FRANKA_TUTORIAL_BASE_ROTATION = (0.0, 0.0, 180.0) _DEFAULT_GRIPPER_TCP_Z = 0.17 +_GRIPPER_CONTACT_LINK_PATTERN = ( + r"(?:.*_)?(?:gripper_finger[12]_link_1|" + r"(?:left|right)_(?:outer|inner)_(?:finger(?:_pad)?|knuckle))" +) _GRIPPER_TCP = ( (1.0, 0.0, 0.0, 0.0), (0.0, 1.0, 0.0, 0.0), @@ -177,6 +199,32 @@ def create_tutorial_argument_parser( return parser +def _tutorial_physics_cfg( + backend: Literal["default", "newton"], +) -> PhysicsBackendCfg: + """Build the shared physics configuration for atomic-action tutorials.""" + physics_cfg = physics_cfg_for_backend(backend) + if isinstance(physics_cfg, NewtonPhysicsCfg): + # Follow Newton's brick-stacking grasp profile. Keep the tutorial's + # smaller 0.5 ms solver step, but align the contact path, nonlinear + # solver, friction cone, impedance ratio, and contact capacities. + contact_max = 16_384 + physics_cfg.num_substeps = 10 + physics_cfg.collision_cfg.reduce_contacts = True + physics_cfg.collision_cfg.rigid_contact_max = contact_max + physics_cfg.solver_cfg = { + "solver_type": "mujoco_warp", + "solver": "newton", + "integrator": "implicitfast", + "iterations": 15, + "ls_iterations": 100, + "nconmax": contact_max, + "njmax": contact_max * 2, + "use_mujoco_contacts": False, + } + return physics_cfg + + def create_tutorial_simulation( args: argparse.Namespace, *, @@ -200,7 +248,8 @@ def create_tutorial_simulation( height=VIEWER_HEIGHT, headless=True, num_envs=args.num_envs, - sim_device=args.device, + device=args.device, + physics_cfg=_tutorial_physics_cfg(args.physics), render_cfg=RenderCfg(renderer=args.renderer), physics_dt=1.0 / 100.0, arena_space=arena_space, @@ -270,13 +319,13 @@ def add_ur5_gripper_robot( Returns: The added robot instance. """ - return sim.add_robot( - cfg=create_ur5_gripper_robot_cfg( - init_pos=init_pos, - init_qpos=init_qpos, - tcp_z=tcp_z, - ) + robot_cfg = create_ur5_gripper_robot_cfg( + init_pos=init_pos, + init_qpos=init_qpos, + tcp_z=tcp_z, ) + # configure_newton_gripper_contacts(sim, robot_cfg) + return sim.add_robot(cfg=robot_cfg) def add_tutorial_robot( @@ -300,14 +349,13 @@ def add_tutorial_robot( Raises: ValueError: If ``robot_type`` is not supported. """ - return sim.add_robot( - cfg=create_tutorial_robot_cfg( - robot_type, - init_pos=init_pos, - init_qpos=init_qpos, - **kwargs, - ) + robot_cfg = create_tutorial_robot_cfg( + robot_type, + init_pos=init_pos, + init_qpos=init_qpos, + **kwargs, ) + return sim.add_robot(cfg=robot_cfg) def create_toppra_motion_generator(robot: Robot) -> MotionGenerator: @@ -324,17 +372,114 @@ def create_toppra_motion_generator(robot: Robot) -> MotionGenerator: ) -def create_curobo_motion_generator(robot: Robot) -> MotionGenerator: +def create_tutorial_rigid_body_physics( + *, + mass: float | None = None, + static_friction: float | None = None, + dynamic_friction: float | None = None, + restitution: float | None = None, + linear_damping: float | None = None, + angular_damping: float | None = None, + max_depenetration_velocity: float | None = None, + enable_ccd: bool | None = None, + min_position_iters: int | None = None, + min_velocity_iters: int | None = None, + contact_offset: float | None = None, + rest_offset: float | None = None, + newton_contact: bool = False, +) -> RigidBodyPhysicsCfg: + """Create portable rigid-body physics for an atomic-action tutorial. + + Material and mass values apply to both physics backends. The remaining + values are retained in the Default-backend configuration group; Newton + safely ignores those properties because it has no equivalent controls. + Set ``newton_contact`` only for a manipulation contact surface in a Newton + scene to use the same less-compliant response as the drawer tutorial. + + Args: + newton_contact: Whether to add the Newton-only contact stiffness and + damping used on grasped or directly manipulated objects. + + Returns: + Grouped physics configuration accepted by both tutorial backends. + """ + rigid_values = ( + linear_damping, + angular_damping, + max_depenetration_velocity, + enable_ccd, + min_position_iters, + min_velocity_iters, + ) + collision_values = (contact_offset, rest_offset) + material_values = (static_friction, dynamic_friction, restitution) + return RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=mass) if mass is not None else None, + rigid_props=( + DefaultRigidBodyPropertiesCfg( + linear_damping=linear_damping, + angular_damping=angular_damping, + max_depenetration_velocity=max_depenetration_velocity, + enable_ccd=enable_ccd, + min_position_iters=min_position_iters, + min_velocity_iters=min_velocity_iters, + ) + if any(value is not None for value in rigid_values) + else None + ), + collision_props=( + CollisionPropertiesCfg( + contact_offset=contact_offset, + rest_offset=rest_offset, + ) + if any(value is not None for value in collision_values) + else None + ), + material_props=( + ( + NewtonRigidBodyMaterialCfg( + static_friction=static_friction, + dynamic_friction=dynamic_friction, + restitution=restitution, + ke=NEWTON_GRASP_CONTACT_STIFFNESS, + kd=NEWTON_GRASP_CONTACT_DAMPING, + ) + if newton_contact + else RigidBodyMaterialCfg( + static_friction=static_friction, + dynamic_friction=dynamic_friction, + restitution=restitution, + ) + ) + if newton_contact or any(value is not None for value in material_values) + else None + ), + ) + + +def create_curobo_motion_generator( + robot: Robot, + *, + use_cuda_graph: bool = True, +) -> MotionGenerator: """Create a cuRobo-backed motion generator for a tutorial robot. Args: robot: Robot whose trajectories will be planned. + use_cuda_graph: Whether cuRobo may capture CUDA graphs. Disable this + when the tutorial uses Newton physics, which owns CUDA graph + capture on the same device. Returns: The configured motion generator with an empty external collision world. """ return MotionGenerator( - cfg=MotionGenCfg(planner_cfg=CuroboPlannerCfg(robot_uid=robot.uid)) + cfg=MotionGenCfg( + planner_cfg=CuroboPlannerCfg( + robot_uid=robot.uid, + use_cuda_graph=use_cuda_graph, + ) + ) ) @@ -916,7 +1061,7 @@ def create_ur5_gripper_robot_cfg( "control_parts": { "hand": [GRIPPER_HAND_JOINT_PATTERN], }, - "drive_pros": { + "joint_drive_props": { "stiffness": { GRIPPER_HAND_JOINT_PATTERN: 1e3, }, @@ -980,7 +1125,7 @@ def create_franka_panda_robot_cfg( ], }, "control_parts": {"hand": [GRIPPER_HAND_JOINT_PATTERN]}, - "drive_pros": { + "joint_drive_props": { "stiffness": {GRIPPER_HAND_JOINT_PATTERN: 1e3}, "damping": {GRIPPER_HAND_JOINT_PATTERN: 1e2}, "max_effort": {GRIPPER_HAND_JOINT_PATTERN: 1e4}, @@ -998,9 +1143,9 @@ def create_franka_panda_robot_cfg( if init_qpos is None: cfg.init_qpos[-2:] = [0.0, 0.0] for drive_values in ( - cfg.drive_pros.stiffness, - cfg.drive_pros.damping, - cfg.drive_pros.max_effort, + cfg.joint_drive_props.stiffness, + cfg.joint_drive_props.damping, + cfg.joint_drive_props.max_effort, ): drive_values.pop("fr3_finger_joint[1-2]", None) return cfg @@ -1051,7 +1196,7 @@ def create_ur10_robotiq_robot_cfg( "control_parts": { "hand": [ROBOTIQ_HAND_JOINT_PATTERN], }, - "drive_pros": { + "joint_drive_props": { "stiffness": {ROBOTIQ_HAND_JOINT_PATTERN: 1e3}, "damping": {ROBOTIQ_HAND_JOINT_PATTERN: 1e2}, "max_effort": {ROBOTIQ_HAND_JOINT_PATTERN: 1e3}, @@ -1117,6 +1262,8 @@ def create_tutorial_robot_cfg( "ROBOTIQ_2F_140_TCP", "ROBOTIQ_2F_140_URDF_PATH", "ROBOTIQ_HAND_JOINT_PATTERN", + "NEWTON_GRASP_CONTACT_DAMPING", + "NEWTON_GRASP_CONTACT_STIFFNESS", "TOP_DOWN_EEF_ROTATION", "TutorialCliFeature", "TutorialRobot", @@ -1126,6 +1273,8 @@ def create_tutorial_robot_cfg( "broadcast_pose_batch", "broadcast_waypoint_pose_batch", "clone_local_pose_from_first_env", + "configure_newton_gripper_contacts", + "configure_newton_link_contacts", "create_antipodal_semantics", "create_parallel_jaw_grasp_pose_generator", "create_curobo_motion_generator", diff --git a/scripts/tutorials/atomic_action/twist.py b/scripts/tutorials/atomic_action/twist.py index 88cccab65..5b23daba3 100644 --- a/scripts/tutorials/atomic_action/twist.py +++ b/scripts/tutorials/atomic_action/twist.py @@ -48,8 +48,10 @@ from embodichain.utils import logger from scripts.tutorials.atomic_action.tutorial_utils import ( add_ur5_gripper_robot, + configure_newton_link_contacts, create_toppra_motion_generator, create_tutorial_argument_parser, + create_tutorial_rigid_body_physics, create_tutorial_simulation, get_hand_open_close_qpos, prepare_tutorial_scene, @@ -87,18 +89,26 @@ def parse_arguments() -> argparse.Namespace: def create_microwave(sim) -> Articulation: """Create the fixed-base microwave articulation used by the demo.""" - microwave = sim.add_articulation( - cfg=ArticulationCfg( - uid="microwave", - fpath=get_data_path(MICROWAVE_ASSET), - init_pos=MICROWAVE_POSITION, - init_rot=MICROWAVE_ORIENTATION, - drive_pros=JointDrivePropertiesCfg( - stiffness=1e-3, damping=1e2, max_effort=1e-2 - ), - fix_base=True, - ) + microwave_cfg = ArticulationCfg( + uid="microwave", + fpath=get_data_path(MICROWAVE_ASSET), + asset_physics_mode="overlay", + init_pos=MICROWAVE_POSITION, + init_rot=MICROWAVE_ORIENTATION, + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + stiffness=1e-3, + damping=1e2, + max_effort=1e-2, + ), ) + configure_newton_link_contacts( + sim, + microwave_cfg, + group_name="newton_knob_contacts", + link_names_expr=[KNOB_LINK_NAME], + ) + microwave = sim.add_articulation(cfg=microwave_cfg) sim.update(step=10) return microwave @@ -109,6 +119,9 @@ def create_rigid_knob(sim) -> RigidObject: cfg=RigidObjectCfg( uid="rigid_knob", shape=CubeCfg(size=list(RIGID_KNOB_SIZE)), + attrs=create_tutorial_rigid_body_physics( + newton_contact=sim.is_newton_backend, + ), body_type="static", init_pos=RIGID_KNOB_POSITION, ) @@ -165,6 +178,7 @@ def main() -> None: sim, init_qpos=[0.0, -1.57, 1.57, -3.14, -1.57, 0.0, 0.0, 0.0] ) target = create_rigid_knob(sim) if args.rigid_object else create_microwave(sim) + sim.prepare() hand_open, hand_close = get_hand_open_close_qpos(robot) motion_gen = create_toppra_motion_generator(robot) semantics, target_pose = create_knob_semantics(target) diff --git a/scripts/tutorials/grasp/grasp_generator.py b/scripts/tutorials/grasp/grasp_generator.py index afea548d1..b94b2dd4d 100644 --- a/scripts/tutorials/grasp/grasp_generator.py +++ b/scripts/tutorials/grasp/grasp_generator.py @@ -30,7 +30,7 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.utility.action_utils import interpolate_with_distance -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.lab.sim.solvers import URSolverCfg from embodichain.data import get_data_path from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -39,10 +39,11 @@ from embodichain.utils import logger from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, LightCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, RigidObjectCfg, URDFCfg, ) @@ -80,8 +81,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, visualization=visualization_cfg_from_args(args), @@ -122,7 +124,7 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]) -> Robot: {"component_type": "hand", "urdf_path": gripper_urdf_path}, ] ), - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness={"Joint[0-9]": 1e4, "FINGER[1-2]": 1e3}, damping={"Joint[0-9]": 1e3, "FINGER[1-2]": 1e2}, max_effort={"Joint[0-9]": 1e5, "FINGER[1-2]": 1e4}, @@ -154,14 +156,18 @@ def create_obj(sim: SimulationManager): uid="table", shape=MeshCfg( fpath=get_resources_data_path("Model", "BakeTexture", "hdr_color_mesh.ply"), + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=16, + acd_method="coacd", + ), ), - attrs=RigidBodyAttributesCfg( - mass=0.01, - dynamic_friction=0.97, - static_friction=0.99, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 0.01}, + "material_props": {"dynamic_friction": 0.97, "static_friction": 0.99}, + } ), - max_convex_hull_num=16, - acd_method="vhacd", init_pos=[0.55, 0.0, 0.08], init_rot=[0.0, 0.0, 0.0], ) @@ -223,6 +229,7 @@ def get_grasp_traj(sim: SimulationManager, robot: Robot, grasp_xpos: torch.Tenso sim = initialize_simulation(args) robot = create_robot(sim, position=[0.0, 0.0, 0.0]) obj = create_obj(sim) + sim.prepare() # get mug grasp pose if not args.headless: diff --git a/scripts/tutorials/gym/modular_env.py b/scripts/tutorials/gym/modular_env.py index c882b258d..4a8a7eab0 100644 --- a/scripts/tutorials/gym/modular_env.py +++ b/scripts/tutorials/gym/modular_env.py @@ -36,11 +36,12 @@ from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, ArticulationCfg, RobotCfg, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) from embodichain.data import get_data_path from embodichain.utils import configclass @@ -133,11 +134,15 @@ class ExampleCfg(EmbodiedEnvCfg): fpath=get_data_path("CircleTableSimple/circle_table_simple.ply"), compute_uv=True, ), - attrs=RigidBodyAttributesCfg( - mass=10.0, - static_friction=0.95, - dynamic_friction=0.85, - restitution=0.01, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 10.0}, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.85, + "restitution": 0.01, + }, + } ), body_type="kinematic", init_pos=(0.80, 0, 0.8), @@ -196,7 +201,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), visualization=visualization_cfg_from_args(args), ), num_envs=args.num_envs, diff --git a/scripts/tutorials/gym/random_reach.py b/scripts/tutorials/gym/random_reach.py index cd30fd0a8..0b813942d 100644 --- a/scripts/tutorials/gym/random_reach.py +++ b/scripts/tutorials/gym/random_reach.py @@ -28,9 +28,11 @@ from embodichain.lab.sim.objects import RigidObject, Robot from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RobotCfg, RigidObjectCfg, - RigidBodyAttributesCfg, + CollisionPropertiesCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.gym.utils.registration import register_env @@ -48,6 +50,7 @@ def __init__( headless=False, device="cpu", renderer="hybrid", + physics_cfg="default", visualization: VisualizationCfg | None = None, **kwargs, ) -> None: @@ -55,8 +58,9 @@ def __init__( 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), visualization=visualization or VisualizationCfg(), ), num_envs=num_envs, @@ -67,12 +71,12 @@ def __init__( **kwargs, ) - def _setup_robot(self, **kwargs) -> Robot: + def _declare_robot(self, **kwargs) -> Robot: from embodichain.data import get_data_path file_path = get_data_path("UniversalRobots/UR10/UR10.urdf") - robot: Robot = self.sim.add_robot( + return self.sim.add_robot( cfg=RobotCfg( uid="ur10", fpath=file_path, @@ -81,6 +85,11 @@ def _setup_robot(self, **kwargs) -> Robot: ) ) + def _setup_robot(self, **kwargs) -> Robot: + robot = self.robot + if robot is None: + raise RuntimeError("UR10 was not declared before simulation prepare.") + qpos_limits = robot.body_data.qpos_limits[0].cpu().numpy() self.single_action_space = gym.spaces.Box( low=qpos_limits[:, 0], high=qpos_limits[:, 1], dtype=np.float32 @@ -96,7 +105,11 @@ def _prepare_scene(self, **kwargs) -> None: cfg=RigidObjectCfg( uid="cube", shape=CubeCfg(size=[size, size, size]), - attrs=RigidBodyAttributesCfg(enable_collision=False), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg( + collision_enabled=False, + ), + ), init_pos=(0.0, 0.0, 0.5), body_type="kinematic", ), @@ -137,6 +150,7 @@ def _extend_obs(self, obs: EnvObs, **kwargs) -> EnvObs: headless=args.headless, device=args.device, renderer=args.renderer, + physics_cfg=args.physics, visualization=visualization_cfg_from_args(args), ) diff --git a/scripts/tutorials/semantic_skill/hand_over.py b/scripts/tutorials/semantic_skill/hand_over.py index c2edbe686..7fe55c320 100644 --- a/scripts/tutorials/semantic_skill/hand_over.py +++ b/scripts/tutorials/semantic_skill/hand_over.py @@ -105,7 +105,7 @@ OBJECT_SIMULATION_UID = "handover_object" HANDOVER_SAMPLE_COUNT = 140 FINAL_OBJECT_POSITION = (0.0, -0.20, 0.70) -OBJECT_QUATERNION_WXYZ = (0.70710678, 0.70710678, 0.0, 0.0) +OBJECT_QUATERNION_XYZW = (0.70710678, 0.0, 0.0, 0.70710678) HANDOVER_CALL_ID = "tutorial.hand_over" HANDOVER_PRE_GRASP_DISTANCE = 0.08 HANDOVER_LIFT_HEIGHT = 0.08 @@ -156,7 +156,7 @@ def lower( final_pose = ( SemanticPose( FINAL_OBJECT_POSITION, - OBJECT_QUATERNION_WXYZ, + OBJECT_QUATERNION_XYZW, ) .to_matrix() .to(device) @@ -423,6 +423,7 @@ def main() -> None: robot = create_dual_robot(sim, args.robot) create_support_surface(sim) obj = create_handover_object(sim) + sim.prepare() settle_object(sim, obj, step=0) clone_local_pose_from_first_env(obj) obj.clear_dynamics() diff --git a/scripts/tutorials/semantic_skill/place.py b/scripts/tutorials/semantic_skill/place.py index 8df92a2fa..f255bf847 100644 --- a/scripts/tutorials/semantic_skill/place.py +++ b/scripts/tutorials/semantic_skill/place.py @@ -89,7 +89,7 @@ OBJECT_ID = "workpiece" OBJECT_SIMULATION_UID = "cube" TARGET_OBJECT_POSITION = (-0.40, 0.48, 0.025) -TARGET_OBJECT_QUATERNION_WXYZ = (1.0, 0.0, 0.0, 0.0) +TARGET_OBJECT_QUATERNION_XYZW = (0.0, 0.0, 0.0, 1.0) PICK_SAMPLE_COUNT = 120 PLACE_SAMPLE_COUNT = 120 TRAJECTORY_SIM_STEPS = 4 @@ -189,7 +189,7 @@ def create_place_task() -> tuple[Pick, Place]: object=object_ref, at=SemanticPose( TARGET_OBJECT_POSITION, - TARGET_OBJECT_QUATERNION_WXYZ, + TARGET_OBJECT_QUATERNION_XYZW, ), ), ) diff --git a/scripts/tutorials/sim/create_articulation.py b/scripts/tutorials/sim/create_articulation.py index 2b2d08129..04c947928 100644 --- a/scripts/tutorials/sim/create_articulation.py +++ b/scripts/tutorials/sim/create_articulation.py @@ -27,14 +27,25 @@ from embodichain.data import get_data_path from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import ArticulationCfg, RenderCfg +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + ArticulationRootPropertiesCfg, + DefaultRigidBodyPropertiesCfg, + JointDrivePropertiesCfg, + RenderCfg, + RigidBodyPhysicsCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.objects import Articulation from embodichain.lab.visualization import visualization_cfg_from_args DRAWER_ASSET = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" DRAWER_USER_QPOS_LIMITS = {"slide_rails": [0.0, 0.18]} -DRAWER_JOINT_FORCE = 1.0 -JOINT_LIMIT_TOLERANCE = 1.0e-3 +DRAWER_JOINT_FORCE_LIMIT = 1.0 +DRAWER_POSITION_GAIN = 20.0 +DRAWER_VELOCITY_GAIN = 4.0 +JOINT_POSITION_TOLERANCE = 1.0e-3 +JOINT_VELOCITY_TOLERANCE = 1.0e-2 def create_articulation(sim: SimulationManager) -> Articulation: @@ -49,19 +60,30 @@ def create_articulation(sim: SimulationManager) -> Articulation: Raises: RuntimeError: If the constructed backend joints are not passive. """ - # Resolve the drawer URDF and configure its initial pose. ``drive_pros`` is - # intentionally omitted: ArticulationCfg defaults to drive_type="none". + # Resolve the drawer URDF and explicitly request the passive drive used by + # this tutorial while retaining all unconfigured asset properties. articulation_cfg = ArticulationCfg( uid="drawer", fpath=get_data_path(DRAWER_ASSET), + asset_physics_mode="overlay", init_pos=(0.0, 0.0, 0.05), - fix_base=True, + root_props=ArticulationRootPropertiesCfg(fixed_base=True), + joint_drive_props=JointDrivePropertiesCfg(drive_type="none"), # The asset limit is [0.0, 0.2]; keep 90% of its travel range. qpos_limits=DRAWER_USER_QPOS_LIMITS, + # Newton currently has no body-level damping setting. Remove the + # Default backend's damping so both passive models use zero damping. + attrs=RigidBodyPhysicsCfg( + rigid_props=DefaultRigidBodyPropertiesCfg( + linear_damping=0.0, + angular_damping=0.0, + ) + ), ) # Load one articulation instance into every simulation environment. articulation: Articulation = sim.add_articulation(cfg=articulation_cfg) + sim.prepare() # Query the constructed DexSim entities, not only the config object. backend_drive_types = articulation.get_joint_drive_type() @@ -77,7 +99,7 @@ def create_articulation(sim: SimulationManager) -> Articulation: print(f"[INFO]: Loaded articulation with {articulation.dof} joint(s)", flush=True) print(f"[INFO]: Joint names: {articulation.joint_names}", flush=True) print( - f"[INFO]: Config drive type: {articulation.cfg.drive_pros.drive_type}", + f"[INFO]: Config drive type: {articulation.cfg.joint_drive_props.drive_type}", flush=True, ) print(f"[INFO]: Backend drive types: {backend_drive_types}", flush=True) @@ -87,15 +109,26 @@ def create_articulation(sim: SimulationManager) -> Articulation: return articulation -def apply_drawer_force(articulation: Articulation, opening: bool) -> None: - """Apply a joint force that opens or closes the drawer. +def apply_drawer_force( + articulation: Articulation, + target_qpos: torch.Tensor, +) -> None: + """Apply effort-limited PD control toward a drawer position. Args: articulation: Drawer articulation receiving the force. - opening: If True, apply positive force; otherwise apply negative force. + target_qpos: Target joint positions for every environment and joint. """ - force = DRAWER_JOINT_FORCE if opening else -DRAWER_JOINT_FORCE - joint_forces = torch.full_like(articulation.get_qpos(), force) + position_error = target_qpos - articulation.get_qpos() + joint_forces = ( + DRAWER_POSITION_GAIN * position_error + - DRAWER_VELOCITY_GAIN * articulation.get_qvel() + ) + joint_forces = torch.clamp( + joint_forces, + min=-DRAWER_JOINT_FORCE_LIMIT, + max=DRAWER_JOINT_FORCE_LIMIT, + ) articulation.set_qf(joint_forces) @@ -104,47 +137,48 @@ def run_simulation( articulation: Articulation, max_steps: int | None = None, ) -> None: - """Open and close the drawer by reversing force at its joint limits. + """Open and close the drawer with effort-limited position tracking. Args: sim: Simulation manager to advance. articulation: Drawer articulation whose joints are updated. max_steps: Optional number of steps to run before returning. """ - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - qpos_limits = articulation.get_qpos_limits() closed_qpos = qpos_limits[..., 0] open_qpos = qpos_limits[..., 1] opening = True + target_qpos = open_qpos step_count = 0 print( - f"[INFO]: Applying +{DRAWER_JOINT_FORCE:.1f} N to open the drawer", + "[INFO]: Tracking the open position with joint effort limited to " + f"+/-{DRAWER_JOINT_FORCE_LIMIT:.1f} N", flush=True, ) try: while max_steps is None or step_count < max_steps: qpos = articulation.get_qpos() - if opening and torch.all(qpos >= open_qpos - JOINT_LIMIT_TOLERANCE).item(): - print(f"[INFO]: Drawer reached open limit: {qpos}", flush=True) - opening = False + qvel = articulation.get_qvel() + settled = torch.all( + (torch.abs(qpos - target_qpos) <= JOINT_POSITION_TOLERANCE) + & (torch.abs(qvel) <= JOINT_VELOCITY_TOLERANCE) + ).item() + if settled: + reached_position = "open" if opening else "closed" print( - f"[INFO]: Applying -{DRAWER_JOINT_FORCE:.1f} N to close the drawer", + f"[INFO]: Drawer settled at {reached_position} position: " + f"qpos={qpos}, qvel={qvel}", flush=True, ) - elif ( - not opening - and torch.all(qpos <= closed_qpos + JOINT_LIMIT_TOLERANCE).item() - ): - print(f"[INFO]: Drawer reached closed limit: {qpos}", flush=True) - opening = True + opening = not opening + target_qpos = open_qpos if opening else closed_qpos + target_position = "open" if opening else "closed" print( - f"[INFO]: Applying +{DRAWER_JOINT_FORCE:.1f} N to open the drawer", + f"[INFO]: Tracking the {target_position} position", flush=True, ) - apply_drawer_force(articulation, opening=opening) + apply_drawer_force(articulation, target_qpos=target_qpos) sim.update(step=1) step_count += 1 except KeyboardInterrupt: @@ -169,13 +203,17 @@ def main() -> None: if args.max_steps is not None and args.max_steps < 1: parser.error("--max-steps must be at least 1") - # Configure the simulation. Window creation is deferred until the asset is loaded. + open_native_window = not args.headless and not args.viser + + # Construct the World without a window so Spawn can finish first. The + # requested native window is opened explicitly after create_articulation(). sim_cfg = SimulationManagerCfg( - headless=args.headless, + headless=True, sim_device=args.device, num_envs=args.num_envs, arena_space=2.0, physics_dt=1.0 / 100.0, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg(renderer=args.renderer), visualization=visualization_cfg_from_args(args), ) @@ -185,14 +223,17 @@ def main() -> None: articulation = create_articulation(sim) print(f"[INFO]: Initial joint positions: {articulation.get_qpos()}", flush=True) - if not args.headless and not args.viser: + if open_native_window: sim.open_window() print("[INFO]: Running simulation. Press Ctrl+C to stop.", flush=True) run_simulation(sim, articulation, max_steps=args.max_steps) finally: - sim.destroy() + sim.destroy(exit_process=False) if __name__ == "__main__": - main() + try: + main() + finally: + SimulationManager.flush_cleanup_queue() diff --git a/scripts/tutorials/sim/create_cloth.py b/scripts/tutorials/sim/create_cloth.py index 405f4f89e..d9de77642 100644 --- a/scripts/tutorials/sim/create_cloth.py +++ b/scripts/tutorials/sim/create_cloth.py @@ -33,8 +33,9 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ClothObjectCfg, ClothPhysicalAttributesCfg, ) @@ -93,8 +94,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), visualization=visualization_cfg_from_args(args), ) @@ -121,7 +123,7 @@ def main(): mass=0.01, youngs=1e9, poissons=0.4, - thickness=0.04, + thickness=0.004, bending_stiffness=0.01, bending_damping=0.1, dynamic_friction=0.95, @@ -134,13 +136,16 @@ def main(): shape=CubeCfg( size=[0.1, 0.1, 0.06], ), - attrs=RigidBodyAttributesCfg( - mass=1.0, - static_friction=0.95, - dynamic_friction=0.9, - restitution=0.01, - min_position_iters=32, - min_velocity_iters=8, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "rigid_props": {"min_position_iters": 32, "min_velocity_iters": 8}, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + "restitution": 0.01, + }, + } ), body_type="dynamic", init_pos=[0.5, 0.0, 0.04], @@ -149,6 +154,8 @@ def main(): padding_box = sim.add_rigid_object(cfg=padding_box_cfg) print("[INFO]: Add soft object complete!") + sim.prepare() + # Open window when the scene has been set up if not args.headless: sim.open_window() @@ -168,9 +175,6 @@ def run_simulation(sim: SimulationManager, cloth: ClothObject) -> None: soft_obj: soft object """ - # Initialize GPU physics - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/create_rigid_constraint.py b/scripts/tutorials/sim/create_rigid_constraint.py index 6a7a629e9..617b074bc 100644 --- a/scripts/tutorials/sim/create_rigid_constraint.py +++ b/scripts/tutorials/sim/create_rigid_constraint.py @@ -30,7 +30,7 @@ from embodichain.lab.sim.cfg import ( RigidObjectCfg, RigidConstraintCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, RenderCfg, ) from embodichain.lab.sim.shapes import CubeCfg @@ -64,7 +64,7 @@ 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), num_envs=args.num_envs, arena_space=3.0, @@ -74,11 +74,15 @@ def main(): sim = SimulationManager(sim_cfg) # Shared physics attributes for the two cubes. - physics_attrs = RigidBodyAttributesCfg( - mass=0.2, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + physics_attrs = RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 0.2}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.1, + }, + } ) # Add two dynamic cubes to the scene. cube_a starts higher than cube_b so @@ -101,8 +105,7 @@ def main(): ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() print("[INFO]: Scene setup complete with two cubes (cube_a, cube_b).") diff --git a/scripts/tutorials/sim/create_rigid_object_group.py b/scripts/tutorials/sim/create_rigid_object_group.py index 08f3aac73..ac8c6cd17 100644 --- a/scripts/tutorials/sim/create_rigid_object_group.py +++ b/scripts/tutorials/sim/create_rigid_object_group.py @@ -25,8 +25,12 @@ 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 ( + RigidBodyPhysicsCfg, + RenderCfg, + physics_cfg_for_backend, +) from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg from embodichain.lab.sim.shapes import CubeCfg from embodichain.lab.sim.objects import ( RigidObjectGroup, @@ -51,10 +55,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, visualization=visualization_cfg_from_args(args), @@ -63,11 +68,15 @@ def main(): # Create the simulation instance sim = SimulationManager(sim_cfg) - physics_attrs = RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + physics_attrs = RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 1.0}, + "material_props": { + "dynamic_friction": 0.5, + "static_friction": 0.5, + "restitution": 0.1, + }, + } ) # Add objects to the scene @@ -102,6 +111,7 @@ def main(): print("[INFO]: Press Ctrl+C to stop the simulation") # Open window when the scene has been set up + sim.prepare() if not args.headless: sim.open_window() @@ -116,10 +126,6 @@ def run_simulation(sim: SimulationManager): sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/create_robot.py b/scripts/tutorials/sim/create_robot.py index e4dd591df..9106f3798 100644 --- a/scripts/tutorials/sim/create_robot.py +++ b/scripts/tutorials/sim/create_robot.py @@ -35,6 +35,7 @@ from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, URDFCfg, @@ -54,15 +55,24 @@ def main(): description="Create and simulate a robot in SimulationManager" ) add_env_launcher_args_to_parser(parser) + parser.add_argument( + "--max-steps", + type=int, + default=None, + help="Stop after this many physics steps (default: run until interrupted).", + ) args = parser.parse_args() + if args.max_steps is not None and args.max_steps < 1: + parser.error("--max-steps must be at least 1") # Initialize simulation 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, visualization=visualization_cfg_from_args(args), @@ -72,16 +82,16 @@ def main(): # Create robot configuration robot = create_robot(sim) - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + # Materialize the declared scene before accessing robot metadata. + sim.prepare() + print(f"Robot created successfully with {robot.dof} joints") # Open visualization window if not headless if not args.headless: sim.open_window() # Run simulation loop - run_simulation(sim, robot) + run_simulation(sim, robot, max_steps=args.max_steps) def create_robot(sim): @@ -127,21 +137,45 @@ def create_robot(sim): ] ), control_parts=CONTROL_PARTS, - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", stiffness={"joint[1-6]": 1e4, "LEFT_.*": 1e3}, - damping={"joint[1-6]": 1e3, "LEFT_.*": 1e2}, + damping={"joint[1-6]": 1.5e3, "LEFT_.*": 1e2}, + max_effort={"joint[1-6]": 1e4, "LEFT_.*": 1e4}, ), ) # Add robot to simulation robot: Robot = sim.add_robot(cfg=cfg) - print(f"Robot created successfully with {robot.dof} joints") - return robot -def run_simulation(sim: SimulationManager, robot: Robot): +def _expand_mimic_targets( + robot: Robot, joint_ids: list[int], joint_targets: torch.Tensor +) -> torch.Tensor: + """Expand active-joint targets into mimic-consistent articulation targets.""" + + targets = robot.get_qpos(target=True).clone() + targets[:, joint_ids] = joint_targets + + for mimic_id, parent_id, multiplier, offset in zip( + robot.mimic_ids, + robot.mimic_parents, + robot.mimic_multipliers, + robot.mimic_offsets, + ): + if mimic_id is None or parent_id is None: + continue + targets[:, mimic_id] = offset + multiplier * targets[:, parent_id] + + limits = robot.body_data.qpos_limits + return targets.clamp(min=limits[..., 0], max=limits[..., 1]) + + +def run_simulation( + sim: SimulationManager, robot: Robot, max_steps: int | None = None +) -> None: """Run the simulation loop with robot control.""" print("Starting simulation...") @@ -170,14 +204,29 @@ def run_simulation(sim: SimulationManager, robot: Robot): # Get joint IDs for the hand. hand_joint_ids = robot.get_joint_ids("hand") - # Define hand open and close positions based on joint limits. - hand_position_open = robot.body_data.qpos_limits[:, hand_joint_ids, 1] - hand_position_close = robot.body_data.qpos_limits[:, hand_joint_ids, 0] + active_hand_joint_ids = robot.get_joint_ids("hand", remove_mimic=True) + # Drive mimic joints toward the pose implied by their active parent instead of + # sending each joint to its independent limit. Newton keeps drives on mimic + # joints, so inconsistent targets otherwise compete with the mimic constraints. + hand_position_open = _expand_mimic_targets( + robot, + active_hand_joint_ids, + robot.body_data.qpos_limits[:, active_hand_joint_ids, 1], + )[:, hand_joint_ids] + hand_position_close = _expand_mimic_targets( + robot, + active_hand_joint_ids, + robot.body_data.qpos_limits[:, active_hand_joint_ids, 0], + )[:, hand_joint_ids] + + # The reset pose is zero for every DOF, but this hand has non-zero mimic + # offsets. Start from a valid closed pose so the initial state and drive + # targets satisfy the same mimic equations. + robot.set_qpos(qpos=hand_position_close, joint_ids=hand_joint_ids, target=False) + robot.set_qpos(qpos=hand_position_close, joint_ids=hand_joint_ids) try: - while True: - # Update physics - sim.update(step=1) + while max_steps is None or step_count < max_steps: cycle_step = step_count % ACTION_CYCLE_STEPS if cycle_step == 0: @@ -196,6 +245,9 @@ def run_simulation(sim: SimulationManager, robot: Robot): robot.set_qpos(qpos=hand_position_open, joint_ids=hand_joint_ids) print(f"Opening hand") + # Apply commands before advancing physics so both backends observe the + # target change on the same simulation step. + sim.update(step=1) step_count += 1 except KeyboardInterrupt: diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index fa4d82ea5..e66f9506b 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -25,8 +25,14 @@ import time from embodichain.lab.sim import SimulationManager, SimulationManagerCfg -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg, RenderCfg -from embodichain.lab.sim.shapes import CubeCfg, MeshCfg +from embodichain.lab.sim.cfg import ( + MassPropertiesCfg, + RenderCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + physics_cfg_for_backend, +) +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, MeshCollisionCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.visualization import visualization_cfg_from_args @@ -42,7 +48,7 @@ def main() -> None: ) add_env_launcher_args_to_parser(parser) parser.add_argument( - "--record-steps", + "--max_steps", type=int, default=1000, help=( @@ -68,7 +74,8 @@ def main() -> None: height=1080, headless=True, physics_dt=1.0 / 100.0, # Physics timestep (100 Hz) - sim_device=args.device, + device=args.device, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg( renderer=args.renderer, ), @@ -86,11 +93,13 @@ def main() -> None: uid="cube", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="dynamic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=0.1), + material_props=RigidBodyMaterialCfg( + dynamic_friction=0.5, + static_friction=0.5, + restitution=0.1, + ), ), init_pos=[0, 0.0, 1.0], ) @@ -101,17 +110,26 @@ def main() -> None: chair: RigidObject = sim.add_rigid_object( cfg=RigidObjectCfg( uid="chair", - shape=MeshCfg(fpath=path), + shape=MeshCfg( + fpath=path, + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=32, + ), + ), body_type="dynamic", - attrs=RigidBodyAttributesCfg( - mass=3.0, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(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], ) ) + # Materialize the complete initial scene before exposing it to the viewer. + sim.prepare() + print("[INFO]: Scene setup complete!") print(f"[INFO]: Running simulation with {args.num_envs} environment(s)") print("[INFO]: Press Ctrl+C to stop the simulation") @@ -133,13 +151,10 @@ def main() -> None: 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( @@ -153,10 +168,6 @@ def run_simulation( max_steps: Optional maximum number of simulation steps to execute. """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 try: @@ -167,6 +178,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 60fa82c3e..42a0aa482 100644 --- a/scripts/tutorials/sim/create_sensor.py +++ b/scripts/tutorials/sim/create_sensor.py @@ -37,6 +37,7 @@ from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, JointDrivePropertiesCfg, RobotCfg, URDFCfg, @@ -98,9 +99,10 @@ def main() -> None: 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, visualization=visualization_cfg_from_args(args), @@ -110,8 +112,6 @@ def main() -> None: # Create robot configuration robot = create_robot(sim) - sensor = create_sensor(sim, args) - # Add a cube to the scene cube_cfg = RigidObjectCfg( uid="cube", @@ -121,9 +121,12 @@ def main() -> None: ) sim.add_rigid_object(cfg=cube_cfg) - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + # Materialize all physical assets before reading robot metadata or + # constructing render-only sensors. + sim.prepare() + print(f"Robot created successfully with {robot.dof} joints") + + sensor = create_sensor(sim, args) # Open visualization window if not headless if not args.headless: @@ -147,7 +150,8 @@ def create_sensor(sim: SimulationManager, args): # extrinsics params pos = [0.09, 0.05, 0.04] - quat = R.from_euler("xyz", [-35, 135, 0], degrees=True).as_quat().tolist() + # CameraCfg uses xyzw; this rotation preserves the intended wrist-camera view. + quat = R.from_euler("xyz", [180, -45, 35], degrees=True).as_quat().tolist() # If attach_sensor is True, attach to robot end-effector; otherwise, place it in the scene if args.attach_sensor: @@ -156,7 +160,6 @@ def create_sensor(sim: SimulationManager, args): parent = None pos = [1.2, -0.2, 1.5] quat = R.from_euler("xyz", [0, 180, 0], degrees=True).as_quat().tolist() - quat = [quat[3], quat[0], quat[1], quat[2]] # Convert to (w, x, y, z) # create camera sensor and attach to robot end-effector camera: Camera = sim.add_sensor( @@ -223,17 +226,17 @@ def create_robot(sim): ] ), control_parts=CONTROL_PARTS, - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", stiffness={"joint[1-6]": 1e4, "LEFT_.*": 1e3}, - damping={"joint[1-6]": 1e3, "LEFT_.*": 1e2}, + damping={"joint[1-6]": 1.5e3, "LEFT_.*": 1e2}, + max_effort={"joint[1-6]": 1e4, "LEFT_.*": 1e4}, ), ) # Add robot to simulation robot: Robot = sim.add_robot(cfg=cfg) - print(f"Robot created successfully with {robot.dof} joints") - return robot diff --git a/scripts/tutorials/sim/create_softbody.py b/scripts/tutorials/sim/create_softbody.py index 83b15b662..aab5b4112 100644 --- a/scripts/tutorials/sim/create_softbody.py +++ b/scripts/tutorials/sim/create_softbody.py @@ -29,6 +29,7 @@ from embodichain.lab.visualization import visualization_cfg_from_args from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, SoftbodyVoxelAttributesCfg, SoftbodyPhysicalAttributesCfg, ) @@ -56,10 +57,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), visualization=visualization_cfg_from_args(args), ) @@ -91,6 +93,8 @@ def main(): ) print("[INFO]: Add soft object complete!") + sim.prepare() + # Open window when the scene has been set up if not args.headless: sim.open_window() @@ -110,9 +114,6 @@ def run_simulation(sim: SimulationManager, soft_obj: SoftObject) -> None: soft_obj: soft object """ - # Initialize GPU physics - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/export_usd.py b/scripts/tutorials/sim/export_usd.py index a8c7de46b..a96f7aede 100644 --- a/scripts/tutorials/sim/export_usd.py +++ b/scripts/tutorials/sim/export_usd.py @@ -28,13 +28,14 @@ from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.cfg import ( RenderCfg, + physics_cfg_for_backend, LightCfg, JointDrivePropertiesCfg, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ArticulationCfg, ) -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.data import get_data_path from embodichain.utils import logger @@ -67,8 +68,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, @@ -180,11 +182,12 @@ def create_table(sim: SimulationManager) -> RigidObject: uid="table", shape=MeshCfg( fpath=get_data_path("MultiW1Data/table_a.obj"), + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=8, + ), ), - attrs=RigidBodyAttributesCfg( - mass=0.5, - ), - max_convex_hull_num=8, + attrs=RigidBodyPhysicsCfg.from_dict({"mass_props": {"mass": 0.5}}), body_type="kinematic", init_pos=[1.1, -0.5, 0.08], init_rot=[0.0, 0.0, 0.0], @@ -206,12 +209,11 @@ def create_caffe(sim: SimulationManager) -> Robot: container_cfg = ArticulationCfg( uid="caffe", fpath=get_data_path("MultiW1Data/cafe/cafe.urdf"), + asset_physics_mode="overlay", init_pos=[1.05, -0.5, 0.79], init_rot=[0, 0, -30], - attrs=RigidBodyAttributesCfg( - mass=1.0, - ), - drive_pros=JointDrivePropertiesCfg( + attrs=RigidBodyPhysicsCfg.from_dict({"mass_props": {"mass": 1.0}}), + joint_drive_props=JointDrivePropertiesCfg( stiffness=1.0, damping=0.1, max_effort=100.0, drive_type="force" ), ) @@ -235,10 +237,7 @@ def create_cup(sim: SimulationManager) -> RigidObject: shape=MeshCfg( fpath=get_data_path("MultiW1Data/paper_cup_2.obj"), ), - attrs=RigidBodyAttributesCfg( - mass=0.3, - ), - max_convex_hull_num=1, + attrs=RigidBodyPhysicsCfg.from_dict({"mass_props": {"mass": 0.3}}), body_type="dynamic", init_pos=[0.86, -0.76, 0.841], init_rot=[0.0, 0.0, 0.0], @@ -261,7 +260,9 @@ def main(): caffe = create_caffe(sim) cup = create_cup(sim) - sim.export_usd("w1_coffee_scene.usda") + sim.prepare() + + sim.export_usd("w1_coffee_scene.usd") logger.log_info("Scene exported successfully.") diff --git a/scripts/tutorials/sim/gizmo_robot.py b/scripts/tutorials/sim/gizmo_robot.py index 3c04f4eb4..c5c2ec453 100644 --- a/scripts/tutorials/sim/gizmo_robot.py +++ b/scripts/tutorials/sim/gizmo_robot.py @@ -29,6 +29,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, @@ -55,8 +56,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), visualization=visualization_cfg_from_args(args), ) @@ -84,7 +86,8 @@ def main(): dt=0.1, ) }, - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", stiffness={"Joint[1-6]": 1e4}, damping={"Joint[1-6]": 1e3}, ), @@ -97,6 +100,8 @@ def main(): dtype=torch.float32, device="cpu", ) + + sim.prepare() joint_ids = robot.get_joint_ids("arm") robot.set_qpos(qpos=initial_qpos, joint_ids=joint_ids) diff --git a/scripts/tutorials/sim/import_usd.py b/scripts/tutorials/sim/import_usd.py index b30639bb1..968840ff2 100644 --- a/scripts/tutorials/sim/import_usd.py +++ b/scripts/tutorials/sim/import_usd.py @@ -28,7 +28,13 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.visualization import visualization_cfg_from_args 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 ( + MassPropertiesCfg, + RenderCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.objects import ( RigidObject, @@ -55,10 +61,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, visualization=visualization_cfg_from_args(args), @@ -72,11 +79,13 @@ def main(): uid="cube", shape=CubeCfg(size=[0.1, 0.1, 0.1]), body_type="dynamic", - attrs=RigidBodyAttributesCfg( - mass=1.0, - dynamic_friction=0.5, - static_friction=0.5, - restitution=0.1, + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg( + dynamic_friction=0.5, + static_friction=0.5, + restitution=0.1, + ), ), init_pos=[0.0, 0.0, 1.0], ) @@ -90,7 +99,7 @@ def main(): shape=MeshCfg(fpath=sugar_box_path), body_type="dynamic", init_pos=[0.2, 0.2, 1.0], - use_usd_properties=True, + asset_physics_mode="preserve", ) ) @@ -103,11 +112,12 @@ def main(): fpath=h1_path, build_pk_chain=False, init_pos=[-0.2, -0.2, 1.05], - use_usd_properties=False, + asset_physics_mode="overlay", ) ) # Open window when the scene has been set up + sim.prepare() if not args.headless: sim.open_window() @@ -125,10 +135,6 @@ def run_simulation(sim: SimulationManager): sim: The SimulationManager instance to run """ - # Initialize GPU physics if using CUDA - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - step_count = 0 try: diff --git a/scripts/tutorials/sim/motion_generator.py b/scripts/tutorials/sim/motion_generator.py index 776ca9dec..351b00009 100644 --- a/scripts/tutorials/sim/motion_generator.py +++ b/scripts/tutorials/sim/motion_generator.py @@ -25,8 +25,8 @@ from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import RenderCfg, physics_cfg_for_backend from embodichain.lab.visualization import visualization_cfg_from_args -from embodichain.lab.sim.cfg import RenderCfg from embodichain.lab.sim.objects import Robot from embodichain.lab.sim.planners import ( MotionGenCfg, @@ -226,7 +226,8 @@ def main() -> None: height=RECORD_HEIGHT, headless=True, physics_dt=1.0 / 100.0, - sim_device=args.device, + device=args.device, + physics_cfg=physics_cfg_for_backend(args.physics), render_cfg=RenderCfg(renderer=args.renderer), num_envs=args.num_envs, arena_space=args.arena_space, @@ -237,8 +238,7 @@ def main() -> None: robot: Robot = sim.add_robot(cfg=CobotMagicCfg.from_dict({"uid": "CobotMagic"})) arm_name = "left_arm" - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() if not args.headless: sim.open_window() diff --git a/scripts/tutorials/sim/open_drawer.py b/scripts/tutorials/sim/open_drawer.py index 9e9e9ec3b..13bccc57a 100644 --- a/scripts/tutorials/sim/open_drawer.py +++ b/scripts/tutorials/sim/open_drawer.py @@ -20,6 +20,7 @@ import argparse from collections.abc import Sequence +from typing import Literal import torch @@ -28,9 +29,15 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( ArticulationCfg, + ArticulationRootPropertiesCfg, JointDrivePropertiesCfg, + LinkPhysicsOverrideCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + PhysicsBackendCfg, RenderCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, + physics_cfg_for_backend, ) from embodichain.lab.sim.objects import Articulation, Robot from embodichain.lab.sim.planners import ( @@ -59,10 +66,15 @@ ARM_NAME = "arm" HAND_NAME = "hand" HANDLE_LINK_NAME = "handle_xpos" +DRAWER_CONTACT_LINK_NAME = "inner_box" +LEFT_FINGER_LINK_NAME = "fr3_leftfinger" +RIGHT_FINGER_LINK_NAME = "fr3_rightfinger" DRAWER_ASSET = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" APPROACH_DISTANCE = 0.10 PULL_DISTANCE = 0.16 +NEWTON_GRASP_CONTACT_STIFFNESS = 4.0e4 +NEWTON_GRASP_CONTACT_DAMPING = 4.0e2 DRAWER_SUCCESS_THRESHOLD = 0.10 HALF_OPEN_FRACTION = 0.5 HALF_OPEN_TOLERANCE = 0.02 @@ -75,6 +87,21 @@ ) +def _newton_grasp_contact_override( + link_names_expr: str, +) -> LinkPhysicsOverrideCfg: + """Build the Newton contact material used at the drawer grasp.""" + return LinkPhysicsOverrideCfg( + link_names_expr=[link_names_expr], + attrs=RigidBodyPhysicsCfg( + material_props=NewtonRigidBodyMaterialCfg( + ke=NEWTON_GRASP_CONTACT_STIFFNESS, + kd=NEWTON_GRASP_CONTACT_DAMPING, + ) + ), + ) + + def create_scene(sim: SimulationManager) -> tuple[Robot, Articulation]: """Add a Franka Panda and a passive sliding drawer to the scene. @@ -94,31 +121,46 @@ def create_scene(sim: SimulationManager) -> tuple[Robot, Articulation]: "uid": "tutorial_franka", "robot_type": "panda", "attrs": { - "static_friction": 1.0, - "dynamic_friction": 1.0, + "material_props": { + "static_friction": 1.0, + "dynamic_friction": 1.0, + }, }, } ) + if sim.is_newton_backend: + robot_cfg.link_attrs = { + **(robot_cfg.link_attrs or {}), + "newton_gripper_contacts": _newton_grasp_contact_override( + f"(?:{LEFT_FINGER_LINK_NAME}|{RIGHT_FINGER_LINK_NAME})" + ), + } robot = sim.add_robot(cfg=robot_cfg) if robot is None: raise RuntimeError("Failed to add the Franka Panda robot.") # Keep the drawer base fixed while leaving its prismatic joint passive. The # 180-degree yaw makes the drawer's opening direction point toward Franka. - drawer = sim.add_articulation( - cfg=ArticulationCfg( - uid="drawer", - fpath=get_data_path(DRAWER_ASSET), - init_pos=(0.72, 0.0, 0.42), - init_rot=(0.0, 0.0, 180.0), - fix_base=True, - drive_pros=JointDrivePropertiesCfg(drive_type="none"), - attrs=RigidBodyAttributesCfg( - static_friction=1.0, - dynamic_friction=1.0, - ), - ) + drawer_cfg = ArticulationCfg( + uid="drawer", + fpath=get_data_path(DRAWER_ASSET), + asset_physics_mode="overlay", + init_pos=(0.72, 0.0, 0.42), + init_rot=(0.0, 0.0, 180.0), + root_props=ArticulationRootPropertiesCfg(fixed_base=True), + joint_drive_props=JointDrivePropertiesCfg(drive_type="none"), + attrs=RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": 1.0, "dynamic_friction": 1.0}} + ), ) + if sim.is_newton_backend: + # The handle marker has no geometry; its collision belongs to inner_box. + drawer_cfg.link_attrs = { + "newton_handle_contacts": _newton_grasp_contact_override( + DRAWER_CONTACT_LINK_NAME + ) + } + drawer = sim.add_articulation(cfg=drawer_cfg) return robot, drawer @@ -228,6 +270,34 @@ def play_arm_trajectory( sim.update(step=physics_steps_per_waypoint) +def _move_arm_to_poses( + sim: SimulationManager, + robot: Robot, + motion_generator: MotionGenerator, + target_poses: Sequence[torch.Tensor], + *, + sample_count: int, + physics_steps_per_waypoint: int = 4, + wait_for_input: bool = False, +) -> None: + """Plan and execute arm motion through Cartesian target poses.""" + start_qpos = robot.get_qpos(name=ARM_NAME) + trajectory = generate_arm_trajectory( + motion_generator, + qpos_waypoints=solve_ik_waypoints(robot, target_poses, start_qpos), + start_qpos=start_qpos, + sample_count=sample_count, + ) + if wait_for_input: + input("[READY]: Trajectory planned. Press Enter to start execution...") + play_arm_trajectory( + sim, + robot, + trajectory, + physics_steps_per_waypoint=physics_steps_per_waypoint, + ) + + def move_gripper( sim: SimulationManager, robot: Robot, @@ -322,21 +392,14 @@ def open_drawer( approach_pose = grasp_pose.clone() approach_pose[:, :3, 3] -= grasp_pose[:, :3, 2] * APPROACH_DISTANCE - start_qpos = robot.get_qpos(name=ARM_NAME) - approach_waypoints = solve_ik_waypoints( + _move_arm_to_poses( + sim, robot, - target_poses=[approach_pose, grasp_pose], - start_qpos=start_qpos, - ) - approach_trajectory = generate_arm_trajectory( motion_generator, - qpos_waypoints=approach_waypoints, - start_qpos=start_qpos, + [approach_pose, grasp_pose], sample_count=60, + wait_for_input=wait_for_input, ) - if wait_for_input: - input("[READY]: Trajectory planned. Press Enter to start execution...") - play_arm_trajectory(sim, robot, approach_trajectory) # Close around the handle, then allow contacts to settle before pulling. move_gripper(sim, robot, hand_closed_qpos) @@ -348,22 +411,12 @@ def open_drawer( pull_pose = grasped_handle_pose.clone() pull_pose[:, :3, 3] -= grasped_handle_pose[:, :3, 2] * PULL_DISTANCE - pull_start_qpos = robot.get_qpos(name=ARM_NAME) - pull_waypoints = solve_ik_waypoints( + _move_arm_to_poses( + sim, robot, - target_poses=[pull_pose], - start_qpos=pull_start_qpos, - ) - pull_trajectory = generate_arm_trajectory( motion_generator, - qpos_waypoints=pull_waypoints, - start_qpos=pull_start_qpos, + [pull_pose], sample_count=80, - ) - play_arm_trajectory( - sim, - robot, - pull_trajectory, physics_steps_per_waypoint=5, ) sim.update(step=50) @@ -388,36 +441,27 @@ def open_drawer( push_pose = pushed_handle_pose.clone() push_pose[:, :3, 3] += pushed_handle_pose[:, :3, 2] * push_distance.unsqueeze(-1) - push_start_qpos = robot.get_qpos(name=ARM_NAME) - push_waypoints = solve_ik_waypoints( + _move_arm_to_poses( + sim, robot, - target_poses=[push_pose], - start_qpos=push_start_qpos, - ) - push_trajectory = generate_arm_trajectory( motion_generator, - qpos_waypoints=push_waypoints, - start_qpos=push_start_qpos, + [push_pose], sample_count=50, - ) - play_arm_trajectory( - sim, - robot, - push_trajectory, physics_steps_per_waypoint=5, ) sim.update(step=50) drawer_qpos = drawer.get_qpos() final_opening = drawer_qpos[:, 0] + half_open_error = torch.abs(final_opening - half_open_target) print( - "[INFO]: Drawer opening after half push (m): " - f"{final_opening.detach().cpu().tolist()}", + "[INFO]: Drawer opening after half push (m): final=" + f"{final_opening.detach().cpu().tolist()}, target=" + f"{half_open_target.detach().cpu().tolist()}, abs_error=" + f"{half_open_error.detach().cpu().tolist()}", flush=True, ) - if not torch.all( - torch.abs(final_opening - half_open_target) <= HALF_OPEN_TOLERANCE - ).item(): + if not torch.all(half_open_error <= HALF_OPEN_TOLERANCE).item(): raise RuntimeError( "The drawer did not return to half of its pulled opening. " f"Expected an error no greater than {HALF_OPEN_TOLERANCE:.2f} m." @@ -425,6 +469,24 @@ def open_drawer( return drawer_qpos +def _tutorial_physics_cfg( + backend: Literal["default", "newton"], +) -> PhysicsBackendCfg: + """Build the physics configuration used by this tutorial.""" + physics_cfg = physics_cfg_for_backend(backend) + if isinstance(physics_cfg, NewtonPhysicsCfg): + # Use a finer step and multi-point contacts for the Newton grasp. + physics_cfg.num_substeps = 20 + physics_cfg.solver_cfg = { + "solver_type": "mujoco_warp", + "njmax": 8192, + "nconmax": 8192, + "cone": "elliptic", + "enable_multiccd": True, + } + return physics_cfg + + def main() -> None: """Run the Franka drawer-manipulation tutorial.""" parser = argparse.ArgumentParser( @@ -464,15 +526,23 @@ def main() -> None: if args.record_save_path is not None and not args.headless: parser.error("--record-save-path requires --headless") + open_native_window = not args.headless and not args.viser + + # PytorchSolver samples multiple IK seeds; make the tutorial trajectory + # reproducible across repeated runs of the same backend. + torch.manual_seed(0) + + # Construct the World without a window so Spawn can finish first. sim = SimulationManager( SimulationManagerCfg( width=RECORD_WIDTH, height=RECORD_HEIGHT, - headless=args.headless, + headless=True, sim_device=args.device, num_envs=args.num_envs, arena_space=args.arena_space, physics_dt=1.0 / 100.0, + physics_cfg=_tutorial_physics_cfg(args.physics), render_cfg=RenderCfg(renderer=args.renderer), visualization=visualization_cfg_from_args(args), ) @@ -481,9 +551,8 @@ def main() -> None: try: robot, drawer = create_scene(sim) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() - if not args.headless and not args.viser: + sim.prepare() + if open_native_window: sim.open_window() sim.update(step=5) diff --git a/scripts/tutorials/sim/srs_solver.py b/scripts/tutorials/sim/srs_solver.py index 4be55f561..59c5c7146 100644 --- a/scripts/tutorials/sim/srs_solver.py +++ b/scripts/tutorials/sim/srs_solver.py @@ -73,7 +73,7 @@ def main( # Keep the native window closed while planning so renderer/window # lifecycle events cannot terminate or perturb timed CUDA IK calls. headless=True, - sim_device=device, + device=device, width=2200, height=1200, visualization=visualization or VisualizationCfg(), @@ -93,6 +93,7 @@ def main( [2.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0] ) robot: Robot = sim.add_robot(cfg=robot_cfg) + sim.prepare() joint_ids = robot.get_joint_ids(arm_name) qpos_seed = torch.tensor( [[np.pi / 6, 0.0, 0.0, -np.pi / 2, 0.0, 0.0, np.pi / 6]], diff --git a/scripts/tutorials/visualization/README.md b/scripts/tutorials/visualization/README.md index 47c2d2713..5dbca9bb5 100644 --- a/scripts/tutorials/visualization/README.md +++ b/scripts/tutorials/visualization/README.md @@ -82,7 +82,7 @@ Viser is configured. It also rejects Viser startup while the native window is already open. Cloth uses its welded physical surface topology. DexSim does not currently -expose the PhysX soft-body collision topology, so the soft-body preview uses +expose the DexSim soft-body collision topology, so the soft-body preview uses a convex-hull surface over the live collision vertices. It follows deformation but intentionally omits concave render-mesh details. diff --git a/scripts/tutorials/visualization/viser_scene.py b/scripts/tutorials/visualization/viser_scene.py index a3d329932..9350391fc 100644 --- a/scripts/tutorials/visualization/viser_scene.py +++ b/scripts/tutorials/visualization/viser_scene.py @@ -159,8 +159,7 @@ def main() -> None: build_pk_chain=False, ) ) - if sim.is_use_gpu_physics: - sim.init_gpu_physics() + sim.prepare() visualization_cfg = VisualizationCfg( backend="viser", diff --git a/tests/docs/test_check_api_docs.py b/tests/docs/test_check_api_docs.py index d3eb209de..abaf1db7e 100644 --- a/tests/docs/test_check_api_docs.py +++ b/tests/docs/test_check_api_docs.py @@ -71,6 +71,7 @@ def test_discover_public_modules_uses_static_all(tmp_path: Path) -> None: _write(package_path / "feature" / "__init__.py", '__all__ = ["Feature"]\n') _write(package_path / "module.py", '__all__ = ["NotPackageLevel"]\n') _write(package_path / "_private" / "__init__.py", '__all__ = ["Hidden"]\n') + _write(package_path / ".generated" / "__init__.py", "__all__ = build_exports()\n") modules = discover_public_modules((PackageRoot("sample", package_path),)) diff --git a/tests/gen_sim/gradio_ui/test_app_articraft.py b/tests/gen_sim/gradio_ui/test_app_articraft.py index 2400484f0..5474dae04 100644 --- a/tests/gen_sim/gradio_ui/test_app_articraft.py +++ b/tests/gen_sim/gradio_ui/test_app_articraft.py @@ -356,7 +356,8 @@ def start_pipeline(command: list[str]): str(artifact.resolve()), "--asset_type", "articulation", - "--use_usd_properties", + "--asset-physics-mode", + "preserve", "--viser", "--viser-host", "0.0.0.0", diff --git a/tests/gen_sim/scene_engine/test_gravity_settler.py b/tests/gen_sim/scene_engine/test_gravity_settler.py index 5829fa39a..2dffebfba 100644 --- a/tests/gen_sim/scene_engine/test_gravity_settler.py +++ b/tests/gen_sim/scene_engine/test_gravity_settler.py @@ -19,7 +19,10 @@ import pytest -from embodichain.gen_sim.scene_engine.core.scene_object import SceneObject +from embodichain.gen_sim.scene_engine.core.scene_object import ( + ObjectPhysics, + SceneObject, +) from embodichain.gen_sim.scene_engine.pipeline.utils.gravity_settler import ( GravitySettleBody, GravitySettler, @@ -79,3 +82,27 @@ def test_gravity_settler_rejects_dynamic_assets_outside_participants() -> None: dynamic_asset_ids={_ASSET_ID}, static_asset_ids=set(), ).settle() + + +@pytest.mark.parametrize( + ("max_hulls", "expected_approximation", "expected_max_hulls"), + [ + (1, "convex_hull", None), + (8, "convex_decomposition", 8), + ], +) +def test_gravity_settler_normalizes_legacy_hull_budget( + max_hulls: int, + expected_approximation: str, + expected_max_hulls: int | None, +) -> None: + collision = GravitySettler._mesh_collision_cfg( + ObjectPhysics( + body_type="dynamic", + attrs={"mass_props": {"mass": 1.0}}, + max_convex_hull_num=max_hulls, + ) + ) + + assert collision.approximation == expected_approximation + assert collision.max_hulls == expected_max_hulls diff --git a/tests/gen_sim/scene_engine/test_scene_core_and_export.py b/tests/gen_sim/scene_engine/test_scene_core_and_export.py index 95658cf16..dc21f853e 100644 --- a/tests/gen_sim/scene_engine/test_scene_core_and_export.py +++ b/tests/gen_sim/scene_engine/test_scene_core_and_export.py @@ -62,7 +62,10 @@ def _scene_object( def _physics(body_type: str) -> ObjectPhysics: return ObjectPhysics( body_type=body_type, # type: ignore[arg-type] - attrs={"mass": 1.0, "static_friction": 0.8}, + attrs={ + "mass_props": {"mass": 1.0}, + "material_props": {"static_friction": 0.8}, + }, max_convex_hull_num=16, ) diff --git a/tests/gen_sim/scene_engine/test_scene_edit.py b/tests/gen_sim/scene_engine/test_scene_edit.py index fd1a9e933..7591989c3 100644 --- a/tests/gen_sim/scene_engine/test_scene_edit.py +++ b/tests/gen_sim/scene_engine/test_scene_edit.py @@ -50,7 +50,7 @@ def _write_scene_export( "shape_type": "Mesh", "fpath": "mesh_assets/table/table.glb", }, - "attrs": {"mass": 1.0}, + "attrs": {"mass_props": {"mass": 1.0}}, "body_type": "kinematic", "init_pos": [0.0, 0.0, 0.0], "init_rot": [0.0, 0.0, 0.0], @@ -67,7 +67,7 @@ def _write_scene_export( "shape_type": "Mesh", "fpath": "mesh_assets/cup/cup.glb", }, - "attrs": {"mass": 1.0}, + "attrs": {"mass_props": {"mass": 1.0}}, "body_type": "dynamic", "init_pos": [1.0, -3.0, 2.0], "init_rot": [0.0, 0.0, 0.0], diff --git a/tests/gym/envs/expert_program/test_compiler.py b/tests/gym/envs/expert_program/test_compiler.py index d4e8c7d77..b30da3789 100644 --- a/tests/gym/envs/expert_program/test_compiler.py +++ b/tests/gym/envs/expert_program/test_compiler.py @@ -131,10 +131,10 @@ def _integration() -> ExpertProgramIntegrationCfg: def _pose(x: float, y: float = 0.0, z: float = 0.2) -> PoseCfg: - """Build one target pose with an identity WXYZ quaternion.""" + """Build one target pose with an identity XYZW quaternion.""" return PoseCfg( position=(x, y, z), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ) @@ -156,7 +156,7 @@ def _program( def _assert_pose_equal(actual: SemanticPose, expected: SemanticPose) -> None: """Compare owned pose tensor values.""" assert torch.allclose(actual.position, expected.position) - assert torch.allclose(actual.quaternion_wxyz, expected.quaternion_wxyz) + assert torch.allclose(actual.quaternion_xyzw, expected.quaternion_xyzw) def _assert_semantic_call_equal( @@ -241,7 +241,7 @@ def test_compiler_matches_direct_python_semantic_calls_and_sequence_order() -> N ), HandOver( object=SceneObjectRef("cube"), - final_target=SemanticPose(target.position, target.quaternion_wxyz), + final_target=SemanticPose(target.position, target.quaternion_xyzw), resources={"destination": "right_actor"}, ), RegisteredSemanticCall( @@ -312,7 +312,7 @@ def test_repeat_expands_independent_segments_with_cyclic_targets() -> None: assert place.call.at is not None _assert_pose_equal( place.call.at, - SemanticPose(pose.position, pose.quaternion_wxyz), + SemanticPose(pose.position, pose.quaternion_xyzw), ) assert place.target_selections[0].value_index == index validator = segment.validators[0] diff --git a/tests/gym/envs/expert_program/test_decoder.py b/tests/gym/envs/expert_program/test_decoder.py index 52dab18cf..6292a3bc3 100644 --- a/tests/gym/envs/expert_program/test_decoder.py +++ b/tests/gym/envs/expert_program/test_decoder.py @@ -62,15 +62,15 @@ def _program_data() -> dict[str, object]: "values": [ { "position": [0.45, -0.20, 0.20], - "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "quaternion_xyzw": [0.0, 0.0, 0.0, 1.0], }, { "position": [0.45, 0.00, 0.20], - "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "quaternion_xyzw": [0.0, 0.0, 0.0, 1.0], }, { "position": [0.45, 0.20, 0.20], - "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "quaternion_xyzw": [0.0, 0.0, 0.0, 1.0], }, ], } diff --git a/tests/gym/envs/expert_program/test_environment.py b/tests/gym/envs/expert_program/test_environment.py index b7028239c..d0a3248aa 100644 --- a/tests/gym/envs/expert_program/test_environment.py +++ b/tests/gym/envs/expert_program/test_environment.py @@ -479,7 +479,7 @@ def _program_with_later_segment_hooks( values=( PoseCfg( position=(0.4, 0.1, 0.2), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ), ) ) @@ -629,7 +629,7 @@ def test_preflight_preserves_pick_target_lookahead_across_explicit_segments( values=( PoseCfg( position=(0.4, 0.1, 0.2), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ), ) ) diff --git a/tests/gym/envs/expert_program/test_expert_program_cfg.py b/tests/gym/envs/expert_program/test_expert_program_cfg.py index a11597de4..a6f4d7f57 100644 --- a/tests/gym/envs/expert_program/test_expert_program_cfg.py +++ b/tests/gym/envs/expert_program/test_expert_program_cfg.py @@ -183,7 +183,7 @@ def test_pose_rejects_zero_quaternion() -> None: with pytest.raises(ValueError, match="non-zero magnitude"): PoseCfg( position=(0.0, 0.0, 0.0), - quaternion_wxyz=(0.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 0.0), ) diff --git a/tests/gym/envs/expert_program/test_loader.py b/tests/gym/envs/expert_program/test_loader.py index 964d4b185..d1d732c9b 100644 --- a/tests/gym/envs/expert_program/test_loader.py +++ b/tests/gym/envs/expert_program/test_loader.py @@ -219,7 +219,7 @@ def test_loads_expert_program_json_normalizes_oversized_integer() -> None: "values": [ { "position": [10**400, 0, 0], - "quaternion_wxyz": [1, 0, 0, 0], + "quaternion_xyzw": [0, 0, 0, 1], } ], } diff --git a/tests/gym/envs/expert_program/test_simulation_environment.py b/tests/gym/envs/expert_program/test_simulation_environment.py index 1e85dcb71..ce51bf9ab 100644 --- a/tests/gym/envs/expert_program/test_simulation_environment.py +++ b/tests/gym/envs/expert_program/test_simulation_environment.py @@ -149,7 +149,7 @@ _RELEASE_SEPARATION = 0.2 _DIRECT_PLACE_TARGET = SemanticPose( position=(0.0, 0.0, 0.0), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ) _QUICKSTART_MAX_LINES = 15 @@ -556,7 +556,7 @@ def resolve( del call, context, bound pose = SemanticPose( position=(0.0, 0.0, 0.5), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ) return HandOverPoseTargets( final=SemanticObjectTarget(pose=pose), @@ -1254,8 +1254,8 @@ def _pick_place_program_data() -> dict[str, object]: "values": [ { "position": _DIRECT_PLACE_TARGET.position.tolist(), - "quaternion_wxyz": ( - _DIRECT_PLACE_TARGET.quaternion_wxyz.tolist() + "quaternion_xyzw": ( + _DIRECT_PLACE_TARGET.quaternion_xyzw.tolist() ), } ], @@ -1881,7 +1881,7 @@ def without_in_flight_guards( object=SceneObjectRef("cube"), at=SemanticPose( position=(0.0, 0.0, 0.0), - quaternion_wxyz=(1.0, 0.0, 0.0, 0.0), + quaternion_xyzw=(0.0, 0.0, 0.0, 1.0), ), ), ), diff --git a/tests/gym/envs/expert_program/test_simulation_handover.py b/tests/gym/envs/expert_program/test_simulation_handover.py index 7766182f2..7ba91276b 100644 --- a/tests/gym/envs/expert_program/test_simulation_handover.py +++ b/tests/gym/envs/expert_program/test_simulation_handover.py @@ -28,7 +28,7 @@ def _provider() -> ConfiguredHandOverPoseProvider: """Return one deterministic dual-arm transfer declaration.""" return ConfiguredHandOverPoseProvider( final_position=(0.0, -0.2, 0.7), - final_quaternion_wxyz=(1.0, 1.0, 0.0, 0.0), + final_quaternion_xyzw=(1.0, 0.0, 0.0, 1.0), ) @@ -59,7 +59,7 @@ def test_configured_handover_provider_normalizes_and_owns_targets() -> None: ("overrides", "error_type"), [ ({"final_position": (0.0, 0.0)}, TypeError), - ({"final_quaternion_wxyz": (0.0, 0.0, 0.0, 0.0)}, ValueError), + ({"final_quaternion_xyzw": (0.0, 0.0, 0.0, 0.0)}, ValueError), ], ) def test_configured_handover_provider_rejects_invalid_declarations( @@ -69,7 +69,7 @@ def test_configured_handover_provider_rejects_invalid_declarations( """Malformed provider declarations fail before simulation construction.""" values: dict[str, object] = { "final_position": (0.0, -0.2, 0.7), - "final_quaternion_wxyz": (1.0, 0.0, 0.0, 0.0), + "final_quaternion_xyzw": (0.0, 0.0, 0.0, 1.0), } values.update(overrides) diff --git a/tests/gym/envs/expert_program/test_simulation_policies.py b/tests/gym/envs/expert_program/test_simulation_policies.py index 483b32156..0b749a769 100644 --- a/tests/gym/envs/expert_program/test_simulation_policies.py +++ b/tests/gym/envs/expert_program/test_simulation_policies.py @@ -151,7 +151,7 @@ def _compiled_segment(*, settle_preset: str = "fast"): "values": [ { "position": [0.0, 0.0, 0.0], - "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "quaternion_xyzw": [0.0, 0.0, 0.0, 1.0], } ], } diff --git a/tests/gym/envs/expert_program/test_task_hand_over.py b/tests/gym/envs/expert_program/test_task_hand_over.py index 3b60a8ebd..6888cf1d7 100644 --- a/tests/gym/envs/expert_program/test_task_hand_over.py +++ b/tests/gym/envs/expert_program/test_task_hand_over.py @@ -191,7 +191,7 @@ def test_hand_over_gym_config_builds_dual_ur5_pgi_scene() -> None: ) assert [item.uid for item in cfg.background] == [_SUPPORT_SURFACE_UID] assert [item.uid for item in cfg.rigid_object] == [_CAN_SIMULATION_UID] - assert cfg.rigid_object[0].max_convex_hull_num == 16 + assert cfg.rigid_object[0].shape.collision.max_hulls == 16 assert cfg.expert_program is not None assert cfg.expert_program.program_id == "dual_ur5_hand_over" @@ -200,8 +200,8 @@ def test_hand_over_config_owns_tuned_can_and_pgi_physics() -> None: """The sole config source retains the tuned object and gripper dynamics.""" cfg = _configured_env_cfg() - assert cfg.rigid_object[0].attrs.mass == pytest.approx(0.33) - drive = cfg.robot.drive_pros + assert cfg.rigid_object[0].attrs.mass_props.mass == pytest.approx(0.33) + drive = cfg.robot.joint_drive_props expected_values = { "stiffness": 1e3, "damping": 1e2, @@ -215,8 +215,8 @@ def test_hand_over_config_owns_tuned_can_and_pgi_physics() -> None: ) assert values[f"{side}_gripper_finger2_joint_1"] == pytest.approx(0.0) finger_attrs = cfg.robot.link_attrs["gripper_fingers"].attrs - assert finger_attrs.dynamic_friction == pytest.approx(2.0) - assert finger_attrs.static_friction == pytest.approx(2.0) + assert finger_attrs.material_props.dynamic_friction == pytest.approx(2.0) + assert finger_attrs.material_props.static_friction == pytest.approx(2.0) def test_hand_over_runtime_owns_scene_pose_and_evidence_services() -> None: diff --git a/tests/gym/envs/expert_program/test_task_vertical_slices.py b/tests/gym/envs/expert_program/test_task_vertical_slices.py index 414a9e287..78e92545c 100644 --- a/tests/gym/envs/expert_program/test_task_vertical_slices.py +++ b/tests/gym/envs/expert_program/test_task_vertical_slices.py @@ -408,11 +408,11 @@ def test_cube_variant_extends_by_data_without_motion_generation_code() -> None: ( { "position": [-0.25, -0.20, 0.10], - "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "quaternion_xyzw": [0.0, 0.0, 0.0, 1.0], }, { "position": [-0.25, 0.20, 0.10], - "quaternion_wxyz": [1.0, 0.0, 0.0, 0.0], + "quaternion_xyzw": [0.0, 0.0, 0.0, 1.0], }, ) ) diff --git a/tests/gym/envs/managers/test_action_manager.py b/tests/gym/envs/managers/test_action_manager.py index c617fee58..3a0cb16b3 100644 --- a/tests/gym/envs/managers/test_action_manager.py +++ b/tests/gym/envs/managers/test_action_manager.py @@ -160,11 +160,10 @@ def test_eef_pose_term_process_action_7d(): cfg = ActionTermCfg(func=EefPoseTerm, params={"scale": 1.0, "pose_dim": 7}) term = EefPoseTerm(cfg, env) - # 7D: position + quaternion (w,x,y,z) + # 7D: position + quaternion (x,y,z,w) action = torch.zeros(2, 7) action[:, :3] = 0.1 - action[:, 3] = 1.0 # quat w - action[:, 4:7] = 0.0 # quat x,y,z (identity) + action[:, 6] = 1.0 # xyzw identity result = term.process_action(action) assert "qpos" in result diff --git a/tests/gym/envs/managers/test_event_functors.py b/tests/gym/envs/managers/test_event_functors.py index 750fa2399..ac7443150 100644 --- a/tests/gym/envs/managers/test_event_functors.py +++ b/tests/gym/envs/managers/test_event_functors.py @@ -58,17 +58,20 @@ def __init__( self.cfg.shape = Mock() self.cfg.shape.fpath = "test.obj" self.cfg.attrs = Mock() - self.cfg.attrs.mass = 1.0 + self.cfg.attrs.mass_props = Mock(mass=1.0) # Default pose at origin self._pose = torch.eye(4).unsqueeze(0).repeat(num_envs, 1, 1) self._mass = torch.ones(num_envs) * 1.0 + self._inertia = torch.ones(num_envs, 3) self._com = torch.zeros(num_envs, 3) # Mock body_data self.body_data = Mock() + self.body_data.default_mass = self._mass.clone() + self.body_data.default_inertia = self._inertia.clone() self.body_data.default_com_pose = torch.zeros(num_envs, 7) - self.body_data.default_com_pose[:, 3] = 1.0 # quaternion w + self.body_data.default_com_pose[:, 6] = 1.0 # xyzw quaternion w self.body_data.lin_vel = torch.zeros(num_envs, 3) self.body_data.ang_vel = torch.zeros(num_envs, 3) @@ -92,6 +95,17 @@ def set_mass(self, mass, env_ids=None): else: self._mass = mass + def get_inertia(self, env_ids=None): + if env_ids is not None: + return self._inertia[env_ids] + return self._inertia + + def set_inertia(self, inertia, env_ids=None): + if env_ids is not None: + self._inertia[env_ids] = inertia + else: + self._inertia = inertia + class MockRigidObjectGroup: """Mock rigid object group for event functor tests.""" @@ -221,12 +235,17 @@ def __init__( # Default pose at origin (position + quaternion) # Format: (N, 7) - position (3) + quaternion (4) self._pose = torch.zeros(num_envs, 7) - self._pose[:, 3] = 1.0 # quaternion w = 1 (identity rotation) + self._pose[:, 6] = 1.0 # xyzw quaternion w = 1 (identity rotation) - self.default_link_masses = torch.ones( - (self.num_envs, len(self.link_names)), device=self.device + self._inertia = torch.ones( + (self.num_envs, len(self.link_names), 3), device=self.device ) self.body_data = Mock() + self.body_data.default_mass = torch.ones( + (self.num_envs, len(self.link_names)), device=self.device + ) + self.body_data.default_inertia = self._inertia.clone() + self.default_link_masses = self.body_data.default_mass self.body_data.body_link_vel = torch.zeros( self.num_envs, len(self.link_names), 6, device=self.device ) @@ -306,6 +325,30 @@ def set_mass(self, mass, link_names, env_ids=None): for j, name in enumerate(link_names): self._entities[env_idx]._link_masses[name] = mass[i, j].item() + def get_inertia(self, link_names=None, env_ids=None): + """Get link inertia diagonals, matching Articulation API.""" + env_index = torch.as_tensor( + list(range(self.num_envs)) if env_ids is None else env_ids, + dtype=torch.long, + ) + names = self.link_names if link_names is None else list(link_names) + link_index = torch.as_tensor( + [self.link_names.index(name) for name in names], dtype=torch.long + ) + return self._inertia[env_index[:, None], link_index[None, :]] + + def set_inertia(self, inertia, link_names=None, env_ids=None): + """Set link inertia diagonals, matching Articulation API.""" + env_index = torch.as_tensor( + list(range(self.num_envs)) if env_ids is None else env_ids, + dtype=torch.long, + ) + names = self.link_names if link_names is None else list(link_names) + link_index = torch.as_tensor( + [self.link_names.index(name) for name in names], dtype=torch.long + ) + self._inertia[env_index[:, None], link_index[None, :]] = inertia + class MockSim: """Mock simulation for event functor tests.""" @@ -533,6 +576,81 @@ def test_relative_mass_randomization(self): assert torch.all(masses >= 0.5) assert torch.all(masses <= 1.5) + def test_relative_mass_randomization_does_not_accumulate(self): + """Test repeated relative randomization uses the initial mass.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + # The backend-resolved mass is the baseline, not stale config metadata. + env.test_object.cfg.attrs.mass_props.mass = 10.0 + + for _ in range(2): + randomize_rigid_object_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="cube"), + mass_range=(0.5, 0.5), + relative=True, + ) + + masses = env.test_object.get_mass().reshape(-1) + assert torch.allclose(masses, torch.full((4,), 1.5)) + + def test_mass_randomization_recomputes_inertia_from_defaults(self): + """Test inertia scaling uses the initial mass-property snapshot.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + env.test_object._inertia.fill_(9.0) + + randomize_rigid_object_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="cube"), + mass_range=(2.0, 2.0), + ) + + assert torch.allclose(env.test_object.get_inertia(), torch.full((4, 3), 2.0)) + + def test_mass_randomization_enforces_positive_mass(self): + """Test relative offsets cannot produce a non-positive mass.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_rigid_object_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="cube"), + mass_range=(-2.0, -2.0), + relative=True, + min_mass=0.25, + ) + + assert torch.allclose(env.test_object.get_mass(), torch.full((4, 1), 0.25)) + + def test_sampling_uses_rigid_object_device(self, monkeypatch): + """Test samples are allocated on the rigid object's device.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + sampled_device = None + + def fake_sample_uniform(*, lower, upper, size, device): + nonlocal sampled_device + sampled_device = device + return torch.zeros(size, device=device) + + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.randomization.physics.sample_uniform", + fake_sample_uniform, + ) + + randomize_rigid_object_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="cube"), + mass_range=(0.5, 2.0), + ) + + assert sampled_device == env.test_object.device + def test_handles_nonexistent_object(self): """Test that function handles non-existent object gracefully.""" env = MockEnv(num_envs=4) @@ -781,7 +899,7 @@ def test_sets_specific_link_with_list(self): assert torch.all(randomized <= 2.0) def test_relative_mass_randomization(self): - """Test relative mass randomization adds to current mass.""" + """Test relative mass randomization adds to the initial mass.""" env = MockEnv(num_envs=4) env_ids = torch.tensor([0, 1, 2, 3]) @@ -802,6 +920,90 @@ def test_relative_mass_randomization(self): assert torch.all(masses >= 0.5) assert torch.all(masses <= 1.5) + def test_relative_mass_randomization_does_not_accumulate(self): + """Repeated relative randomization uses initialization-time link mass.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + + for _ in range(2): + randomize_articulation_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="articulation"), + mass_range=(0.5, 0.5), + link_names=["base_link"], + relative=True, + ) + + masses = env.test_articulation.get_mass( + link_names=["base_link"], env_ids=env_ids + ) + assert torch.allclose(masses, torch.full((4, 1), 1.5)) + + def test_mass_randomization_recomputes_link_inertia_from_defaults(self): + """Inertia scaling uses initialization snapshots rather than current values.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + env.test_articulation._inertia.fill_(9.0) + + randomize_articulation_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="articulation"), + mass_range=(2.0, 2.0), + link_names=["base_link"], + ) + + inertia = env.test_articulation.get_inertia( + link_names=["base_link"], env_ids=env_ids + ) + assert torch.allclose(inertia, torch.full((4, 1, 3), 2.0)) + + def test_mass_randomization_enforces_positive_link_mass(self): + """Relative offsets cannot produce a non-positive link mass.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + + randomize_articulation_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="articulation"), + mass_range=(-2.0, -2.0), + link_names=["base_link"], + relative=True, + min_mass=0.25, + ) + + masses = env.test_articulation.get_mass( + link_names=["base_link"], env_ids=env_ids + ) + assert torch.allclose(masses, torch.full((4, 1), 0.25)) + + def test_sampling_uses_articulation_device(self, monkeypatch): + """Test tuple-range samples use the articulation's device.""" + env = MockEnv(num_envs=4) + env_ids = torch.tensor([0, 1, 2, 3]) + sampled_device = None + + def fake_sample_uniform(*, lower, upper, size, device): + nonlocal sampled_device + sampled_device = device + return torch.zeros(size, device=device) + + monkeypatch.setattr( + "embodichain.lab.gym.envs.managers.randomization.physics.sample_uniform", + fake_sample_uniform, + ) + + randomize_articulation_mass( + env, + env_ids, + entity_cfg=MagicMock(uid="articulation"), + mass_range=(0.5, 2.0), + ) + + assert sampled_device == env.test_articulation.device + def test_handles_nonexistent_articulation(self): """Test that function handles non-existent articulation gracefully.""" env = MockEnv(num_envs=4) diff --git a/tests/gym/envs/managers/test_observation_functors.py b/tests/gym/envs/managers/test_observation_functors.py index ced6e1f7e..aae6f4c5c 100644 --- a/tests/gym/envs/managers/test_observation_functors.py +++ b/tests/gym/envs/managers/test_observation_functors.py @@ -104,7 +104,7 @@ def get_local_pose(self, to_matrix=True): pos = self._pose[:, :3, 3] # Simple quaternion from identity rotation quat = torch.zeros(self.num_envs, 4) - quat[:, 0] = 1.0 # w=1 (identity) + quat[:, 3] = 1.0 # xyzw identity return torch.cat([pos, quat], dim=-1) def get_mass(self): diff --git a/tests/gym/envs/managers/test_randomize_anchor_height.py b/tests/gym/envs/managers/test_randomize_anchor_height.py index 1c6acf17e..a3f4c9972 100644 --- a/tests/gym/envs/managers/test_randomize_anchor_height.py +++ b/tests/gym/envs/managers/test_randomize_anchor_height.py @@ -41,7 +41,7 @@ def __init__(self, uid: str, num_envs: int = 4): self.cfg = MagicMock() self.cfg.init_pos = [0.0, 0.0, 0.0] self._pose = torch.zeros(num_envs, 7) - self._pose[:, 3] = 1.0 # identity quaternion + self._pose[:, 6] = 1.0 # xyzw identity quaternion self._cleared = False self._cleared_env_ids = None diff --git a/tests/gym/envs/test_base_env.py b/tests/gym/envs/test_base_env.py index b156c9ec5..7ea53dc77 100644 --- a/tests/gym/envs/test_base_env.py +++ b/tests/gym/envs/test_base_env.py @@ -29,7 +29,7 @@ RobotCfg, JointDrivePropertiesCfg, RigidObjectCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.gym.utils.registration import register_env from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -58,7 +58,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, ) @@ -68,19 +68,24 @@ def __init__( **kwargs, ) - def _setup_robot(self, **kwargs): + def _declare_robot(self, **kwargs) -> Robot: file_path = get_data_path("UniversalRobots/UR10/UR10.urdf") - robot: Robot = self.sim.add_robot( + return self.sim.add_robot( cfg=RobotCfg( uid="UR10", fpath=file_path, init_pos=(0, 0, 1), init_qpos=self.robot_init_qpos, - drive_pros=JointDrivePropertiesCfg(drive_type=self.drive_type), + joint_drive_props=JointDrivePropertiesCfg(drive_type=self.drive_type), ) ) + def _setup_robot(self, **kwargs) -> Robot: + robot = self.robot + if robot is None: + raise RuntimeError("UR10 was not declared before simulation prepare.") + qpos_limits = robot.body_data.qpos_limits[0].cpu().numpy() self.single_action_space = gym.spaces.Box( low=qpos_limits[:, 0], high=qpos_limits[:, 1], dtype=np.float32 @@ -96,7 +101,9 @@ def _prepare_scene(self, **kwargs): cfg=RigidObjectCfg( uid="cube", shape=CubeCfg(size=[size, size, size]), - attrs=RigidBodyAttributesCfg(enable_collision=False), + attrs=RigidBodyPhysicsCfg.from_dict( + {"collision_props": {"collision_enabled": False}} + ), init_pos=(0.0, 0.0, 0.5), body_type="kinematic", ), @@ -121,14 +128,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") @@ -221,12 +228,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_embodied_env.py b/tests/gym/envs/test_differentiable_embodied_env.py new file mode 100644 index 000000000..1594b110f --- /dev/null +++ b/tests/gym/envs/test_differentiable_embodied_env.py @@ -0,0 +1,1731 @@ +# ---------------------------------------------------------------------------- +# 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 + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +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.sim.cfg import DefaultPhysicsCfg, NewtonPhysicsCfg +from embodichain.lab.sim.diff import NewtonStepFunc, differentiable_step +from embodichain.lab.sim.diff.runtime import NewtonDifferentiableRuntime +import embodichain.lab.sim.diff.bridge as diff_bridge +from embodichain.lab.sim.sim_manager import SimulationManagerCfg + +_CONTROL_SUBSTEPS = 3 + + +@wp.kernel +def _write_bridge_joint_force_kernel( + action: wp.array(dtype=wp.float32), + joint_f: wp.array(dtype=wp.float32), +) -> None: + """Write one tape-tracked action value into Newton joint force.""" + joint_f[0] = action[0] + + +@wp.kernel +def _bridge_terminal_loss_kernel( + body_q: wp.array(dtype=wp.transform), + body_id: int, + target: wp.vec3, + loss: wp.array(dtype=wp.float32), +) -> None: + """Measure a terminal body-position loss inside the Warp tape.""" + delta = wp.transform_get_translation(body_q[body_id]) - target + loss[0] = wp.dot(delta, delta) + + +class _FakeModel: + """Keep the pre-contract bridge path runnable for clean RED failures.""" + + def __init__(self) -> None: + self.states: list[_FakeState] = [] + + def state(self) -> "_FakeState": + state = _FakeState(f"trajectory-{len(self.states)}") + self.states.append(state) + return state + + +class _FakeState: + """State buffer with explicit detached-copy observability.""" + + def __init__(self, name: str, value: int = 0) -> None: + self.name = name + self.value = value + self.assign_sources: list[_FakeState] = [] + + def assign(self, other: "_FakeState") -> None: + """Copy state and retain every publication source for assertions.""" + self.value = other.value + self.assign_sources.append(other) + + +class _FakeStepper: + """Fallback used only while proving the old private route is rejected.""" + + def __init__(self, *, raise_on_step: bool = False) -> None: + self.calls: list[tuple[object, object, object, float]] = [] + self._raise_on_step = raise_on_step + + def create_contacts(self) -> object: + return object() + + def step( + self, + state_in: object, + state_out: object, + *, + contacts: object, + dt: float, + ) -> None: + self.calls.append((state_in, state_out, contacts, dt)) + if self._raise_on_step: + raise RuntimeError("injected trajectory-step failure") + state_out.value = state_in.value + 1 + + +class _RecordingTape: + """Expose construction, exit, and recording ownership of the fake tape.""" + + def __init__(self, warp: "_RecordingWarp") -> None: + self._warp = warp + + def __enter__(self) -> "_RecordingTape": + assert not self._warp.tape_active + self._warp.tape_active = True + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: Any, + ) -> bool: + del exc_type, exc_value, traceback + assert self._warp.tape_active + self._warp.tape_active = False + self._warp.events.append("tape.exit") + return False + + def reset(self) -> None: + """Model the tape cleanup required before releasing a trajectory.""" + assert not self._warp.tape_active + self._warp.events.append("tape.reset") + + def backward(self, *_args: Any, **_kwargs: Any) -> None: + """Provide the minimal action gradient required by bridge tests.""" + self._warp.events.append("tape.backward") + if self._warp.raise_on_tape_backward: + raise RuntimeError("injected tape backward failure") + if self._warp.last_action is not None: + self._warp.last_action.grad = torch.ones_like(self._warp.last_action) + + def zero(self) -> None: + """Keep the current bridge executable until it migrates to reset().""" + self._warp.events.append("tape.zero") + + +class _RecordingWarp: + """Tiny Warp fake that makes tape ownership observable to a manager.""" + + float32 = object() + + def __init__(self) -> None: + self.tape_active = False + self.events: list[str] = [] + self.last_action: torch.Tensor | None = None + self.raise_on_tape_backward = False + + def Tape(self) -> _RecordingTape: + """Record construction before returning a tape context manager.""" + self.events.append("tape.construct") + return _RecordingTape(self) + + def from_torch( + self, tensor: torch.Tensor, *, requires_grad: bool = False, **_: Any + ) -> torch.Tensor: + """Preserve the test tensor as the fake Warp action array.""" + action = tensor.detach().clone().requires_grad_(requires_grad) + self.last_action = action + return action + + def to_torch(self, tensor: torch.Tensor) -> torch.Tensor: + """Expose a fake Warp array to the PyTorch bridge.""" + if self.last_action is not None and tensor is self.last_action.grad: + self.events.append("action-gradient.capture") + return tensor + + +class _ManagerOwnedTrajectory: + """Fake public trajectory whose stepping must occur inside the tape.""" + + def __init__( + self, + manager: "_TrajectoryNewtonManager", + *, + physics_steps: int, + physics_dt: float, + ) -> None: + self._manager = manager + self.control = object() + self.physics_steps = physics_steps + self.physics_dt = physics_dt + self.total_solver_steps = physics_steps * manager.num_substeps + self.states = [ + _FakeState(f"trajectory-state-{index}") + for index in range(self.total_solver_steps + 1) + ] + self.states[0].assign(manager._state_0) + self.contacts = [object() for _ in range(self.total_solver_steps)] + self.step_calls = 0 + self._released = False + + @property + def final_state(self) -> _FakeState: + """Return the terminal state owned by this one taped trajectory.""" + return self.states[-1] + + def step(self) -> _FakeState: + """Advance the owned trajectory and expose tape placement.""" + self._manager.events.append("trajectory.step") + assert self._manager.warp.tape_active + self.step_calls += 1 + if self._manager.raise_on_trajectory_step: + raise RuntimeError("injected trajectory-step failure") + for state_in, state_out in zip(self.states, self.states[1:]): + state_out.value = state_in.value + 1 + return self.final_state + + def release(self) -> None: + """Release this trajectory's model lease after its tape is reset.""" + if self._released: + return + self._manager._release_differentiable_trajectory(self) + self._released = True + + +class _TrajectoryNewtonManager: + """Fake Newton manager for the manager-owned trajectory bridge contract.""" + + def __init__(self, warp: _RecordingWarp, *, num_substeps: int = 1) -> None: + self.warp = warp + self.events = warp.events + self._state_0 = _FakeState("live-state-0") + self._state_1 = _FakeState("live-state-1") + # Keep the old private path runnable so each regression fails on the + # missing public trajectory contract rather than a fake-only error. + self._model = _FakeModel() + self._control = object() + self.num_substeps = num_substeps + self.solver_dt = 0.01 + self._dt = self.solver_dt * self.num_substeps + self.physics_dt = self._dt + self.trajectory_requests: list[dict[str, Any]] = [] + self.trajectories: list[_ManagerOwnedTrajectory] = [] + self.commits: list[_ManagerOwnedTrajectory] = [] + self.commit_assignment_counts: list[tuple[int, int]] = [] + self._active_trajectory: _ManagerOwnedTrajectory | None = None + self.raise_on_trajectory_step = False + + def create_differentiable_trajectory( + self, *, physics_steps: int, physics_dt: float + ) -> _ManagerOwnedTrajectory: + """Create the public trajectory before the bridge opens its tape.""" + if physics_steps < 1: + raise ValueError("physics_steps must be positive") + if self._active_trajectory is not None: + raise RuntimeError( + "A differentiable trajectory is still active; release it after " + "backward before creating another trajectory." + ) + self.events.append("create") + trajectory = _ManagerOwnedTrajectory( + self, + physics_steps=physics_steps, + physics_dt=physics_dt, + ) + self.trajectories.append(trajectory) + self._active_trajectory = trajectory + self.trajectory_requests.append( + { + "physics_steps": physics_steps, + "physics_dt": physics_dt, + "tape_active": self.warp.tape_active, + } + ) + return trajectory + + def commit_differentiable_trajectory( + self, trajectory: _ManagerOwnedTrajectory + ) -> None: + """Record a detached post-tape publication through the manager API.""" + assert not self.warp.tape_active + assert trajectory in self.trajectories + before = (len(self._state_0.assign_sources), len(self._state_1.assign_sources)) + self._state_0.assign(trajectory.final_state) + self._state_1.assign(trajectory.final_state) + self.commits.append(trajectory) + self.commit_assignment_counts.append( + ( + len(self._state_0.assign_sources) - before[0], + len(self._state_1.assign_sources) - before[1], + ) + ) + self.events.append("commit") + + def _release_differentiable_trajectory( + self, trajectory: _ManagerOwnedTrajectory + ) -> None: + """Release the one active trajectory once tape ownership has ended.""" + assert self._active_trajectory is trajectory + self._active_trajectory = None + self.events.append("trajectory.release") + + +class _TrajectorySimulationManager: + """Bridge-facing manager exposing public and legacy test doubles.""" + + def __init__(self, warp: _RecordingWarp, *, num_substeps: int = 1) -> None: + self.is_newton_backend = True + self.physics = SimpleNamespace( + newton_manager=_TrajectoryNewtonManager(warp, num_substeps=num_substeps) + ) + self.steppers: list[_FakeStepper] = [] + + def create_differentiable_stepper(self) -> _FakeStepper: + """Keep the pre-contract bridge executable for a clean RED failure.""" + stepper = _FakeStepper( + raise_on_step=self.physics.newton_manager.raise_on_trajectory_step + ) + self.steppers.append(stepper) + return stepper + + +class _RealBridgeManager: + """Expose only the Spawn-owned differentiable runtime to the bridge.""" + + def __init__(self, runtime: Any) -> None: + self.is_newton_backend = True + self.differentiable_runtime = runtime + + def create_differentiable_stepper(self) -> None: + """Fail if the bridge retains the removed SimulationManager route.""" + raise AssertionError( + "NewtonStepFunc must use the Spawn differentiable runtime, " + "not SimulationManager.create_differentiable_stepper()." + ) + + +def _route_env( + manager: Any, + *, + mode: str | None = None, + control_substeps: int = _CONTROL_SUBSTEPS, +) -> tuple[DifferentiableEmbodiedEnv, list[object]]: + """Build an uninitialized environment with only the route dependencies.""" + env = object.__new__(DifferentiableEmbodiedEnv) + env.sim = manager + env.cfg = SimpleNamespace(sim_steps_per_control=control_substeps) + if mode is not None: + env.differentiable_step_mode = mode + final_states: list[object] = [] + + def _apply_dynamics_action( + _action_wp: torch.Tensor, _control: Any, tape: Any + ) -> None: + del tape + + def _apply_kinematic_action(_action_wp: torch.Tensor, tape: Any) -> None: + del tape + + def _read_outputs(final_state: object) -> dict[str, Any]: + final_states.append(final_state) + return { + "obs": torch.zeros(1, 1), + "reward": torch.zeros(1), + "terminated": torch.zeros(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + env._apply_dynamics_action_kernel = _apply_dynamics_action + env._apply_action_kernel = _apply_kinematic_action + env._read_outputs = _read_outputs + return env, final_states + + +def _manager_owned_trajectory_sim_state( + manager: _TrajectorySimulationManager, + *, + action_to_control_kernel: Any, + step_mode: str | None = None, + step_fn: Any | None = None, +) -> dict[str, Any]: + """Build the narrow bridge input used by manager-owned trajectory tests.""" + nm = manager.physics.newton_manager + + def _read_outputs(final_state: _FakeState) -> dict[str, Any]: + del final_state + assert nm.warp.tape_active + nm.events.append("outputs") + return { + "obs": torch.zeros(1, 1), + "reward": torch.zeros(1), + "terminated": torch.zeros(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + sim_state: dict[str, Any] = { + "manager": manager, + "substeps": _CONTROL_SUBSTEPS, + "physics_dt": nm.physics_dt, + "action_to_control_kernel": action_to_control_kernel, + "kernel_args": ("kernel-argument",), + "obs_reward_fn": _read_outputs, + } + if step_mode is not None: + sim_state["step_mode"] = step_mode + if step_fn is not None: + sim_state["step_fn"] = step_fn + return sim_state + + +def _assert_tape_reset_then_trajectory_release(events: list[str]) -> None: + """Require one terminal tape reset followed immediately by release.""" + assert events.count("tape.reset") == 1 + assert events.count("trajectory.release") == 1 + reset_index = events.index("tape.reset") + release_index = events.index("trajectory.release") + assert events.index("tape.exit") < reset_index < release_index + assert events[-2:] == ["tape.reset", "trajectory.release"] + + +def _assert_backward_captures_gradient_then_releases(events: list[str]) -> None: + """Require gradient capture before terminal tape and trajectory cleanup.""" + tracked_events = { + "tape.backward", + "action-gradient.capture", + "tape.reset", + "trajectory.release", + } + assert [event for event in events if event in tracked_events] == [ + "tape.backward", + "action-gradient.capture", + "tape.reset", + "trajectory.release", + ] + + +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_default_dynamics_manager_trajectory_lifecycle_is_fully_ordered( + monkeypatch, +) -> None: + """Allocate, tape, action, step, output, and commit stay in one order.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + + def _apply_action(_action: torch.Tensor, *_args: Any) -> None: + assert warp.tape_active + nm.events.append("action") + + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_apply_action, + step_mode="dynamics", + ), + ) + + assert nm.trajectory_requests == [ + { + "physics_steps": _CONTROL_SUBSTEPS, + "physics_dt": nm.physics_dt, + "tape_active": False, + } + ] + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert trajectory.step_calls == 1 + assert nm.events == [ + "create", + "tape.construct", + "action", + "trajectory.step", + "outputs", + "tape.exit", + "commit", + ] + assert nm.commits == [trajectory] + assert nm.commit_assignment_counts == [(1, 1)] + assert [len(state.assign_sources) for state in (nm._state_0, nm._state_1)] == [ + 1, + 1, + ] + + +def test_default_dynamics_action_hook_receives_trajectory_local_control( + monkeypatch, +) -> None: + """The taped action write never targets the manager's shared control.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + received: list[tuple[tuple[Any, ...], bool]] = [] + + def _apply_action(_action: torch.Tensor, *args: Any) -> None: + received.append((args, warp.tape_active)) + + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_apply_action, + step_mode="dynamics", + ), + ) + + nm = manager.physics.newton_manager + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert received == [((trajectory.control, "kernel-argument"), True)] + assert received[0][0][0] is not nm._control + + +def test_dynamics_legacy_action_type_error_is_not_retried_after_creation( + monkeypatch, +) -> None: + """Dynamics propagates a legacy callback error instead of falling back.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + legacy_calls: list[tuple[torch.Tensor, Any]] = [] + + def _legacy_action(action_wp: torch.Tensor, tape: Any) -> None: + legacy_calls.append((action_wp, tape)) + raise TypeError("original legacy action TypeError") + + sim_state = _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_legacy_action, + step_mode="dynamics", + ) + # With no extra kernel arguments, the legacy two-argument callback is + # entered once with local control in its obsolete ``tape`` position. + # Retrying after its body raises TypeError would invoke it a second time. + sim_state["kernel_args"] = () + + with pytest.raises(TypeError) as exc_info: + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + assert str(exc_info.value) == "original legacy action TypeError" + assert len(legacy_calls) == 1 + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert legacy_calls[0][1] is trajectory.control + assert trajectory.step_calls == 0 + assert nm.commits == [] + assert nm.commit_assignment_counts == [] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [], + [], + ] + assert nm._active_trajectory is None + assert trajectory._released + _assert_tape_reset_then_trajectory_release(nm.events) + + +def test_default_dynamics_commits_manager_trajectory_once_after_tape_closes( + monkeypatch, +) -> None: + """A public commit is the sole detached publication of live state.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + nm = manager.physics.newton_manager + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm.events[-1] == "commit" + assert nm.commit_assignment_counts == [(1, 1)] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [trajectory.final_state], + [trajectory.final_state], + ] + + +@pytest.mark.parametrize("failure_site", ("action", "trajectory_step")) +def test_failed_manager_trajectory_forward_resets_and_releases_without_commit( + monkeypatch, failure_site: str +) -> None: + """A failed taped forward releases its manager lease without publishing it.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + + if failure_site == "action": + + def _apply_action(_action: torch.Tensor, *_args: Any) -> None: + assert warp.tape_active + nm.events.append("action.error") + raise RuntimeError("injected action failure") + + error_match = "injected action failure" + else: + nm.raise_on_trajectory_step = True + + def _apply_action(_action: torch.Tensor, *_args: Any) -> None: + assert warp.tape_active + nm.events.append("action") + + error_match = "injected trajectory-step failure" + + with pytest.raises(RuntimeError, match=error_match): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=_apply_action, + step_mode="dynamics", + ), + ) + + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [] + assert nm.commit_assignment_counts == [] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [], + [], + ] + assert nm._active_trajectory is None + assert trajectory._released + _assert_tape_reset_then_trajectory_release(nm.events) + + +def test_backward_resets_tape_then_releases_manager_trajectory(monkeypatch) -> None: + """A grad-tracked trajectory resets its tape before releasing after backward.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + action = torch.zeros(1, requires_grad=True) + + outputs = NewtonStepFunc.apply( + action, + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + outputs[0].sum().backward() + + assert action.grad is not None + _assert_backward_captures_gradient_then_releases(nm.events) + _assert_tape_reset_then_trajectory_release(nm.events) + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm._active_trajectory is None + assert trajectory._released + + +def test_backward_exception_resets_tape_then_releases_manager_trajectory( + monkeypatch, +) -> None: + """A tape-backward failure cannot leave a manager trajectory leased.""" + warp = _RecordingWarp() + warp.raise_on_tape_backward = True + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + action = torch.zeros(1, requires_grad=True) + + outputs = NewtonStepFunc.apply( + action, + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + with pytest.raises(RuntimeError, match="injected tape backward failure"): + outputs[0].sum().backward() + + assert nm.events.count("tape.backward") == 1 + assert "action-gradient.capture" not in nm.events + _assert_tape_reset_then_trajectory_release(nm.events) + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm._active_trajectory is None + assert trajectory._released + + +def test_obs_reward_failure_releases_manager_trajectory_before_fresh_forward( + monkeypatch, +) -> None: + """An output-read error rolls back its lease so the next trajectory starts.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + failing_state = _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ) + + def _raise_from_outputs(_final_state: _FakeState) -> dict[str, Any]: + assert warp.tape_active + nm.events.append("outputs.error") + raise RuntimeError("injected output-read failure") + + failing_state["obs_reward_fn"] = _raise_from_outputs + with pytest.raises(RuntimeError, match="injected output-read failure"): + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), failing_state) + + failure_events = list(nm.events) + assert "trajectory.step" in failure_events + assert "outputs.error" in failure_events + assert failure_events.index("trajectory.step") < failure_events.index( + "outputs.error" + ) + _assert_tape_reset_then_trajectory_release(failure_events) + assert len(nm.trajectories) == 1 + failed_trajectory = nm.trajectories[0] + assert failed_trajectory.step_calls == 1 + assert nm.commits == [] + assert nm.commit_assignment_counts == [] + assert [state.assign_sources for state in (nm._state_0, nm._state_1)] == [ + [], + [], + ] + assert nm._active_trajectory is None + assert failed_trajectory._released + + with torch.no_grad(): + outputs = NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + assert len(outputs) == 4 + assert len(nm.trajectories) == 2 + fresh_trajectory = nm.trajectories[1] + assert fresh_trajectory is not failed_trajectory + assert nm.commits == [fresh_trajectory] + assert nm._active_trajectory is None + assert fresh_trajectory._released + + +def test_no_grad_forward_resets_tape_then_releases_manager_trajectory( + monkeypatch, +) -> None: + """A non-grad forward cannot retain a trajectory lease for backward.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + + with torch.no_grad(): + outputs = NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ), + ) + + assert len(outputs) == 4 + assert not outputs[0].requires_grad + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert nm.commits == [trajectory] + assert nm._active_trajectory is None + assert trajectory._released + assert "tape.backward" not in nm.events + _assert_tape_reset_then_trajectory_release(nm.events) + + +def test_legacy_dynamics_step_fn_is_rejected_before_opening_tape(monkeypatch) -> None: + """An untrusted callback cannot silently bypass default solver dynamics.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + legacy_calls: list[None] = [] + + def _legacy_step() -> _FakeState: + legacy_calls.append(None) + return _FakeState("legacy-dynamics-final") + + with pytest.raises(ValueError, match=r"step_fn.*kinematics"): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + step_fn=_legacy_step, + ), + ) + + assert legacy_calls == [] + assert warp.events == [] + + +def test_missing_step_mode_with_step_fn_is_rejected_before_opening_tape( + monkeypatch, +) -> None: + """Historical implicit-FK dictionaries cannot bypass solver dynamics.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + + with pytest.raises(ValueError, match=r"step_mode.*kinematics"): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_fn=lambda: _FakeState("implicit-legacy-final"), + ), + ) + + assert warp.events == [] + + +def test_bridge_rejects_invalid_step_mode_before_opening_tape(monkeypatch) -> None: + """Direct bridge callers cannot open a tape for an unsupported mode.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + + with pytest.raises(ValueError, match=r"step_mode.*dynamics.*kinematics"): + NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="unsupported", + ), + ) + + assert warp.events == [] + assert manager.physics.newton_manager.trajectory_requests == [] + assert manager.steppers == [] + + +def test_explicit_kinematics_step_fn_remains_a_supported_bridge_route( + monkeypatch, +) -> None: + """The deliberate kinematics escape hatch does not request a trajectory.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + kinematic_calls: list[None] = [] + final_state = _FakeState("kinematic-final") + + def _kinematic_step() -> _FakeState: + kinematic_calls.append(None) + return final_state + + outputs = NewtonStepFunc.apply( + torch.zeros(1, requires_grad=True), + _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="kinematics", + step_fn=_kinematic_step, + ), + ) + + assert len(outputs) == 4 + assert kinematic_calls == [None] + assert manager.physics.newton_manager.trajectory_requests == [] + + +def test_environment_sim_state_marks_default_and_explicit_kinematics_routes() -> None: + """The bridge can distinguish an explicit FK request from legacy bypasses.""" + dynamics_env, _ = _route_env(SimpleNamespace()) + kinematics_env, _ = _route_env(SimpleNamespace(), mode="kinematics") + kinematics_env._make_kinematic_step_fn = lambda: (lambda: _FakeState("fk")) + + dynamics_state = dynamics_env._build_sim_state_dict(torch.zeros(1)) + kinematics_state = kinematics_env._build_sim_state_dict(torch.zeros(1)) + + assert dynamics_state["step_mode"] == "dynamics" + assert kinematics_state["step_mode"] == "kinematics" + + +def test_environment_dynamics_hook_receives_local_control_with_migration_api() -> None: + """The default environment wrapper calls only the v1 dynamics hook.""" + env, _ = _route_env(SimpleNamespace()) + dynamics_calls: list[tuple[object, object, object]] = [] + legacy_calls: list[tuple[object, object]] = [] + action = object() + control = object() + + def _dynamics_action( + action_wp: object, trajectory_control: object, tape: object + ) -> None: + dynamics_calls.append((action_wp, trajectory_control, tape)) + + def _legacy_action(action_wp: object, tape: object) -> None: + legacy_calls.append((action_wp, tape)) + + env._apply_dynamics_action_kernel = _dynamics_action + env._apply_action_kernel = _legacy_action + sim_state = env._build_sim_state_dict(torch.zeros(1)) + sim_state["action_to_control_kernel"](action, control, "kernel-argument") + + assert dynamics_calls == [(action, control, None)] + assert legacy_calls == [] + + +def test_environment_dynamics_hook_observes_only_its_active_tape( + monkeypatch, +) -> None: + """The bridge binds the tape through a per-step wrapper closure.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + observed_tapes: list[object | None] = [] + + def _dynamics_action( + _action_wp: object, + _trajectory_control: object, + tape: object | None, + ) -> None: + observed_tapes.append(tape) + + env._apply_dynamics_action_kernel = _dynamics_action + sim_state = env._build_sim_state_dict(torch.zeros(1)) + + with torch.no_grad(): + NewtonStepFunc.apply(torch.zeros(1), sim_state) + + assert len(observed_tapes) == 1 + assert isinstance(observed_tapes[0], _RecordingTape) + + sim_state["action_to_control_kernel"](object(), object()) + assert observed_tapes[-1] is None + + +def test_environment_rejects_legacy_dynamics_action_hook_with_migration_error( + monkeypatch, +) -> None: + """Default dynamics cannot silently keep the pre-local-control hook.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + env = object.__new__(DifferentiableEmbodiedEnv) + env.sim = manager + env.cfg = SimpleNamespace(sim_steps_per_control=_CONTROL_SUBSTEPS) + env._apply_dynamics_action_kernel = None + env._apply_action_kernel = lambda _action, tape: None + env._read_outputs = lambda _state: { + "obs": torch.zeros(1, 1), + "reward": torch.zeros(1), + "terminated": torch.zeros(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + with pytest.raises( + NotImplementedError, match=r"legacy.*_apply_dynamics_action_kernel" + ): + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + assert warp.events == [] + + +def test_environment_kinematics_hook_keeps_its_strict_legacy_signature( + monkeypatch, +) -> None: + """FK-only bridge execution receives action and tape, never local control.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager, mode="kinematics") + monkeypatch.setattr(diff_bridge, "wp", warp) + calls: list[tuple[torch.Tensor, object]] = [] + final_state = _FakeState("kinematic-final") + + def _kinematic_action(action_wp: torch.Tensor, tape: object) -> None: + calls.append((action_wp, tape)) + + env._apply_action_kernel = _kinematic_action + env._make_kinematic_step_fn = lambda: (lambda: final_state) + action = torch.zeros(1, requires_grad=True) + sim_state = env._build_sim_state_dict(action) + outputs = NewtonStepFunc.apply(action, sim_state) + + assert len(outputs) == 4 + assert len(calls) == 1 + assert torch.equal(calls[0][0], action) + assert isinstance(calls[0][1], _RecordingTape) + assert manager.physics.newton_manager.trajectory_requests == [] + + +def test_grad_terminal_step_defers_reset_until_after_backward(monkeypatch) -> None: + """A terminal grad step must return before touching fenced live state.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + reset_calls: list[torch.Tensor] = [] + + def _terminal_outputs(_final_state: object) -> dict[str, Any]: + return { + "obs": torch.full((1, 1), 7.0), + "reward": torch.full((1,), 3.0), + "terminated": torch.ones(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + def _reset(*, options: dict[str, Any]): + if nm._active_trajectory is not None: + raise RuntimeError("reset crossed an active Newton trajectory fence") + reset_ids = torch.as_tensor(options["reset_ids"]).clone() + reset_calls.append(reset_ids) + return torch.full((1, 1), -1.0), {} + + env._read_outputs = _terminal_outputs + env.reset = _reset + action = torch.zeros(1, requires_grad=True) + + obs, reward, terminated, truncated, info = env.step(action) + + assert torch.equal(obs.detach(), torch.full((1, 1), 7.0)) + assert terminated.tolist() == [True] + assert truncated.tolist() == [False] + assert reset_calls == [] + assert info["requires_reset_after_backward"] is True + assert torch.equal(info["deferred_reset_ids"], torch.tensor([0])) + assert nm._active_trajectory is not None + + reward.sum().backward() + + assert action.grad is not None + assert nm._active_trajectory is None + env.reset(options={"reset_ids": info["deferred_reset_ids"]}) + assert len(reset_calls) == 1 + assert torch.equal(reset_calls[0], torch.tensor([0])) + + +def test_no_grad_terminal_step_keeps_synchronous_auto_reset(monkeypatch) -> None: + """A terminal no-grad step may reset after its tape is released.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, _ = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + reset_calls: list[torch.Tensor] = [] + + env._read_outputs = lambda _state: { + "obs": torch.full((1, 1), 7.0), + "reward": torch.full((1,), 3.0), + "terminated": torch.ones(1, dtype=torch.bool), + "truncated": torch.zeros(1, dtype=torch.bool), + "_order": ("obs", "reward", "terminated", "truncated"), + "_grad_track": {}, + } + + def _reset(*, options: dict[str, Any]): + assert nm._active_trajectory is None + reset_ids = torch.as_tensor(options["reset_ids"]).clone() + reset_calls.append(reset_ids) + return torch.full((1, 1), -1.0), {} + + env.reset = _reset + with torch.no_grad(): + obs, reward, terminated, truncated, info = env.step( + torch.zeros(1, requires_grad=True) + ) + + assert torch.equal(obs, torch.full((1, 1), -1.0)) + assert not reward.requires_grad + assert terminated.tolist() == [True] + assert truncated.tolist() == [False] + assert len(reset_calls) == 1 + assert torch.equal(reset_calls[0], torch.tensor([0])) + assert "deferred_reset_ids" not in info + assert "requires_reset_after_backward" not in info + + +def test_default_dynamics_route_uses_manager_trajectory_without_bypass(monkeypatch): + """Default state construction delegates one control step to Newton.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, final_states = _route_env(manager) + monkeypatch.setattr(diff_bridge, "wp", warp) + + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + outputs = NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + nm = manager.physics.newton_manager + assert sim_state["step_mode"] == "dynamics" + assert "step_fn" not in sim_state + assert len(outputs) == 4 + assert len(nm.trajectories) == 1 + assert nm.trajectories[0].total_solver_steps == _CONTROL_SUBSTEPS + assert final_states[0].value == _CONTROL_SUBSTEPS + + +def test_dynamics_bridge_keeps_an_odd_continuous_horizon_in_one_trajectory( + monkeypatch, +): + """A continuous odd horizon is one lease-owning manager trajectory.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + env, final_states = _route_env(manager, control_substeps=5) + monkeypatch.setattr(diff_bridge, "wp", warp) + + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + nm = manager.physics.newton_manager + assert [state.value for state in final_states] == [5] + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert trajectory.total_solver_steps == 5 + assert nm._state_0.value == 5 + assert nm._state_1.value == 5 + assert nm._state_0.assign_sources == [trajectory.final_state] + assert nm._state_1.assign_sources == [trajectory.final_state] + assert len({id(state) for state in trajectory.states}) == 6 + assert all(state not in {nm._state_0, nm._state_1} for state in trajectory.states) + assert len({id(contact) for contact in trajectory.contacts}) == 5 + + +def test_dynamics_bridge_rejects_a_second_outstanding_manager_trajectory( + monkeypatch, +) -> None: + """A second grad forward requires release of the first trajectory lease.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + sim_state = _manager_owned_trajectory_sim_state( + manager, + action_to_control_kernel=lambda _action, *_args: None, + step_mode="dynamics", + ) + + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + with pytest.raises(RuntimeError, match=r"trajectory.*active.*release"): + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + +def test_dynamics_bridge_multiplies_control_and_newton_substeps(monkeypatch): + """One control step preserves EmbodiChain and Newton time semantics.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp, num_substeps=3) + env, _ = _route_env(manager, control_substeps=2) + monkeypatch.setattr(diff_bridge, "wp", warp) + + sim_state = env._build_sim_state_dict(torch.zeros(1, requires_grad=True)) + NewtonStepFunc.apply(torch.zeros(1, requires_grad=True), sim_state) + + nm = manager.physics.newton_manager + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert trajectory.physics_steps == 2 + assert trajectory.physics_dt == nm.physics_dt + assert trajectory.total_solver_steps == 6 + + +def test_differentiable_step_uses_manager_owned_trajectory_and_local_control( + monkeypatch, +): + """The low-level helper also delegates state publication to Newton.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp) + monkeypatch.setattr(diff_bridge, "wp", warp) + received: list[tuple[Any, ...]] = [] + + result = differentiable_step( + manager, + apply_control_fn=lambda *args: received.append(args), + substeps=_CONTROL_SUBSTEPS, + ) + nm = manager.physics.newton_manager + + assert len(nm.trajectories) == 1 + trajectory = nm.trajectories[0] + assert any(trajectory.control in args for args in received) + assert result["trajectory"] is trajectory + assert nm.commits == [trajectory] + assert nm.commit_assignment_counts == [(1, 1)] + + +def test_differentiable_step_rejects_substeps_not_divisible_by_newton_substeps( + monkeypatch, +) -> None: + """A low-level solver horizon must map to whole Newton physics steps.""" + warp = _RecordingWarp() + manager = _TrajectorySimulationManager(warp, num_substeps=2) + monkeypatch.setattr(diff_bridge, "wp", warp) + nm = manager.physics.newton_manager + control_calls: list[tuple[Any, ...]] = [] + + with pytest.raises(ValueError, match=r"substeps.*divisible.*num_substeps"): + differentiable_step( + manager, + apply_control_fn=lambda *args: control_calls.append(args), + substeps=3, + ) + + assert control_calls == [] + assert nm.trajectory_requests == [] + assert manager.steppers == [] + assert warp.events == [] + + +@pytest.mark.parametrize("substeps", (0, -1)) +def test_differentiable_step_rejects_nonpositive_substeps(substeps: int) -> None: + """The public helper rejects an invalid empty solver trajectory.""" + manager = _TrajectorySimulationManager(_RecordingWarp()) + + with pytest.raises(ValueError, match=r"positive"): + differentiable_step( + manager, + apply_control_fn=lambda *_args: None, + substeps=substeps, + ) + + +def test_cpu_spawn_trajectory_retains_local_control_gradient_and_fd(tmp_path): + """The Spawn bridge keeps a local control trajectory across two steps.""" + newton = pytest.importorskip("newton") + pytest.importorskip("dexsim.engine.newton_physics") + from dexsim.engine.newton_physics import ( + NewtonCfg, + NewtonCollisionPipelineCfg, + SemiImplicitSolverCfg, + ) + from dexsim.engine.newton_physics.newton_backend import NewtonBackend + + previous_kernel_cache_dir = wp.config.kernel_cache_dir + previous_verify_access = wp.config.verify_autograd_array_access + backend = None + wp.config.kernel_cache_dir = str(tmp_path / "warp_cache") + wp.config.verify_autograd_array_access = True + try: + cfg = NewtonCfg() + cfg.device = "cpu" + cfg.dt = 1.0 / 60.0 + cfg.num_substeps = 2 + cfg.requires_grad = True + cfg.use_cuda_graph = False + cfg.solver_cfg = SemiImplicitSolverCfg() + cfg.collision_pipeline_cfg = NewtonCollisionPipelineCfg( + broad_phase="explicit", + requires_grad=True, + ) + backend = NewtonBackend(cfg) + shape_cfg = newton.ModelBuilder.ShapeConfig( + ke=1.0e4, + kd=1.0e1, + kf=0.0, + mu=0.0, + ) + body_id = backend.builder.add_body( + xform=wp.transform(wp.vec3(0.0, 0.0, 0.5), wp.quat_identity()), + mass=1.0, + label="embodichain_manager_trajectory_gradient_ball", + ) + backend.builder.add_shape_sphere(body=body_id, radius=0.1, cfg=shape_cfg) + backend.builder.add_ground_plane(cfg=shape_cfg) + backend.finalize() + nm = NewtonDifferentiableRuntime(lambda: backend) + assert nm._model.joint_count == 1 + + manager = _RealBridgeManager(nm) + initial_state = nm._model.state() + initial_state.assign(nm._state_0) + target = wp.vec3(0.5, 0.0, 0.5) + + def _restore_initial_state() -> None: + nm._state_0.assign(initial_state) + nm._state_1.assign(initial_state) + + def _run( + action_value: float, *, requires_grad: bool + ) -> tuple[torch.Tensor, torch.Tensor, list[Any]]: + loss_wp = wp.zeros( + 1, + dtype=wp.float32, + device=nm._state_0.body_q.device, + requires_grad=True, + ) + local_controls: list[Any] = [] + + def _apply_control(action_wp: Any, *args: Any) -> None: + assert len(args) == 1, "Bridge must pass exactly one local control." + control = args[0] + assert control.joint_f is not None + local_controls.append(control) + wp.launch( + _write_bridge_joint_force_kernel, + dim=1, + inputs=[action_wp, control.joint_f], + device=control.joint_f.device, + ) + + def _read_reward(final_state: Any) -> dict[str, Any]: + loss_wp.zero_() + wp.launch( + _bridge_terminal_loss_kernel, + dim=1, + inputs=[final_state.body_q, body_id, target, loss_wp], + device=final_state.body_q.device, + ) + return { + "reward": wp.to_torch(loss_wp), + "_order": ("reward",), + "_grad_track": {"reward": loss_wp}, + } + + action = torch.tensor( + [action_value], dtype=torch.float32, requires_grad=requires_grad + ) + sim_state = { + "manager": manager, + "step_mode": "dynamics", + "substeps": 2, + "physics_dt": cfg.dt, + "action_to_control_kernel": _apply_control, + "kernel_args": (), + "obs_reward_fn": _read_reward, + } + return NewtonStepFunc.apply(action, sim_state)[0], action, local_controls + + reward, action, local_controls = _run(1.0, requires_grad=True) + reward.backward() + + assert len(local_controls) == 1 + assert action.grad is not None + analytic_gradient = float(action.grad[0]) + assert np.isfinite(analytic_gradient) + assert not np.isclose(analytic_gradient, 0.0) + first_final_state = nm._state_0.body_q.numpy().copy() + assert np.allclose( + nm._state_0.body_q.numpy(), nm._state_1.body_q.numpy(), atol=1.0e-6 + ) + + continuation_reward, continuation_action, continuation_controls = _run( + 1.0, requires_grad=True + ) + continuation_reward.backward() + assert len(continuation_controls) == 1 + assert continuation_action.grad is not None + assert np.isfinite(continuation_action.grad).all() + assert not np.allclose(first_final_state, nm._state_0.body_q.numpy()) + + def _reward_value(action_value: float) -> float: + _restore_initial_state() + value, _action, controls = _run(action_value, requires_grad=False) + assert len(controls) == 1 + return float(value.detach()) + + epsilon = 1.0e-3 + finite_difference_gradient = ( + _reward_value(1.0 + epsilon) - _reward_value(1.0 - epsilon) + ) / (2.0 * epsilon) + assert np.isclose( + analytic_gradient, + finite_difference_gradient, + rtol=2.0e-2, + atol=1.0e-4, + ) + finally: + if backend is not None: + backend.close() + wp.config.verify_autograd_array_access = previous_verify_access + if previous_kernel_cache_dir is None: + from warp._src.build import init_kernel_cache + + init_kernel_cache() + else: + wp.config.kernel_cache_dir = previous_kernel_cache_dir + + +def test_dynamics_environment_does_not_expose_generic_step_helper(): + """Only the low-level bridge may accept an arbitrary dynamics callback.""" + assert "_make_step_fn" not in DifferentiableEmbodiedEnv.__dict__ + + +def test_kinematics_route_uses_only_named_kinematic_hook(): + """FK stepping is selected only through the explicit kinematics mode.""" + manager = SimpleNamespace() + env, _ = _route_env(manager, mode="kinematics") + expected_state = object() + kinematic_calls: list[None] = [] + + def _kinematic_step() -> object: + kinematic_calls.append(None) + return expected_state + + def _generic_step_fn() -> object: + raise AssertionError("The generic step helper must not route kinematics.") + + env._make_kinematic_step_fn = lambda: _kinematic_step + env._make_step_fn = _generic_step_fn + + sim_state = env._build_sim_state_dict(torch.zeros(1)) + + assert sim_state["step_mode"] == "kinematics" + assert sim_state["step_fn"]() is expected_state + assert kinematic_calls == [None] + + +def test_kinematics_route_requires_named_hook(): + """Kinematics mode rejects environments that do not define its hook.""" + manager = SimpleNamespace() + env, _ = _route_env(manager, mode="kinematics") + + with pytest.raises( + NotImplementedError, match=r"kinematics.*_make_kinematic_step_fn" + ): + env._build_sim_state_dict(torch.zeros(1)) + + +def test_invalid_differentiable_step_mode_raises_clear_error(): + """Unsupported stepping modes fail before creating a bridge callback.""" + manager = SimpleNamespace() + env, _ = _route_env(manager, mode="unsupported") + + with pytest.raises( + ValueError, match=r"differentiable_step_mode.*dynamics.*kinematics" + ): + env._build_sim_state_dict(torch.zeros(1)) + + +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_tasks.special.franka_reach_apg import FrankaReachApgEnv + + return FrankaReachApgEnv + + +def test_franka_kinematics_build_snapshots_live_primal_before_bridge( + monkeypatch, +) -> None: + """Franka must detach taped FK inputs before the parent opens a tape.""" + from embodichain_tasks.special import franka_reach_apg + + env = object.__new__(franka_reach_apg.FrankaReachApgEnv) + live_joint_q = object() + snapshot_joint_q = object() + fresh_fk_state = object() + events: list[str] = [] + env.sim = SimpleNamespace( + differentiable_runtime=SimpleNamespace( + current_state=SimpleNamespace(joint_q=live_joint_q), + model=SimpleNamespace( + state=lambda: (events.append("state"), fresh_fk_state)[1] + ), + ) + ) + + def _clone(array: object) -> object: + assert array is live_joint_q + events.append("clone") + return snapshot_joint_q + + def _parent_build(_self: object, _action: torch.Tensor) -> dict[str, Any]: + events.append("parent") + assert env._current_joint_q_snapshot is snapshot_joint_q + assert env._fk_state is fresh_fk_state + return {"prepared": True} + + monkeypatch.setattr(franka_reach_apg.wp, "clone", _clone) + monkeypatch.setattr( + DifferentiableEmbodiedEnv, + "_build_sim_state_dict", + _parent_build, + ) + + result = env._build_sim_state_dict(torch.zeros(1, 7)) + + assert result == {"prepared": True} + assert events == ["clone", "state", "parent"] + + +def test_franka_action_kernel_reads_snapshot_instead_of_live_state(monkeypatch) -> None: + """The recorded action kernel must not capture mutable manager state.""" + from embodichain_tasks.special import franka_reach_apg + + env = object.__new__(franka_reach_apg.FrankaReachApgEnv) + live_joint_q = object() + snapshot_joint_q = object() + target_joint_q = object() + action_wp = object() + launch_inputs: list[object] = [] + env.sim = SimpleNamespace( + num_envs=1, + physics=SimpleNamespace( + newton_manager=SimpleNamespace( + _state_0=SimpleNamespace(joint_q=live_joint_q) + ) + ), + ) + env._current_joint_q_snapshot = snapshot_joint_q + env._n_joints_per_env = 9 + env._wp_device = "cpu" + env._limit_lo_wp = object() + env._limit_hi_wp = object() + env._action_scale = 0.2 + + monkeypatch.setattr( + franka_reach_apg.wp, + "zeros", + lambda *_args, **_kwargs: target_joint_q, + ) + + def _launch(*_args: Any, inputs: list[object], **_kwargs: Any) -> None: + launch_inputs.extend(inputs) + + monkeypatch.setattr(franka_reach_apg.wp, "launch", _launch) + + env._apply_action_kernel(action_wp, tape=object()) + + assert launch_inputs[0] is action_wp + assert launch_inputs[1] is snapshot_joint_q + assert launch_inputs[1] is not live_joint_q + assert launch_inputs[2] is target_joint_q + + +def test_franka_snapshot_keeps_gradient_after_live_state_mutation_and_matches_fd( + monkeypatch, + tmp_path, +) -> None: + """Detached FK input survives live writes before backward under strict mode.""" + from embodichain_tasks.special import franka_reach_apg + + env = object.__new__(franka_reach_apg.FrankaReachApgEnv) + device = "cpu" + live_joint_q = wp.zeros(7, dtype=wp.float32, device=device) + env.sim = SimpleNamespace( + num_envs=1, + differentiable_runtime=SimpleNamespace( + current_state=SimpleNamespace(joint_q=live_joint_q), + model=SimpleNamespace(state=lambda: object()), + ), + ) + env._wp_device = device + env._n_joints_per_env = 7 + env._limit_lo_wp = wp.array( + np.full(7, -10.0, dtype=np.float32), + dtype=wp.float32, + device=device, + ) + env._limit_hi_wp = wp.array( + np.full(7, 10.0, dtype=np.float32), + dtype=wp.float32, + device=device, + ) + env._action_scale = 0.2 + monkeypatch.setattr( + DifferentiableEmbodiedEnv, + "_build_sim_state_dict", + lambda _self, _action: {}, + ) + env._build_sim_state_dict(torch.zeros(1, 7)) + + previous_verify_access = wp.config.verify_autograd_array_access + previous_kernel_cache_dir = wp.config.kernel_cache_dir + wp.config.verify_autograd_array_access = True + wp.config.kernel_cache_dir = str(tmp_path / "warp_cache") + tape = wp.Tape() + try: + action_wp = wp.array( + np.zeros(7, dtype=np.float32), + dtype=wp.float32, + device=device, + requires_grad=True, + ) + with tape: + env._apply_action_kernel(action_wp, tape=tape) + analytic_output = env._new_joint_q + + wp.copy( + live_joint_q, + wp.array( + np.full(7, 5.0, dtype=np.float32), + dtype=wp.float32, + device=device, + ), + ) + tape.backward(grads={analytic_output: wp.ones_like(analytic_output)}) + analytic_gradient = action_wp.grad.numpy().copy() + + assert np.isfinite(analytic_gradient).all() + assert np.all(np.abs(analytic_gradient) > 0.0) + + def _loss(action_value: float) -> float: + values = np.zeros(7, dtype=np.float32) + values[0] = action_value + finite_difference_action = wp.array( + values, + dtype=wp.float32, + device=device, + ) + env._apply_action_kernel(finite_difference_action, tape=object()) + return float(env._new_joint_q.numpy().sum()) + + epsilon = 1.0e-3 + finite_difference_gradient = (_loss(epsilon) - _loss(-epsilon)) / ( + 2.0 * epsilon + ) + assert np.isclose( + analytic_gradient[0], + finite_difference_gradient, + rtol=1.0e-4, + atol=1.0e-5, + ) + finally: + tape.reset() + wp.config.verify_autograd_array_access = previous_verify_access + wp.config.kernel_cache_dir = previous_kernel_cache_dir + + +@pytest.mark.requires_sim +@pytest.mark.gpu +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) + try: + 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() + finally: + env.close() + + +@pytest.mark.requires_sim +@pytest.mark.gpu +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) + try: + 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}" + finally: + env.close() diff --git a/tests/gym/envs/test_embodied_env.py b/tests/gym/envs/test_embodied_env.py index 2ecacc0a9..1a16d0dd5 100644 --- a/tests/gym/envs/test_embodied_env.py +++ b/tests/gym/envs/test_embodied_env.py @@ -75,7 +75,7 @@ ], "robot": { "fpath": urdf_path, - "drive_pros": {"stiffness": {"joint[1-6]": 200.0}}, + "joint_drive_props": {"stiffness": {"joint[1-6]": 200.0}}, "solver_cfg": { "class_type": "PytorchSolver", "end_link_name": "ee_link", @@ -101,9 +101,12 @@ "shape": { "shape_type": "Mesh", "fpath": "ShopTableSimple/shop_table_simple.ply", + "collision": { + "approximation": "convex_decomposition", + "max_hulls": 2, + }, }, - "max_convex_hull_num": 2, - "attrs": {"mass": 10.0}, + "attrs": {"mass_props": {"mass": 10.0}}, "body_scale": (2, 1.6, 1), } ], @@ -145,14 +148,14 @@ def test_visual_randomization_filter_keeps_deterministic_material_events(): 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/envs/test_replay.py b/tests/gym/envs/test_replay.py index c2e3a7830..91a985cc6 100644 --- a/tests/gym/envs/test_replay.py +++ b/tests/gym/envs/test_replay.py @@ -56,7 +56,7 @@ def __init__( uid="UR10", fpath=get_data_path("UniversalRobots/UR10/UR10.urdf"), init_pos=(0.0, 0.0, 1.0), - drive_pros=JointDrivePropertiesCfg(drive_type="force"), + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), ) cfg.rigid_object = [ RigidObjectCfg( @@ -469,7 +469,7 @@ def __init__( uid="UR10", fpath=get_data_path("UniversalRobots/UR10/UR10.urdf"), init_pos=(0.0, 0.0, 1.0), - drive_pros=JointDrivePropertiesCfg(drive_type="force"), + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), ) cfg.rigid_object = [ RigidObjectCfg( diff --git a/tests/gym/utils/test_gym_utils.py b/tests/gym/utils/test_gym_utils.py index ce90bcfc1..8ae069c98 100644 --- a/tests/gym/utils/test_gym_utils.py +++ b/tests/gym/utils/test_gym_utils.py @@ -54,6 +54,29 @@ CUBE_SCENE_REGISTRY_ID = "expert_program_repeated_pick_place" +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.""" @@ -284,6 +307,7 @@ def test_merge_args_with_gym_config_overrides_max_episodes(): device="cpu", headless=False, renderer="auto", + physics="default", gpu_id=0, arena_space=5.0, max_episodes=12, @@ -303,6 +327,7 @@ def test_merge_args_with_gym_config_keeps_default_max_episodes(): device="cpu", headless=False, renderer="auto", + physics="default", gpu_id=0, arena_space=5.0, max_episodes=None, diff --git a/tests/lab/scripts/test_preview_asset.py b/tests/lab/scripts/test_preview_asset.py index 2c17c900e..e5e926ec5 100644 --- a/tests/lab/scripts/test_preview_asset.py +++ b/tests/lab/scripts/test_preview_asset.py @@ -69,6 +69,21 @@ def test_joint_control_is_enabled_by_default_and_can_be_disabled() -> None: assert disabled.joint_control is False +def test_asset_physics_mode_accepts_cli_spelling_variants() -> None: + parser = _create_parser() + default = parser.parse_args(["--asset_path", ASSET_PATH]) + hyphenated = parser.parse_args( + ["--asset_path", ASSET_PATH, "--asset-physics-mode", "preserve"] + ) + underscored = parser.parse_args( + ["--asset_path", ASSET_PATH, "--asset_physics_mode", "preserve"] + ) + + assert default.asset_physics_mode == "overlay" + assert hyphenated.asset_physics_mode == "preserve" + assert underscored.asset_physics_mode == "preserve" + + def test_loaded_assets_are_published_immediately_in_viser() -> None: """Assets added after manager construction should be captured before waiting.""" sim = Mock() diff --git a/tests/learning/test_shared_rollout.py b/tests/learning/test_shared_rollout.py index 907325948..388baae65 100644 --- a/tests/learning/test_shared_rollout.py +++ b/tests/learning/test_shared_rollout.py @@ -192,7 +192,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/sim/atomic_actions/test_curobo_motion_strategy_e2e.py b/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py index 7ea946be9..06f2b645a 100644 --- a/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py +++ b/tests/sim/atomic_actions/test_curobo_motion_strategy_e2e.py @@ -33,7 +33,7 @@ pytest.skip("cuRobo V2 requires CUDA", allow_module_level=True) from embodichain.lab.sim import SimulationManager, SimulationManagerCfg # noqa: E402 -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 +from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg # noqa: E402 from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 from embodichain.lab.sim.robots import FrankaPandaCfg # noqa: E402 from embodichain.lab.sim.shapes import CubeCfg # noqa: E402 @@ -71,12 +71,13 @@ def _make_franka_curobo_engine(): cfg=RigidObjectCfg( uid="demo_block", shape=CubeCfg(size=DEMO_BLOCK_DIMS), - attrs=RigidBodyAttributesCfg(), - body_type="kinematic", + attrs=RigidBodyPhysicsCfg(), + body_type="static", init_pos=DEMO_BLOCK_POS, init_rot=[0.0, 0.0, 0.0], ) ) + sim.prepare() mg = MotionGenerator( MotionGenCfg( planner_cfg=CuroboPlannerCfg( diff --git a/tests/sim/atomic_actions/test_motion_strategy_e2e.py b/tests/sim/atomic_actions/test_motion_strategy_e2e.py index 38bff128c..3964916ad 100644 --- a/tests/sim/atomic_actions/test_motion_strategy_e2e.py +++ b/tests/sim/atomic_actions/test_motion_strategy_e2e.py @@ -53,6 +53,7 @@ def _setup(self): } ) ) + sim.prepare() mg = MotionGenerator( MotionGenCfg(planner_cfg=ToppraPlannerCfg(robot_uid=self.ROBOT_UID)) ) diff --git a/tests/sim/atomic_actions/test_tutorial_utils.py b/tests/sim/atomic_actions/test_tutorial_utils.py index 131518d77..6f35bf44e 100644 --- a/tests/sim/atomic_actions/test_tutorial_utils.py +++ b/tests/sim/atomic_actions/test_tutorial_utils.py @@ -31,6 +31,12 @@ import torch from embodichain.lab.sim.atomic_actions import TimedTrajectory +from embodichain.lab.sim.cfg import ( + DefaultPhysicsCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + RigidBodyMaterialCfg, +) from scripts.tutorials.atomic_action.dynamic_obstacle_recovery import ( _animate_obstacle_to_pose, _blocking_obstacle_pose, @@ -44,15 +50,21 @@ create_dual_tutorial_robot_cfg, ) from scripts.tutorials.atomic_action.tutorial_utils import ( + NEWTON_GRASP_CONTACT_DAMPING, + NEWTON_GRASP_CONTACT_STIFFNESS, ROBOTIQ_2F_140_TCP, ROBOTIQ_HAND_JOINT_PATTERN, TUTORIAL_ROBOTS, + add_tutorial_robot, broadcast_pose_batch, broadcast_waypoint_pose_batch, clone_local_pose_from_first_env, + configure_newton_gripper_contacts, + configure_newton_link_contacts, create_antipodal_semantics, create_curobo_motion_generator, create_franka_panda_robot_cfg, + create_tutorial_rigid_body_physics, create_tutorial_argument_parser, create_tutorial_robot_cfg, create_ur10_robotiq_robot_cfg, @@ -60,6 +72,7 @@ create_parallel_jaw_grasp_pose_generator, get_hand_open_close_qpos, replay_trajectory, + run_tutorial, should_open_tutorial_window, should_wait_for_tutorial_input, ) @@ -153,6 +166,48 @@ def _run_obstacle_animation(*, pace_wall_time: bool) -> tuple[MagicMock, MagicMo return obstacle, adapter +def test_atomic_action_tutorial_uses_grasp_stable_newton_solver_settings() -> None: + module = importlib.import_module("scripts.tutorials.atomic_action.tutorial_utils") + + default_cfg = module._tutorial_physics_cfg("default") + newton_cfg = module._tutorial_physics_cfg("newton") + + assert isinstance(default_cfg, DefaultPhysicsCfg) + assert isinstance(newton_cfg, NewtonPhysicsCfg) + assert newton_cfg.num_substeps == 20 + assert newton_cfg.collision_cfg.reduce_contacts is True + assert newton_cfg.collision_cfg.rigid_contact_max == 16_384 + assert newton_cfg.collision_cfg.broad_phase == "nxn" + assert newton_cfg.solver_cfg == { + "solver_type": "mujoco_warp", + "solver": "newton", + "integrator": "implicitfast", + "iterations": 15, + "ls_iterations": 100, + "nconmax": 16_384, + "njmax": 32_768, + "cone": "elliptic", + "impratio": 50.0, + "use_mujoco_contacts": False, + } + + dexsim_cfg = newton_cfg.to_dexsim_cfg(gpu_id=0) + assert dexsim_cfg.solver_cfg.solver_type == "mujoco_warp" + assert dexsim_cfg.solver_cfg.solver == "newton" + assert dexsim_cfg.solver_cfg.integrator == "implicitfast" + assert dexsim_cfg.solver_cfg.iterations == 15 + assert dexsim_cfg.solver_cfg.ls_iterations == 100 + assert dexsim_cfg.solver_cfg.nconmax == 16_384 + assert dexsim_cfg.solver_cfg.njmax == 32_768 + assert dexsim_cfg.solver_cfg.cone == "elliptic" + assert dexsim_cfg.solver_cfg.impratio == pytest.approx(50.0) + assert dexsim_cfg.solver_cfg.use_mujoco_contacts is False + assert dexsim_cfg.solver_cfg.enable_multiccd is False + assert dexsim_cfg.collision_pipeline_cfg.reduce_contacts is True + assert dexsim_cfg.collision_pipeline_cfg.rigid_contact_max == 16_384 + assert dexsim_cfg.collision_pipeline_cfg.broad_phase == "nxn" + + def test_should_wait_for_tutorial_input_is_disabled_for_headless_modes() -> None: assert ( should_wait_for_tutorial_input( @@ -297,8 +352,8 @@ def test_franka_tutorial_config_uses_ur5_gripper_component() -> None: assert franka_cfg.init_qpos[-2:] == [0.0, 0.0] assert franka_cfg.init_rot == FRANKA_TUTORIAL_BASE_ROTATION for property_name in ("stiffness", "damping", "max_effort"): - ur5_values = getattr(ur5_cfg.drive_pros, property_name) - franka_values = getattr(franka_cfg.drive_pros, property_name) + ur5_values = getattr(ur5_cfg.joint_drive_props, property_name) + franka_values = getattr(franka_cfg.joint_drive_props, property_name) assert franka_values["gripper_finger1_joint_1"] == ( ur5_values["gripper_finger1_joint_1"] ) @@ -477,12 +532,95 @@ def test_curobo_motion_generator_factory_selects_curobo_backend() -> None: with patch( "scripts.tutorials.atomic_action.tutorial_utils.MotionGenerator" ) as motion_generator_cls: - result = create_curobo_motion_generator(robot) + result = create_curobo_motion_generator(robot, use_cuda_graph=False) cfg = motion_generator_cls.call_args.kwargs["cfg"] assert result is motion_generator_cls.return_value assert cfg.planner_cfg.planner_type == "curobo" assert cfg.planner_cfg.robot_uid == "tutorial_robot" + assert cfg.planner_cfg.use_cuda_graph is False + + +def test_tutorial_rigid_body_physics_groups_backend_specific_properties() -> None: + physics = create_tutorial_rigid_body_physics( + mass=0.05, + static_friction=0.8, + dynamic_friction=0.4, + restitution=0.1, + linear_damping=0.2, + angular_damping=0.3, + max_depenetration_velocity=1.5, + enable_ccd=True, + min_position_iters=4, + min_velocity_iters=2, + contact_offset=0.01, + rest_offset=0.001, + ) + + assert physics.mass_props.mass == 0.05 + assert physics.material_props.static_friction == 0.8 + assert physics.material_props.dynamic_friction == 0.4 + assert physics.material_props.restitution == 0.1 + assert physics.rigid_props.linear_damping == 0.2 + assert physics.rigid_props.angular_damping == 0.3 + assert physics.rigid_props.max_depenetration_velocity == 1.5 + assert physics.rigid_props.enable_ccd is True + assert physics.rigid_props.min_position_iters == 4 + assert physics.rigid_props.min_velocity_iters == 2 + assert physics.collision_props.contact_offset == 0.01 + assert physics.collision_props.rest_offset == 0.001 + + +def test_tutorial_rigid_body_physics_adds_only_newton_contact_response() -> None: + empty_physics = create_tutorial_rigid_body_physics() + default_physics = create_tutorial_rigid_body_physics( + static_friction=0.8, + dynamic_friction=0.4, + ) + newton_physics = create_tutorial_rigid_body_physics( + static_friction=0.8, + dynamic_friction=0.4, + newton_contact=True, + ) + + assert empty_physics.material_props is None + assert type(default_physics.material_props) is RigidBodyMaterialCfg + assert type(newton_physics.material_props) is NewtonRigidBodyMaterialCfg + assert newton_physics.material_props.static_friction == pytest.approx(0.8) + assert newton_physics.material_props.dynamic_friction == pytest.approx(0.4) + assert newton_physics.material_props.ke == pytest.approx( + NEWTON_GRASP_CONTACT_STIFFNESS + ) + assert newton_physics.material_props.kd == pytest.approx( + NEWTON_GRASP_CONTACT_DAMPING + ) + + +def test_run_tutorial_uses_deferred_simulation_cleanup() -> None: + sim = MagicMock() + sim.is_window_recording.return_value = False + + with ( + patch( + "scripts.tutorials.atomic_action.tutorial_utils." + "SimulationManager.is_instantiated", + return_value=True, + ), + patch( + "scripts.tutorials.atomic_action.tutorial_utils." + "SimulationManager.get_instance", + return_value=sim, + ), + patch( + "scripts.tutorials.atomic_action.tutorial_utils." + "SimulationManager.flush_cleanup_queue" + ) as flush_cleanup_queue, + ): + run_tutorial(lambda: None) + + sim.wait_window_record_saves.assert_called_once_with() + sim.destroy.assert_called_once_with(exit_process=False) + flush_cleanup_queue.assert_called_once_with() def test_shared_robot_selection_keeps_ur5_default_and_accepts_all_variants() -> None: @@ -543,6 +681,7 @@ def test_place_tutorial_registers_pick_object_with_simulation_engine_factory() - ) sim = MagicMock() sim.device = torch.device("cpu") + sim.is_newton_backend = False sim.sim_config.physics_dt = PHYSICS_DT robot = MagicMock() robot.get_qpos.return_value = torch.zeros(1, 8) @@ -582,6 +721,164 @@ def test_place_tutorial_registers_pick_object_with_simulation_engine_factory() - engine.initial_context.assert_called_once_with(control_dt=PHYSICS_DT) +@pytest.mark.parametrize( + "link_name", + ( + "gripper_finger1_link_1", + "left_gripper_finger2_link_1", + "right_gripper_finger1_link_1", + "left_inner_finger_pad", + "right_left_outer_knuckle", + ), +) +def test_shared_tutorial_gripper_uses_newton_contact_material( + link_name: str, +) -> None: + sim = SimpleNamespace(is_newton_backend=True) + robot_cfg = SimpleNamespace(link_attrs=None) + + configure_newton_gripper_contacts(sim, robot_cfg) + + override = robot_cfg.link_attrs["newton_gripper_contacts"] + material = override.attrs.material_props + assert isinstance(material, NewtonRigidBodyMaterialCfg) + assert material.ke == pytest.approx(NEWTON_GRASP_CONTACT_STIFFNESS) + assert material.kd == pytest.approx(NEWTON_GRASP_CONTACT_DAMPING) + assert re.fullmatch(override.link_names_expr[0], link_name) + + +def test_shared_tutorial_preserves_default_gripper_contact_config() -> None: + existing_link_attrs = {"existing": MagicMock()} + sim = SimpleNamespace(is_newton_backend=False) + robot_cfg = SimpleNamespace(link_attrs=existing_link_attrs) + + configure_newton_gripper_contacts(sim, robot_cfg) + + assert robot_cfg.link_attrs is existing_link_attrs + + +def test_add_tutorial_robot_authors_newton_contacts_before_spawn() -> None: + sim = MagicMock() + sim.is_newton_backend = True + robot_cfg = SimpleNamespace(link_attrs=None) + + with patch( + "scripts.tutorials.atomic_action.tutorial_utils.create_tutorial_robot_cfg", + return_value=robot_cfg, + ): + result = add_tutorial_robot(sim, "ur5") + + assert result is sim.add_robot.return_value + sim.add_robot.assert_called_once_with(cfg=robot_cfg) + material = robot_cfg.link_attrs["newton_gripper_contacts"].attrs.material_props + assert material.ke == pytest.approx(NEWTON_GRASP_CONTACT_STIFFNESS) + assert material.kd == pytest.approx(NEWTON_GRASP_CONTACT_DAMPING) + + +def test_shared_tutorial_tunes_selected_newton_articulation_link() -> None: + sim = SimpleNamespace(is_newton_backend=True) + articulation_cfg = SimpleNamespace(link_attrs={"existing": MagicMock()}) + + configure_newton_link_contacts( + sim, + articulation_cfg, + group_name="newton_handle_contacts", + link_names_expr=["door_handle"], + ) + + assert "existing" in articulation_cfg.link_attrs + override = articulation_cfg.link_attrs["newton_handle_contacts"] + assert override.link_names_expr == ["door_handle"] + assert override.attrs.material_props.ke == pytest.approx( + NEWTON_GRASP_CONTACT_STIFFNESS + ) + assert override.attrs.material_props.kd == pytest.approx( + NEWTON_GRASP_CONTACT_DAMPING + ) + + +@pytest.mark.parametrize( + ("module_name", "factory_name", "group_name", "contact_link"), + ( + ("slide", "create_drawer", "newton_handle_contacts", "large_handle_bar"), + ("open_door", "create_microwave", "newton_handle_contacts", "door_handle"), + ("twist", "create_microwave", "newton_knob_contacts", "cap_1"), + ("press", "create_microwave", "newton_button_contacts", "button_cap"), + ), +) +def test_articulation_contact_tutorials_author_newton_material_before_spawn( + module_name: str, + factory_name: str, + group_name: str, + contact_link: str, +) -> None: + module = importlib.import_module(f"scripts.tutorials.atomic_action.{module_name}") + sim = MagicMock() + sim.is_newton_backend = True + + with patch.object(module, "get_data_path", return_value="/tmp/tutorial.urdf"): + result = getattr(module, factory_name)(sim) + + assert result is sim.add_articulation.return_value + cfg = sim.add_articulation.call_args.kwargs["cfg"] + assert cfg.asset_physics_mode == "overlay" + assert cfg.root_props.fixed_base is True + override = cfg.link_attrs[group_name] + assert override.link_names_expr == [contact_link] + assert override.attrs.material_props.ke == pytest.approx( + NEWTON_GRASP_CONTACT_STIFFNESS + ) + assert override.attrs.material_props.kd == pytest.approx( + NEWTON_GRASP_CONTACT_DAMPING + ) + + +@pytest.mark.parametrize( + ("is_newton_backend", "expected_material_type"), + ( + (False, RigidBodyMaterialCfg), + (True, NewtonRigidBodyMaterialCfg), + ), +) +def test_place_cube_uses_backend_scoped_contact_material( + is_newton_backend: bool, + expected_material_type: type[RigidBodyMaterialCfg], +) -> None: + module = importlib.import_module("scripts.tutorials.atomic_action.place") + sim = MagicMock() + sim.is_newton_backend = is_newton_backend + obj = MagicMock() + sim.add_rigid_object.return_value = obj + + with patch.object(module, "clone_local_pose_from_first_env"): + result = module.create_pick_object(sim) + + cfg = sim.add_rigid_object.call_args.kwargs["cfg"] + material = cfg.attrs.material_props + assert type(material) is expected_material_type + assert material.dynamic_friction == pytest.approx(0.97) + assert material.static_friction == pytest.approx(0.99) + if is_newton_backend: + assert material.ke == pytest.approx(NEWTON_GRASP_CONTACT_STIFFNESS) + assert material.kd == pytest.approx(NEWTON_GRASP_CONTACT_DAMPING) + result.clear_dynamics.assert_called_once_with() + + +def test_move_held_object_cup_starts_above_the_ground() -> None: + module = importlib.import_module("scripts.tutorials.atomic_action.move_held_object") + sim = MagicMock() + sim.is_newton_backend = True + obj = MagicMock() + sim.add_rigid_object.return_value = obj + + with patch.object(module, "clone_local_pose_from_first_env"): + module.create_pick_object(sim) + + cfg = sim.add_rigid_object.call_args.kwargs["cfg"] + assert cfg.init_pos == [*module.OBJECT_XY, module.OBJECT_INITIAL_Z] + assert module.OBJECT_INITIAL_Z == pytest.approx(0.05) + + def test_atomic_action_tutorial_scene_strategies_cover_every_entry_point() -> None: classified = ( set(RIGID_SCENE_TUTORIAL_MODULES) diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index 03f07983c..8fb3673ba 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -19,6 +19,7 @@ import os from types import SimpleNamespace +import numpy as np import pytest import torch @@ -32,15 +33,24 @@ ArticulationCfg, JointDrivePropertiesCfg, LinkPhysicsOverrideCfg, - RigidBodyAttributesCfg, - RigidBodyAttributesOverrideCfg, + MassPropertiesCfg, + physics_cfg_for_backend, + RigidBodyPhysicsCfg, ) -from embodichain.lab.sim.utility.sim_utils import _resolve_link_physics_groups from embodichain.data import get_data_path from dexsim.types import ActorType, DriveType ART_PATH = "SlidingBoxDrawer/SlidingBoxDrawer.urdf" NUM_ARENAS = 10 +NEWTON_EFFORT_TARGET_MODE = 4 +DRIVE_TEST_STIFFNESS = 12.0 +DRIVE_TEST_DAMPING = 4.0 + + +def _teardown_newton_physics() -> None: + from dexsim.engine.newton_physics import teardown_newton_physics + + teardown_newton_physics() def test_get_qf_returns_all_articulation_joint_efforts(): @@ -97,7 +107,9 @@ def test_get_parent_joint_chain_returns_backend_neutral_child_to_root_values(): 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 + return art.get_link_physical_attr(link_names=[link_name], env_ids=[env_idx])[ + 0 + ].static_friction class _EntityMethodOverride: @@ -114,56 +126,61 @@ def __getattr__(self, name: str): return getattr(self._entity, name) -class TestRigidBodyAttributesOverride: +class TestLinkPhysicsOverrideCfg: """Pure-Python tests for per-link physics config merging.""" - def test_merge_with_applies_only_set_fields(self): - base = RigidBodyAttributesCfg( - static_friction=0.3, - dynamic_friction=0.25, - linear_damping=0.5, + def test_grouped_override_applies_only_configured_fields(self): + base = RigidBodyPhysicsCfg.from_dict( + { + "rigid_props": {"linear_damping": 0.5}, + "material_props": { + "static_friction": 0.3, + "dynamic_friction": 0.25, + }, + } ) - override = RigidBodyAttributesOverrideCfg(static_friction=0.85) - merged = override.merge_with(base) + override = RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": 0.85}} + ) + merged = override.to_dexsim_physical_attr(base=base.to_dexsim_physical_attr()) assert abs(merged.static_friction - 0.85) < 1e-6 assert abs(merged.dynamic_friction - 0.25) < 1e-6 assert abs(merged.linear_damping - 0.5) < 1e-6 - def test_resolve_link_physics_overlap_raises(self): - link_names = ["outer_box", "handle_xpos", "inner_drawer"] - link_attrs = { - "box": LinkPhysicsOverrideCfg( - link_names_expr=["outer_box", "handle_xpos"], - attrs=RigidBodyAttributesOverrideCfg(static_friction=0.9), - ), - "handle": LinkPhysicsOverrideCfg( - link_names_expr=["handle_xpos"], - attrs=RigidBodyAttributesOverrideCfg(static_friction=0.8), - ), - } - with pytest.raises(ValueError, match="multiple link_attrs groups"): - _resolve_link_physics_groups(link_names, link_attrs) - 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) - cfg_dict = {"fpath": art_path, "drive_pros": {"drive_type": "force"}} + cfg_dict = { + "fpath": art_path, + "asset_physics_mode": "overlay", + "joint_drive_props": {"drive_type": "force"}, + } self.art: Articulation = self.sim.add_articulation( cfg=ArticulationCfg.from_dict(cfg_dict) ) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): - self.sim.init_gpu_physics() + self.sim.prepare() def test_local_pose_behavior(self): """Test set_local_pose and get_local_pose: @@ -171,21 +188,148 @@ def test_local_pose_behavior(self): """ # Set initial poses - pose = torch.eye(4, device=self.sim.device) - pose[2, 3] = 1.0 - pose = pose.unsqueeze(0).repeat(NUM_ARENAS, 1, 1) + distinct_xyzw = torch.tensor( + [1.0, 2.0, 3.0, 4.0], device=self.sim.device + ) / torch.sqrt(torch.tensor(30.0, device=self.sim.device)) + pose = torch.zeros(NUM_ARENAS, 7, device=self.sim.device) + pose[:, 2] = 1.0 + pose[:, 3:7] = distinct_xyzw self.art.set_local_pose(pose, env_ids=None) # --- Check poses immediately after setting - xyz = self.art.get_local_pose()[0, :3] - + actual_pose = self.art.get_local_pose() + xyz = actual_pose[0, :3] expected_pos = torch.tensor( [0.0, 0.0, 1.0], device=self.sim.device, dtype=torch.float32 ) assert torch.allclose( xyz, expected_pos, atol=1e-5 ), f"FAIL: Drawer pose not set correctly: {xyz.tolist()}" + torch.testing.assert_close( + actual_pose[:, 3:7], + distinct_xyzw.unsqueeze(0).expand(NUM_ARENAS, -1), + atol=1e-5, + rtol=1e-5, + ) + + def test_replicated_link_shapes_are_isolated_by_environment(self): + """Every articulation link shape should use its environment group.""" + for env_index, entity in enumerate(self.art._entities): + if self.physics == "newton": + shape_ids = [ + shape_id + for link in entity.physics_articulation.links + for shape_id in link.shape_ids + ] + assert shape_ids + groups = ( + entity.physics_articulation.runtime.model.shape_collision_group.numpy() + ) + assert {int(groups[shape_id]) for shape_id in shape_ids} == { + env_index + 1 + } + continue + + expected = np.asarray([env_index, 1, 0, 0], dtype=np.uint32) + physical_links = [ + link + for link in entity.articulation_desc.links + if link.rigid_body is not None + ] + assert physical_links + for link in physical_links: + np.testing.assert_array_equal( + link.rigid_body.collision_filter_data, + expected, + ) + + def test_body_data_exposes_link_mass_properties(self): + """Current and initialization-time link mass properties share one layout.""" + data = self.art.body_data + + assert data.mass.shape == (NUM_ARENAS, self.art.num_links) + assert data.inertia.shape == (NUM_ARENAS, self.art.num_links, 3) + assert data.com_pose.shape == (NUM_ARENAS, self.art.num_links, 7) + assert data.default_mass.shape == data.mass.shape + assert data.default_inertia.shape == data.inertia.shape + assert data.default_com_pose.shape == data.com_pose.shape + assert torch.allclose(self.art.default_link_masses, data.default_mass) + + def test_reset_restores_default_link_mass_properties(self): + """Partial reset restores mass, inertia, and COM only for selected rows.""" + data = self.art.body_data + link_name = self.art.link_names[0] + link_id = self.art.link_names.index(link_name) + env_ids = [0, 1] + default_mass = data.default_mass[env_ids, link_id : link_id + 1].clone() + default_inertia = data.default_inertia[env_ids, link_id : link_id + 1].clone() + default_com_pose = data.default_com_pose[env_ids, link_id : link_id + 1].clone() + changed_mass = default_mass + 0.5 + changed_inertia = default_inertia * 1.25 + changed_com_pose = default_com_pose.clone() + changed_com_pose[..., 0] += 0.02 + changed_com_pose[..., 3:7] = torch.tensor( + [1.0, 2.0, 3.0, 4.0], device=self.sim.device + ) / torch.sqrt(torch.tensor(30.0, device=self.sim.device)) + + self.art.set_mass(changed_mass, link_names=[link_name], env_ids=env_ids) + self.art.set_inertia( + changed_inertia, + link_names=[link_name], + env_ids=env_ids, + ) + self.art.set_com_pose( + changed_com_pose, + link_names=[link_name], + env_ids=env_ids, + ) + self.sim.prepare() + + assert torch.allclose( + data.default_mass[env_ids, link_id : link_id + 1], default_mass + ) + assert torch.allclose( + data.default_inertia[env_ids, link_id : link_id + 1], default_inertia + ) + assert torch.allclose( + data.default_com_pose[env_ids, link_id : link_id + 1], default_com_pose + ) + + self.art.reset(env_ids=[env_ids[0]]) + self.sim.prepare() + mass_after_partial = self.art.get_mass(link_names=[link_name], env_ids=env_ids) + inertia_after_partial = self.art.get_inertia( + link_names=[link_name], env_ids=env_ids + ) + com_after_partial = self.art.get_com_pose( + link_names=[link_name], env_ids=env_ids + ) + + assert torch.allclose(mass_after_partial[0], default_mass[0], atol=1e-5) + assert torch.allclose(mass_after_partial[1], changed_mass[1], atol=1e-5) + assert torch.allclose(inertia_after_partial[0], default_inertia[0], atol=1e-5) + assert torch.allclose(inertia_after_partial[1], changed_inertia[1], atol=1e-5) + assert torch.allclose(com_after_partial[0], default_com_pose[0], atol=1e-5) + assert torch.allclose(com_after_partial[1], changed_com_pose[1], atol=1e-5) + + self.art.reset(env_ids=[env_ids[1]]) + self.sim.prepare() + assert torch.allclose( + self.art.get_mass(link_names=[link_name], env_ids=env_ids), + default_mass, + atol=1e-5, + ) + assert torch.allclose( + self.art.get_inertia(link_names=[link_name], env_ids=env_ids), + default_inertia, + atol=1e-5, + ) + assert torch.allclose( + self.art.get_com_pose(link_names=[link_name], env_ids=env_ids), + default_com_pose, + atol=1e-5, + ) def test_control_api(self): """Test control API for setting and getting joint positions.""" @@ -368,12 +512,14 @@ def test_get_joint_drive_with_joint_ids(self): armature, expected_armature, atol=1e-5 ), "FAIL: armature does not match expected filtered values" - def test_default_drive_type_is_none_after_construction(self): - """A default ArticulationCfg creates passive backend joint drives.""" + def test_explicit_passive_drive_after_construction(self): + """An explicit passive overlay disables backend joint drives.""" passive_articulation = self.sim.add_articulation( cfg=ArticulationCfg( uid="passive_drawer", fpath=get_data_path(ART_PATH), + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="none"), ) ) @@ -382,6 +528,50 @@ def test_default_drive_type_is_none_after_construction(self): ] assert passive_articulation.get_joint_drive_type() == expected_drive_types + if self.sim.is_newton_backend: + expected_target_modes = [ + [0] * passive_articulation.dof for _ in range(NUM_ARENAS) + ] + assert passive_articulation.get_joint_target_mode() == expected_target_modes + + def test_preserve_mode_ignores_urdf_physics_overrides(self): + """Preserve mode keeps source-resolved URDF link and joint physics.""" + source = self.sim.add_articulation( + cfg=ArticulationCfg( + uid="source_drawer", + fpath=get_data_path(ART_PATH), + asset_physics_mode="preserve", + init_pos=(-1.0, 0.0, 0.0), + ) + ) + preserved = self.sim.add_articulation( + cfg=ArticulationCfg( + uid="preserved_drawer", + fpath=get_data_path(ART_PATH), + asset_physics_mode="preserve", + init_pos=(1.0, 0.0, 0.0), + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=123.0)), + joint_drive_props=JointDrivePropertiesCfg( + drive_type="none", + stiffness=987.0, + damping=654.0, + max_effort=321.0, + max_velocity=123.0, + ), + qpos_limits={".*": [-0.01, 0.01]}, + ) + ) + + assert torch.allclose(preserved.body_data.mass, source.body_data.mass) + assert torch.allclose( + preserved.body_data.qpos_limits, + source.body_data.qpos_limits, + ) + for preserved_value, source_value in zip( + preserved.get_joint_drive(), source.get_joint_drive() + ): + assert torch.allclose(preserved_value, source_value) + def test_joint_limit_getters_support_env_and_joint_filters(self): """Test joint limit getters support joint_ids and env_ids filtering.""" all_qpos_limits = self.art.body_data.qpos_limits @@ -808,7 +998,8 @@ def test_qpos_limits_from_cfg_dict_can_tighten(self): cfg = ArticulationCfg( uid="drawer_cfg_qpos_limits", fpath=get_data_path(ART_PATH), - drive_pros=JointDrivePropertiesCfg(drive_type="force"), + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), qpos_limits={".*": [-0.05, 0.05]}, ) art: Articulation = self.sim.add_articulation(cfg=cfg) @@ -832,7 +1023,8 @@ def test_qpos_limits_from_cfg_can_expand(self): cfg = ArticulationCfg( uid="drawer_expanded_limits", fpath=get_data_path(ART_PATH), - drive_pros=JointDrivePropertiesCfg(drive_type="force"), + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), qpos_limits={joint_name: [expanded_lower, expanded_upper]}, ) art: Articulation = self.sim.add_articulation(cfg=cfg) @@ -874,8 +1066,8 @@ def teardown_method(self): 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) + def setup_simulation(self, device: str) -> None: + config = SimulationManagerCfg(headless=True, device=device, num_envs=2) self.sim = SimulationManager(config) self.art_path = get_data_path(ART_PATH) assert os.path.isfile(self.art_path) @@ -896,10 +1088,14 @@ def test_global_attrs_applied_to_all_links(self): cfg = ArticulationCfg( uid="drawer_global_attrs", fpath=self.art_path, - drive_pros=JointDrivePropertiesCfg(drive_type="force"), - attrs=RigidBodyAttributesCfg(static_friction=global_friction), + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), + attrs=RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": global_friction}} + ), ) art: Articulation = self.sim.add_articulation(cfg=cfg) + self.sim.prepare() for link_name in art.link_names: assert abs(_link_static_friction(art, link_name) - global_friction) < 1e-3 @@ -910,18 +1106,22 @@ def test_link_attrs_override_selected_links(self): cfg = ArticulationCfg( uid="drawer_link_attrs", fpath=self.art_path, - drive_pros=JointDrivePropertiesCfg(drive_type="force"), - attrs=RigidBodyAttributesCfg(static_friction=global_friction), + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), + attrs=RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": global_friction}} + ), link_attrs={ "handle": LinkPhysicsOverrideCfg( link_names_expr=["handle_xpos"], - attrs=RigidBodyAttributesOverrideCfg( - static_friction=handle_friction + attrs=RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": handle_friction}} ), ), }, ) art: Articulation = self.sim.add_articulation(cfg=cfg) + self.sim.prepare() assert abs(_link_static_friction(art, "handle_xpos") - handle_friction) < 1e-3 for link_name in art.link_names: if link_name == "handle_xpos": @@ -934,17 +1134,19 @@ def test_link_attrs_from_dict(self): { "uid": "drawer_link_attrs_dict", "fpath": self.art_path, - "drive_pros": {"drive_type": "force"}, - "attrs": {"static_friction": 0.4}, + "asset_physics_mode": "overlay", + "joint_drive_props": {"drive_type": "force"}, + "attrs": {"material_props": {"static_friction": 0.4}}, "link_attrs": { "handle": { "link_names_expr": ["handle_xpos"], - "attrs": {"static_friction": 0.77}, + "attrs": {"material_props": {"static_friction": 0.77}}, } }, } ) art: Articulation = self.sim.add_articulation(cfg=cfg) + self.sim.prepare() assert abs(_link_static_friction(art, "handle_xpos") - 0.77) < 1e-3 assert abs(_link_static_friction(art, "outer_box") - 0.4) < 1e-3 @@ -953,19 +1155,31 @@ def test_set_link_physical_attr_runtime(self): cfg = ArticulationCfg( uid="drawer_runtime_attrs", fpath=self.art_path, - drive_pros=JointDrivePropertiesCfg(drive_type="force"), + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), ) art: Articulation = self.sim.add_articulation(cfg=cfg) + self.sim.prepare() + source_friction = { + link_name: _link_static_friction(art, link_name) + for link_name in art.link_names + } handle_friction = 0.66 art.set_link_physical_attr( - RigidBodyAttributesOverrideCfg(static_friction=handle_friction), + RigidBodyPhysicsCfg.from_dict( + {"material_props": {"static_friction": handle_friction}} + ), link_names=["handle_xpos"], ) + self.sim.prepare() assert abs(_link_static_friction(art, "handle_xpos") - handle_friction) < 1e-3 for link_name in art.link_names: if link_name == "handle_xpos": continue - assert abs(_link_static_friction(art, link_name) - 0.5) < 1e-3 + assert ( + abs(_link_static_friction(art, link_name) - source_friction[link_name]) + < 1e-3 + ) class TestArticulationLinkPhysicsCPU(BaseArticulationLinkPhysicsTest): @@ -988,6 +1202,117 @@ 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.gpu + def test_runtime_effort_drive_mode(self): + """Newton authors effort mode and removes effective PD gains.""" + shape = (NUM_ARENAS, self.art.dof) + self.art.set_joint_drive( + stiffness=torch.full( + shape, + DRIVE_TEST_STIFFNESS, + dtype=torch.float32, + device=self.sim.device, + ), + damping=torch.full( + shape, + DRIVE_TEST_DAMPING, + dtype=torch.float32, + device=self.sim.device, + ), + drive_type="force", + target_mode="effort", + ) + + assert self.art.get_joint_target_mode() == [ + [NEWTON_EFFORT_TARGET_MODE] * self.art.dof for _ in range(NUM_ARENAS) + ] + stiffness, damping, *_ = self.art.get_joint_drive() + assert torch.count_nonzero(stiffness) == 0 + assert torch.count_nonzero(damping) == 0 + + @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_mass_rebuilds_mass_on_newton(self): + """A retained Newton per-link mass takes effect at prepare().""" + 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_mass( + torch.full( + (NUM_ARENAS, 1), + new_mass, + dtype=torch.float32, + device=self.sim.device, + ), + link_names=[link_name], + ) + self.sim.prepare() + 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 after Newton rebuild (got {live_mass})" + + if __name__ == "__main__": test = TestArticulationCPU() test.setup_method() diff --git a/tests/sim/objects/test_articulation_drive_compat.py b/tests/sim/objects/test_articulation_drive_compat.py new file mode 100644 index 000000000..7bf03ca19 --- /dev/null +++ b/tests/sim/objects/test_articulation_drive_compat.py @@ -0,0 +1,90 @@ +# ---------------------------------------------------------------------------- +# 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 numpy as np +import pytest +import torch +from dexsim.types import DriveType + +from embodichain.lab.sim.objects.articulation import Articulation + +pytestmark = pytest.mark.no_sim + + +def test_newton_target_modes_map_to_portable_drive_types() -> None: + target_modes = np.asarray([0, 1, 2, 3, 4], dtype=np.int32) + entity = SimpleNamespace( + get_newton_drive=lambda: (None, None, None, None, None, None, target_modes) + ) + articulation = object.__new__(Articulation) + articulation._data = SimpleNamespace( + is_newton_backend=True, + dof=len(target_modes), + ) + articulation._all_indices = np.asarray([0], dtype=np.int32) + articulation._entities = [entity] + + assert articulation.get_joint_drive_type() == [ + [ + DriveType.NONE, + DriveType.FORCE, + DriveType.FORCE, + DriveType.FORCE, + DriveType.NONE, + ] + ] + + +def test_newton_drive_type_query_honors_joint_selection() -> None: + target_modes = np.asarray([0, 3, 0], dtype=np.int32) + entity = SimpleNamespace( + get_newton_drive=lambda: (None, None, None, None, None, None, target_modes) + ) + articulation = object.__new__(Articulation) + articulation._data = SimpleNamespace(is_newton_backend=True, dof=3) + articulation._all_indices = np.asarray([0], dtype=np.int32) + articulation._entities = [entity] + + assert articulation.get_joint_drive_type(joint_ids=[2, 1]) == [ + [DriveType.NONE, DriveType.FORCE] + ] + + +def test_runtime_effort_mode_disables_pd_gains_on_newton() -> None: + calls: list[dict[str, object]] = [] + entity = SimpleNamespace(set_newton_drive=lambda **kwargs: calls.append(kwargs)) + articulation = object.__new__(Articulation) + articulation._spawn_result = object() + articulation._entities = [entity] + articulation._all_indices = np.asarray([0], dtype=np.int32) + articulation._data = SimpleNamespace(is_newton_backend=True, dof=1) + articulation.device = torch.device("cpu") + + articulation.set_joint_drive( + stiffness=torch.tensor([[12.0]]), + damping=torch.tensor([[4.0]]), + drive_type="force", + target_mode="effort", + ) + + assert len(calls) == 1 + assert calls[0]["target_mode"] == 4 + assert calls[0]["target_ke"] == 0.0 + assert calls[0]["target_kd"] == 0.0 diff --git a/tests/sim/objects/test_asset_material_initialization.py b/tests/sim/objects/test_asset_material_initialization.py index 6602e811b..c46f64784 100644 --- a/tests/sim/objects/test_asset_material_initialization.py +++ b/tests/sim/objects/test_asset_material_initialization.py @@ -59,6 +59,7 @@ def _make_asset(asset_type, materials): asset = asset_type.__new__(asset_type) asset._entities = [entity] + asset._spawn_result = None asset._all_indices = [0] asset.is_shared_visual_material = False asset.uid = asset_type.__name__ @@ -191,10 +192,14 @@ def test_asset_restores_only_changed_segments(asset_type): def test_asset_reset_restores_selected_environment_material(asset_type): asset = asset_type.__new__(asset_type) + asset._entities = [MagicMock(name="entity")] + asset._declared_num_instances = 1 + asset._spawn_result = MagicMock(name="spawn_result") asset._all_indices = [0] asset.device = torch.device("cpu") asset.cfg = SimpleNamespace( attrs=MagicMock(), + init_local_pose=None, init_pos=(0.0, 0.0, 0.0), init_rot=(0.0, 0.0, 0.0), init_qpos=(0.0,), @@ -203,9 +208,12 @@ def test_asset_reset_restores_selected_environment_material(asset_type): asset.set_local_pose = MagicMock() if asset_type is RigidObject: + asset._data = None asset.set_attrs = MagicMock() asset.clear_dynamics = MagicMock() elif asset_type is Articulation: + asset._data = MagicMock(is_newton_backend=True) + asset._restore_default_physical_properties = MagicMock() asset.set_qpos = MagicMock() asset.clear_dynamics = MagicMock() asset._world = MagicMock() diff --git a/tests/sim/objects/test_cloth_object.py b/tests/sim/objects/test_cloth_object.py index 8ddaecaa6..b5238f1a6 100644 --- a/tests/sim/objects/test_cloth_object.py +++ b/tests/sim/objects/test_cloth_object.py @@ -21,7 +21,12 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ClothPhysicalAttributesCfg from embodichain.lab.sim.shapes import MeshCfg -from embodichain.lab.sim.objects import ClothObjectCfg, ClothObject +from embodichain.lab.sim.objects import ( + ClothObject, + ClothObjectCfg, + DeformableObject, + SurfaceDeformableObject, +) import open3d as o3d import pytest import torch @@ -69,7 +74,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, ) @@ -108,9 +113,9 @@ def setup_simulation(self): ), ) ) + self.sim.prepare() def test_run_simulation(self): - self.sim.init_gpu_physics() for _ in range(100): self.sim.update(step=1) self.cloth.reset() @@ -118,10 +123,9 @@ def test_run_simulation(self): self.sim.update(step=1) def test_remove(self): - self.sim.remove_asset(self.cloth.uid) - assert ( - self.cloth.uid not in self.sim._soft_objects - ), "Cow UID still present after removal" + with pytest.raises(NotImplementedError, match="pending removal"): + self.sim.remove_asset(self.cloth.uid) + assert self.sim.get_deformable_object(self.cloth.uid) is self.cloth def test_get_current_vertex_positions(self): vertex_positions = self.cloth.get_current_vertex_position() @@ -133,7 +137,7 @@ def test_get_current_vertex_positions(self): def test_get_deformable_mesh_geometry(self): """Test current cloth vertices and matching surface triangles.""" - self.sim.init_gpu_physics() + self.sim.prepare() vertices = self.cloth.get_current_vertex_position() triangles = self.cloth.get_triangles(env_ids=[0]) @@ -141,6 +145,40 @@ def test_get_deformable_mesh_geometry(self): assert triangles.ndim == 3 and triangles.shape[0] == 1 assert int(triangles.max()) < vertices.shape[1] + def test_unified_deformable_contract(self): + self.sim.update(step=5) + assert isinstance(self.cloth, DeformableObject) + assert isinstance(self.cloth, SurfaceDeformableObject) + assert self.cloth.deformable_type == "surface" + assert self.sim.get_deformable_object("cloth") is self.cloth + assert self.sim.get_cloth_object("cloth") is self.cloth + assert self.sim.get_deformable_object_uid_list() == ["cloth"] + + positions = self.cloth.get_current_nodal_position() + velocities = self.cloth.get_current_nodal_velocity() + state = self.cloth.get_current_nodal_state() + default_state = self.cloth.get_default_nodal_state() + assert positions.shape[-1] == 3 + assert velocities.shape == positions.shape + assert state.shape == (*positions.shape[:-1], 6) + assert default_state.shape == state.shape + native_velocities = torch.stack( + [ + body.get_velocity_buffer()[:, :3].clone() + for body in self.cloth.body_data.cloth_bodies + ] + ) + assert torch.count_nonzero(native_velocities) > 0 + torch.testing.assert_close(velocities, native_velocities) + torch.testing.assert_close( + self.cloth.get_surface_vertices(), + self.cloth.get_current_vertex_position(), + ) + torch.testing.assert_close( + self.cloth.get_surface_triangles(env_ids=[0]), + self.cloth.get_triangles(env_ids=[0]), + ) + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/sim/objects/test_deformable_object.py b/tests/sim/objects/test_deformable_object.py new file mode 100644 index 000000000..2791fc076 --- /dev/null +++ b/tests/sim/objects/test_deformable_object.py @@ -0,0 +1,123 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Contract tests for the unified deformable-object API.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from embodichain.lab.sim.cfg import ( + ClothObjectCfg, + DeformableObjectCfg, + SoftObjectCfg, + SurfaceDeformableObjectCfg, + VolumeDeformableObjectCfg, +) +from embodichain.lab.sim.objects import ( + ClothBodyData, + ClothObject, + DeformableObject, + DeformableObjectData, + SoftBodyData, + SoftObject, + SurfaceDeformableData, + SurfaceDeformableObject, + VolumeDeformableData, + VolumeDeformableObject, +) +from embodichain.lab.sim.physics import DefaultPhysicsBackend, NewtonPhysicsBackend +from embodichain.lab.sim.sim_manager import SimulationManager + + +class _Data(DeformableObjectData): + def __init__(self) -> None: + self._pos = torch.tensor( + [[[0.0, 0.0, 0.0], [2.0, 4.0, 6.0]]], dtype=torch.float32 + ) + self._vel = torch.tensor( + [[[1.0, 2.0, 3.0], [3.0, 4.0, 5.0]]], dtype=torch.float32 + ) + + @property + def nodal_pos_w(self) -> torch.Tensor: + return self._pos + + @property + def nodal_vel_w(self) -> torch.Tensor: + return self._vel + + @property + def default_nodal_state_w(self) -> torch.Tensor: + return torch.cat((self._pos, torch.zeros_like(self._vel)), dim=-1) + + +def test_legacy_configs_specialize_common_deformable_config() -> None: + assert issubclass(SoftObjectCfg, VolumeDeformableObjectCfg) + assert issubclass(ClothObjectCfg, SurfaceDeformableObjectCfg) + assert issubclass(VolumeDeformableObjectCfg, DeformableObjectCfg) + assert issubclass(SurfaceDeformableObjectCfg, DeformableObjectCfg) + assert SoftObjectCfg().deformable_type == "volume" + assert ClothObjectCfg().deformable_type == "surface" + + +def test_legacy_objects_are_aliases_of_topology_specializations() -> None: + assert SoftObject is VolumeDeformableObject + assert ClothObject is SurfaceDeformableObject + assert SoftBodyData is VolumeDeformableData + assert ClothBodyData is SurfaceDeformableData + assert issubclass(SoftObject, DeformableObject) + assert issubclass(ClothObject, DeformableObject) + + +def test_common_data_contract_combines_and_derives_nodal_state() -> None: + data = _Data() + + assert data.nodal_state_w.shape == (1, 2, 6) + torch.testing.assert_close(data.nodal_state_w[..., :3], data.nodal_pos_w) + torch.testing.assert_close(data.nodal_state_w[..., 3:], data.nodal_vel_w) + torch.testing.assert_close(data.root_pos_w, torch.tensor([[1.0, 2.0, 3.0]])) + torch.testing.assert_close(data.root_vel_w, torch.tensor([[2.0, 3.0, 4.0]])) + + +def test_backend_capabilities_keep_newton_deformable_entry_disabled() -> None: + default = DefaultPhysicsBackend(SimpleNamespace()) + newton = NewtonPhysicsBackend(SimpleNamespace()) + + assert default.supports_volume_deformables + assert default.supports_surface_deformables + assert default.supports_soft_bodies + assert default.supports_cloth + assert not newton.supports_volume_deformables + assert not newton.supports_surface_deformables + assert not newton.supports_soft_bodies + assert not newton.supports_cloth + + +def test_manager_generic_and_legacy_getters_share_one_registry() -> None: + sim = object.__new__(SimulationManager) + volume = object.__new__(VolumeDeformableObject) + surface = object.__new__(SurfaceDeformableObject) + sim._deformable_objects = {"volume": volume, "surface": surface} + + assert sim.get_deformable_object("volume") is volume + assert sim.get_soft_object("volume") is volume + assert sim.get_cloth_object("surface") is surface + assert sim.get_deformable_object_uid_list() == ["volume", "surface"] + assert sim.get_soft_object_uid_list() == ["volume"] + assert sim.get_cloth_object_uid_list() == ["surface"] diff --git a/tests/sim/objects/test_dual_arm.py b/tests/sim/objects/test_dual_arm.py index d4d9febf5..571012d67 100644 --- a/tests/sim/objects/test_dual_arm.py +++ b/tests/sim/objects/test_dual_arm.py @@ -20,6 +20,7 @@ import numpy as np import pytest +from embodichain.lab.sim.cfg import NewtonJointDrivePropertiesCfg from embodichain.lab.sim.robots.dual_arm import ( DualArmRobotCfg, _transform_from_xyz_rpy, @@ -162,6 +163,29 @@ def test_build_dual_arm_dual_part_toggle(): assert "dual_arm" not in cfg.control_parts +def test_build_dual_arm_mirrors_newton_joint_overrides(): + base = URRobotCfg.from_dict({"robot_type": "ur5"}) + base.joint_drive_props = NewtonJointDrivePropertiesCfg( + stiffness={"joint[1-6]": 12.0}, + target_mode={"joint[1-6]": "position"}, + friction=0.2, + ) + mounts = resolve_mounts({"preset": "side_by_side", "separation": 0.6}) + + cfg = build_dual_arm_cfg(base, mounts) + + assert isinstance(cfg.joint_drive_props, NewtonJointDrivePropertiesCfg) + assert cfg.joint_drive_props.stiffness == { + "left_joint[1-6]": 12.0, + "right_joint[1-6]": 12.0, + } + assert cfg.joint_drive_props.target_mode == { + "left_joint[1-6]": "position", + "right_joint[1-6]": "position", + } + assert cfg.joint_drive_props.friction == 0.2 + + # --------------------------------------------------------------------------- # # DualArmRobotCfg from_dict + round-trip # --------------------------------------------------------------------------- # diff --git a/tests/sim/objects/test_light.py b/tests/sim/objects/test_light.py index 0b7bbd794..322d42430 100644 --- a/tests/sim/objects/test_light.py +++ b/tests/sim/objects/test_light.py @@ -25,7 +25,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 @@ -37,6 +37,7 @@ def setup_method(self): "uid": "point_light", } self.light = self.sim.add_light(cfg=LightCfg.from_dict(cfg_dict)) + self.sim.prepare() def test_set_color_with_env_ids(self): """Test set_color with and without env_ids.""" @@ -169,7 +170,7 @@ class TestLightTypes: @pytest.fixture(autouse=True) def setup(self): """Create a SimulationManager for each test.""" - config = SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=4) + config = SimulationManagerCfg(headless=True, device="cpu", num_envs=4) self.sim = SimulationManager(config) yield self.sim.destroy() @@ -214,9 +215,9 @@ def test_create_each_light_type(self, light_type, expected_num_instances): assert light.is_global, f"{light_type} should be a global light" def test_unknown_light_type_errors(self): - """Passing an invalid light_type raises RuntimeError.""" + """Passing an invalid light_type raises ValueError.""" cfg = LightCfg(uid="bad", light_type="invalid") - with pytest.raises(RuntimeError, match="Unsupported light type"): + with pytest.raises(ValueError, match="Unsupported light type"): self.sim.add_light(cfg=cfg) def test_mesh_light_empty_path_warns(self): diff --git a/tests/sim/objects/test_rigid_constraint.py b/tests/sim/objects/test_rigid_constraint.py index 9911135d3..6cd8bb626 100644 --- a/tests/sim/objects/test_rigid_constraint.py +++ b/tests/sim/objects/test_rigid_constraint.py @@ -263,8 +263,7 @@ def __init__(self, num_envs=4, arenas=None): self._robots = {} self._rigid_objects = {} self._rigid_object_groups = {} - self._soft_objects = {} - self._cloth_objects = {} + self._deformable_objects = {} self._articulations = {} self._constraints = {} self.device = torch.device("cpu") diff --git a/tests/sim/objects/test_rigid_object.py b/tests/sim/objects/test_rigid_object.py index 929810cae..d281ce595 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -13,41 +13,73 @@ # See the License for the specific language governing permissions and # limitations under the License. # ---------------------------------------------------------------------------- - from __future__ import annotations import os -import torch + +import numpy as np import pytest +import torch from embodichain.lab.sim import ( SimulationManager, SimulationManagerCfg, VisualMaterialCfg, ) -from embodichain.lab.sim.objects import RigidObject -from embodichain.lab.sim.cfg import RigidObjectCfg, RigidBodyAttributesCfg -from embodichain.lab.sim.shapes import MeshCfg from embodichain.data import get_data_path -from dexsim.types import ActorType - -from embodichain.lab.sim.cfg import RenderCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import ( + MassPropertiesCfg, + NewtonCollisionPropertiesCfg, + NewtonRigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + physics_cfg_for_backend, +) +from embodichain.lab.sim.objects import RigidObject +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg, MeshCollisionCfg +from embodichain.utils.math import matrix_from_quat DUCK_PATH = "ToyDuck/toy_duck.glb" TABLE_PATH = "ShopTableSimple/shop_table_simple.ply" CHAIR_PATH = "Chair/chair.glb" NUM_ARENAS = 2 Z_TRANSLATION = 2.0 +# Newton stores a full inertia tensor and converts it to/from the principal-frame +# diagonal in float32. The two quaternion rotations introduce small round-trip +# error for imported meshes whose COM frame is not axis-aligned. +NEWTON_INERTIA_ROUND_TRIP_ATOL = 2e-4 + + +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, device: str, physics: str = "default"): config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + headless=True, + device=device, + num_envs=NUM_ARENAS, + physics_cfg=physics_cfg_for_backend(physics), ) self.sim = SimulationManager(config) + self.physics = physics self.sim.enable_physics(False) duck_path = get_data_path(DUCK_PATH) assert os.path.isfile(duck_path) @@ -62,9 +94,7 @@ def setup_simulation(self, sim_device): "shape_type": "Mesh", "fpath": duck_path, }, - "attrs": { - "mass": 1.0, - }, + "attrs": {"mass_props": {"mass": 1.0}}, "body_type": "dynamic", } self.duck: RigidObject = self.sim.add_rigid_object( @@ -78,12 +108,13 @@ def setup_simulation(self, sim_device): self.chair: RigidObject = self.sim.add_rigid_object( cfg=RigidObjectCfg( - uid="chair", shape=MeshCfg(fpath=chair_path), body_type="kinematic" + uid="chair", + shape=MeshCfg(fpath=chair_path), + body_type="kinematic", ), ) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): - self.sim.init_gpu_physics() + self.sim.prepare() self.sim.enable_physics(True) @@ -95,6 +126,34 @@ def test_is_static(self): not self.chair.is_static ), "Chair should be kinematic but is marked static" + def test_replicated_collision_shapes_are_isolated_by_environment(self): + """Every rigid shape should use its replicated environment group.""" + for env_index in range(NUM_ARENAS): + for rigid_object in (self.duck, self.table, self.chair): + entity = rigid_object._entities[env_index] + if self.physics == "newton": + shape_ids = entity.physics_body.shape_ids + assert shape_ids + groups = ( + entity.physics_body.runtime.model.shape_collision_group.numpy() + ) + assert {int(groups[shape_id]) for shape_id in shape_ids} == { + env_index + 1 + } + else: + np.testing.assert_array_equal( + entity.object_desc.physics.collision_filter_data, + np.asarray([env_index, 1, 0, 0], dtype=np.uint32), + ) + + 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 {entity.get_name() for entity in self.duck._entities} == {"duck"} + assert len({entity.path for entity in self.duck._entities}) == NUM_ARENAS + def test_local_pose_behavior(self): """Test set_local_pose and get_local_pose: - duck pose is correctly set @@ -160,9 +219,32 @@ 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.chair.body_type == "kinematic" and 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_dynamic_pose_write_persists_across_physics_step(self): + """A dynamic pose reset must update Newton's FREE-joint state too.""" + target_xy = torch.tensor( + [[0.31, -0.27], [-0.42, 0.36]], + dtype=torch.float32, + device=self.sim.device, + ) + pose = torch.eye(4, device=self.sim.device).repeat(NUM_ARENAS, 1, 1) + pose[:, :2, 3] = target_xy + pose[:, 2, 3] = Z_TRANSLATION + + self.duck.set_local_pose(pose) + self.sim.update(step=1) + + torch.testing.assert_close( + self.duck.get_local_pose()[:, :2], + target_xy, + atol=1.0e-4, + rtol=0.0, + ) def test_add_force_torque(self): """Test that add_force applies force correctly to the duck object.""" @@ -333,7 +415,13 @@ def test_add_sdf_mesh(self): sdf = self.sim.add_rigid_object( cfg=RigidObjectCfg( uid="duck_sdf", - shape=MeshCfg(fpath=duck_path, sdf_resolution=128), + shape=MeshCfg( + fpath=duck_path, + collision=MeshCollisionCfg( + approximation="sdf", + sdf_resolution=128, + ), + ), body_type="dynamic", ) ) @@ -361,6 +449,8 @@ def test_body_data(self): """Test the body_data property for dynamic objects.""" # Dynamic object should have body_data assert self.duck.body_data is not None, "Dynamic duck should have body_data" + assert self.duck.body_data.mass.shape == (NUM_ARENAS,) + assert self.duck.body_data.inertia.shape == (NUM_ARENAS, 3) # Static object should return None with warning assert self.table.body_data is None, "Static table should not have body_data" @@ -368,6 +458,29 @@ def test_body_data(self): # Kinematic object should have body_data assert self.chair.body_data is not None, "Kinematic chair should have body_data" + def test_default_physical_properties_remain_at_initialized_values(self): + """Test runtime writes do not mutate the mass-property snapshots.""" + assert self.duck.body_data is not None + data = self.duck.body_data + initial_mass = self.duck.get_mass().clone() + initial_inertia = self.duck.get_inertia().clone() + initial_com_pose = data.com_pose.clone() + + assert torch.allclose(data.default_mass, initial_mass) + assert torch.allclose(data.default_inertia, initial_inertia) + assert torch.allclose(data.default_com_pose, initial_com_pose) + assert torch.allclose(self.duck.default_mass, data.default_mass) + + self.duck.set_mass(initial_mass + 0.5) + self.duck.set_inertia(initial_inertia + 0.1) + changed_com_pose = initial_com_pose.clone() + changed_com_pose[:, :3] += 0.05 + self.duck.set_com_pose(changed_com_pose) + + assert torch.allclose(data.default_mass, initial_mass) + assert torch.allclose(data.default_inertia, initial_inertia) + assert torch.allclose(data.default_com_pose, initial_com_pose) + def test_physical_attributes(self): """Test getting and setting physical attributes and body states.""" # 1. Body state @@ -403,30 +516,108 @@ def test_physical_attributes(self): # 2. is_non_dynamic assert not self.duck.is_non_dynamic, "Dynamic duck should not be is_non_dynamic" 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" + assert self.chair.is_non_dynamic == (self.chair.body_type == "kinematic") + + if self.physics == "newton": + expected_mass = torch.ones(NUM_ARENAS, device=self.sim.device) + 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 self.duck.get_friction().shape == (NUM_ARENAS,) + assert torch.isfinite(self.duck.get_friction()).all() + assert self.duck.get_damping().shape == (NUM_ARENAS, 2) + assert torch.isfinite(self.duck.get_damping()).all() + + self.duck.set_attrs( + RigidBodyPhysicsCfg.from_dict({"mass_props": {"mass": 2.5}}) + ) + assert torch.allclose( + self.duck.get_mass(), + torch.full((NUM_ARENAS,), 2.5, device=self.sim.device), + ) + + # Actor type is topology, not a runtime batch property. + with pytest.raises(NotImplementedError, match="descriptor mutation"): + 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) + actual_inertia = self.duck.get_inertia() + assert torch.allclose( + actual_inertia, + new_inertia, + atol=NEWTON_INERTIA_ROUND_TRIP_ATOL, + rtol=0.0, + ), ( + "Newton set_inertia round-trip failed: " + f"max_abs_error={(actual_inertia - new_inertia).abs().max().item()}" + ) + + # 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" + + # Static Spawn actors do not have dynamic body ids. Their getters + # remain readable from source/backend metadata. Empty grouped cfgs + # intentionally preserve those values rather than authoring defaults. + assert self.table.get_mass().shape == (NUM_ARENAS,) + assert torch.isfinite(self.table.get_mass()).all() + assert self.table.get_friction().shape == (NUM_ARENAS,) + assert torch.isfinite(self.table.get_friction()).all() + assert self.table.get_damping().shape == (NUM_ARENAS, 2) + assert torch.isfinite(self.table.get_damping()).all() + assert torch.equal( + self.table.get_inertia(), + torch.zeros((NUM_ARENAS, 3), device=self.sim.device), + ) + return # 3. body_type assert self.duck.body_type == "dynamic" - self.duck.set_body_type("kinematic") - assert self.duck.body_type == "kinematic" - self.duck.set_body_type("dynamic") + with pytest.raises(NotImplementedError, match="descriptor mutation"): + self.duck.set_body_type("kinematic") assert self.duck.body_type == "dynamic" - assert self.chair.body_type == "kinematic" - self.chair.set_body_type("dynamic") - assert self.chair.body_type == "dynamic" - self.chair.set_body_type("kinematic") - assert self.chair.body_type == "kinematic" + if self.chair.body_type == "kinematic": + with pytest.raises(NotImplementedError, match="descriptor mutation"): + self.chair.set_body_type("dynamic") + assert self.chair.body_type == "kinematic" # 4. attrs - new_attrs = RigidBodyAttributesCfg(mass=2.5, density=1000.0) + new_attrs = RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 2.5, "density": 1000.0}} + ) self.duck.set_attrs(new_attrs) masses = self.duck.get_mass() assert torch.allclose( masses, torch.tensor([2.5] * NUM_ARENAS, device=self.sim.device) ), f"Mass not set correctly: {masses.tolist()}" - partial_attrs = RigidBodyAttributesCfg(mass=3.0) + partial_attrs = RigidBodyPhysicsCfg.from_dict({"mass_props": {"mass": 3.0}}) self.duck.set_attrs(partial_attrs, env_ids=[0]) masses = self.duck.get_mass() assert torch.allclose( @@ -482,29 +673,62 @@ 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) + if self.chair.body_type == "kinematic": + assert torch.allclose( + self.chair.body_data.com_pose, chair_com_pose_before, atol=1e-5 + ), "Kinematic rigid object COM pose should not change" + else: + assert torch.allclose( + self.chair.body_data.com_pose, com_pose, atol=1e-5 + ), "Dynamic rigid object COM pose should 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.""" @@ -579,6 +803,291 @@ 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_reset_restores_default_physical_properties(self): + """Test full and partial reset restore mass, inertia, and COM defaults.""" + assert self.duck.body_data is not None + data = self.duck.body_data + default_mass = data.default_mass.clone() + default_inertia = data.default_inertia.clone() + default_com_pose = data.default_com_pose.clone() + + changed_mass = default_mass + 0.5 + changed_inertia = default_inertia + 0.1 + changed_com_pose = default_com_pose.clone() + changed_com_pose[:, :3] += 0.05 + self.duck.set_mass(changed_mass) + self.duck.set_inertia(changed_inertia) + self.duck.set_com_pose(changed_com_pose) + + self.duck.reset(env_ids=[0]) + + mass_after_partial = self.duck.get_mass() + inertia_after_partial = self.duck.get_inertia() + com_after_partial = data.com_pose + inertia_atol = ( + NEWTON_INERTIA_ROUND_TRIP_ATOL if self.physics == "newton" else 1e-5 + ) + assert torch.allclose(mass_after_partial[0], default_mass[0], atol=1e-5) + assert torch.allclose( + inertia_after_partial[0], + default_inertia[0], + atol=inertia_atol, + rtol=0.0, + ) + assert torch.allclose(com_after_partial[0], default_com_pose[0], atol=1e-5) + assert torch.allclose(mass_after_partial[1], changed_mass[1], atol=1e-5) + assert torch.allclose( + inertia_after_partial[1], + changed_inertia[1], + atol=inertia_atol, + rtol=0.0, + ) + assert torch.allclose(com_after_partial[1], changed_com_pose[1], atol=1e-5) + + self.duck.reset() + + assert torch.allclose(self.duck.get_mass(), default_mass, atol=1e-5) + assert torch.allclose( + self.duck.get_inertia(), + default_inertia, + atol=inertia_atol, + rtol=0.0, + ) + assert torch.allclose(data.com_pose, default_com_pose, atol=1e-5) + + 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 + expected_xyzw = torch.tensor( + [1.0, 2.0, 3.0, 4.0], device=self.sim.device + ) / torch.sqrt(torch.tensor(30.0, device=self.sim.device)) + pose_7[:3, :3] = matrix_from_quat(expected_xyzw.unsqueeze(0))[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}" + torch.testing.assert_close( + pose_vec[:, 3:7], + expected_xyzw.unsqueeze(0).expand(NUM_ARENAS, -1), + atol=1e-5, + rtol=1e-5, + ) + + # 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 test_multi_mesh_geometry_is_combined(self): """GLB render meshes are exported as one complete indexed geometry.""" render_body = self.chair._entities[0].get_render_body() @@ -620,6 +1129,155 @@ class TestRigidObjectCUDA(BaseRigidObjectTest): def setup_method(self): self.setup_simulation("cuda") + def test_kinematic_binding_supports_pose_updates(self): + obj = self.sim.add_rigid_object( + cfg=RigidObjectCfg( + uid="gpu_kinematic", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + body_type="kinematic", + ) + ) + assert obj.body_data is not None + + pose = torch.eye(4, device=self.sim.device).repeat(NUM_ARENAS, 1, 1) + pose[:, :3, 3] = torch.tensor([0.2, -0.1, 0.5], device=self.sim.device) + obj.set_local_pose(pose) + self.sim.update(0.01) + + assert torch.allclose(obj.get_local_pose(to_matrix=True), pose, atol=1e-5) + + +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): + """Typed Newton attributes register through the public Spawn result. + + Newton-native contact/shape parameters are consumed by the descriptor + adapter without an independently owned manager or legacy patch path. + """ + duck_path = get_data_path(DUCK_PATH) + cfg = RigidObjectCfg( + uid="duck_newton_native", + shape=MeshCfg(fpath=duck_path), + body_type="dynamic", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + collision_props=NewtonCollisionPropertiesCfg(margin=0.01), + material_props=NewtonRigidBodyMaterialCfg( + dynamic_friction=0.5, + restitution=0.1, + ke=1e3, + kd=50.0, + ), + ), + ) + obj: RigidObject = self.sim.add_rigid_object(cfg=cfg) + self.sim.prepare() + + assert obj.num_instances == NUM_ARENAS + assert obj.body_type == "dynamic" + result = self.sim.spawn_result + handles = [ + result.get_object(f"{arena_name}/{obj.uid}") + for arena_name in result.arenas.names[1:] + ] + assert len(result.create_rigid_body_batch(handles)) == NUM_ARENAS + assert all(handle.is_valid for handle in handles) + assert all( + handle.desc is not None and handle.desc.physics is not None + for handle in handles + ) + # 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() + + +@pytest.mark.gpu +class TestRigidObjectNewtonMujoco: + """Focused standalone-rigid state synchronization on MuJoCo-Warp.""" + + def setup_method(self): + physics_cfg = physics_cfg_for_backend("newton") + physics_cfg.gravity = (0.0, 0.0, 0.0) + physics_cfg.solver_cfg = {"solver_type": "mujoco_warp"} + self.sim = SimulationManager( + SimulationManagerCfg( + headless=True, + device="cuda", + num_envs=1, + physics_cfg=physics_cfg, + ) + ) + self.obj = self.sim.add_rigid_object( + RigidObjectCfg( + uid="free_body", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + ), + init_pos=(0.0, 0.0, Z_TRANSLATION), + ) + ) + self.sim.prepare() + + def teardown_method(self): + self.sim.destroy() + SimulationManager.flush_cleanup_queue() + _teardown_newton_physics() + + def test_clear_dynamics_persists_across_mujoco_step(self): + """Clearing body velocity also clears its reduced FREE-joint velocity.""" + linear_velocity = torch.tensor( + [[0.4, -0.2, 0.3]], + dtype=torch.float32, + device=self.sim.device, + ) + angular_velocity = torch.tensor( + [[0.1, 0.2, -0.3]], + dtype=torch.float32, + device=self.sim.device, + ) + self.obj.set_velocity( + lin_vel=linear_velocity, + ang_vel=angular_velocity, + ) + self.sim.update(step=1) + assert not torch.allclose( + self.obj.body_data.vel, + torch.zeros((1, 6), device=self.sim.device), + ) + + self.obj.clear_dynamics() + self.sim.update(step=1) + + torch.testing.assert_close( + self.obj.body_data.vel, + torch.zeros((1, 6), device=self.sim.device), + atol=1.0e-5, + rtol=0.0, + ) + if __name__ == "__main__": # pytest.main(["-s", __file__]) diff --git a/tests/sim/objects/test_rigid_object_group.py b/tests/sim/objects/test_rigid_object_group.py index 4fadc3869..be8c9da48 100644 --- a/tests/sim/objects/test_rigid_object_group.py +++ b/tests/sim/objects/test_rigid_object_group.py @@ -17,12 +17,18 @@ from __future__ import annotations import os +from unittest.mock import Mock + import torch import pytest from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import RigidBodyGroupData, RigidObjectGroup -from embodichain.lab.sim.cfg import RigidObjectGroupCfg, RigidObjectCfg +from embodichain.lab.sim.cfg import ( + RigidObjectGroupCfg, + RigidObjectCfg, + physics_cfg_for_backend, +) from embodichain.lab.sim.shapes import MeshCfg from embodichain.data import get_data_path from dexsim.types import ActorType @@ -31,39 +37,51 @@ TABLE_PATH = "ShopTableSimple/shop_table_simple.ply" NUM_ARENAS = 4 Z_TRANSLATION = 2.0 +# Newton converts principal-frame inertia diagonals through a float32 full +# tensor, so imported non-axis-aligned COM frames are not bit-exact on readback. +NEWTON_INERTIA_ROUND_TRIP_ATOL = 2e-4 -@pytest.mark.no_sim -def test_cpu_body_data_reads_angular_velocity_from_angular_api(): - """CPU rigid-object groups must not report linear velocity as angular.""" +def _teardown_newton_physics() -> None: + from dexsim.engine.newton_physics import teardown_newton_physics - class VelocityEntity: - def get_linear_velocity(self): - return [1.0, 2.0, 3.0] + teardown_newton_physics() - def get_angular_velocity(self): - return [4.0, 5.0, 6.0] - body_data = object.__new__(RigidBodyGroupData) - body_data.entities = [[VelocityEntity(), VelocityEntity()]] - body_data.device = torch.device("cpu") +@pytest.mark.no_sim +def test_cpu_body_data_reads_angular_velocity_from_angular_api(): + """CPU rigid-object groups must not report linear velocity as angular.""" + expected = torch.tensor([[[4.0, 5.0, 6.0], [4.0, 5.0, 6.0]]]) + body_view = Mock() + body_view.fetch_angular_velocity.side_effect = lambda out: out.copy_( + expected.reshape(-1, 3) + ) + body_data = RigidBodyGroupData( + body_view, + num_instances=1, + num_objects=2, + device=torch.device("cpu"), + ) angular_velocity = body_data.ang_vel - assert torch.equal( - angular_velocity, - torch.tensor([[[4.0, 5.0, 6.0], [4.0, 5.0, 6.0]]]), - ) + assert torch.equal(angular_velocity, expected) + body_view.fetch_angular_velocity.assert_called_once() + body_view.fetch_linear_velocity.assert_not_called() class BaseRigidObjectGroupTest: """Shared test logic for CPU and CUDA.""" - def setup_simulation(self, sim_device): + def setup_simulation(self, device: str, physics: str = "default") -> None: config = SimulationManagerCfg( - headless=True, sim_device=sim_device, num_envs=NUM_ARENAS + headless=True, + device=device, + num_envs=NUM_ARENAS, + physics_cfg=physics_cfg_for_backend(physics), ) self.sim = SimulationManager(config) + self.physics = physics duck_path = get_data_path(DUCK_PATH) assert os.path.isfile(duck_path) @@ -91,8 +109,7 @@ def setup_simulation(self, sim_device): cfg=RigidObjectGroupCfg.from_dict(cfg_dict) ) - if sim_device == "cuda" and self.sim.is_use_gpu_physics: - self.sim.init_gpu_physics() + self.sim.prepare() self.sim.enable_physics(True) @@ -117,6 +134,118 @@ def test_local_pose_behavior(self): atol=1e-5, ), "FAIL: Local poses do not match after setting." + distinct_xyzw = torch.tensor( + [1.0, 2.0, 3.0, 4.0], device=self.sim.device + ) / torch.sqrt(torch.tensor(30.0, device=self.sim.device)) + vector_pose = torch.zeros( + NUM_ARENAS, + self.obj_group.num_objects, + 7, + device=self.sim.device, + ) + vector_pose[..., :3] = combined_pose[..., :3, 3] + vector_pose[..., 3:7] = distinct_xyzw + + self.obj_group.set_local_pose(vector_pose) + + torch.testing.assert_close( + self.obj_group.get_local_pose(), + vector_pose, + atol=1e-5, + rtol=1e-5, + ) + + def test_body_data_exposes_mass_properties(self): + """Current and initialization-time properties use [env, object] layout.""" + data = self.obj_group.body_data + expected_prefix = (NUM_ARENAS, self.obj_group.num_objects) + + assert data.mass.shape == expected_prefix + assert data.inertia.shape == (*expected_prefix, 3) + assert data.com_pose.shape == (*expected_prefix, 7) + assert data.default_mass.shape == data.mass.shape + assert data.default_inertia.shape == data.inertia.shape + assert data.default_com_pose.shape == data.com_pose.shape + + def test_reset_restores_default_mass_properties(self): + """Partial reset restores Group mass properties only in selected envs.""" + data = self.obj_group.body_data + env_ids = [0, 1] + obj_ids = [0] + default_mass = data.default_mass[env_ids, :1].clone() + default_inertia = data.default_inertia[env_ids, :1].clone() + default_com_pose = data.default_com_pose[env_ids, :1].clone() + changed_mass = default_mass + 0.5 + changed_inertia = default_inertia * 1.25 + changed_com_pose = default_com_pose.clone() + changed_com_pose[..., 0] += 0.02 + changed_com_pose[..., 3:7] = torch.tensor( + [1.0, 2.0, 3.0, 4.0], device=self.sim.device + ) / torch.sqrt(torch.tensor(30.0, device=self.sim.device)) + + self.obj_group.set_mass(changed_mass, env_ids=env_ids, obj_ids=obj_ids) + self.obj_group.set_inertia( + changed_inertia, + env_ids=env_ids, + obj_ids=obj_ids, + ) + self.obj_group.set_com_pose( + changed_com_pose, + env_ids=env_ids, + obj_ids=obj_ids, + ) + + assert torch.allclose(data.default_mass[env_ids, :1], default_mass) + assert torch.allclose(data.default_inertia[env_ids, :1], default_inertia) + assert torch.allclose(data.default_com_pose[env_ids, :1], default_com_pose) + + self.obj_group.reset(env_ids=[env_ids[0]]) + mass_after_partial = self.obj_group.get_mass(env_ids=env_ids, obj_ids=obj_ids) + inertia_after_partial = self.obj_group.get_inertia( + env_ids=env_ids, obj_ids=obj_ids + ) + com_after_partial = self.obj_group.get_com_pose( + env_ids=env_ids, obj_ids=obj_ids + ) + inertia_atol = ( + NEWTON_INERTIA_ROUND_TRIP_ATOL if self.physics == "newton" else 1e-5 + ) + + assert torch.allclose(mass_after_partial[0], default_mass[0], atol=1e-5) + assert torch.allclose(mass_after_partial[1], changed_mass[1], atol=1e-5) + assert torch.allclose( + inertia_after_partial[0], + default_inertia[0], + atol=inertia_atol, + rtol=0.0, + ) + assert torch.allclose( + inertia_after_partial[1], + changed_inertia[1], + atol=inertia_atol, + rtol=0.0, + ) + assert torch.allclose(com_after_partial[0], default_com_pose[0], atol=1e-5) + assert torch.allclose(com_after_partial[1], changed_com_pose[1], atol=1e-5) + + self.obj_group.reset(env_ids=[env_ids[1]]) + assert torch.allclose( + self.obj_group.get_mass(env_ids=env_ids, obj_ids=obj_ids), + default_mass, + atol=1e-5, + ) + assert torch.allclose( + self.obj_group.get_inertia(env_ids=env_ids, obj_ids=obj_ids), + default_inertia, + atol=inertia_atol, + rtol=0.0, + ) + assert torch.allclose( + self.obj_group.get_com_pose(env_ids=env_ids, obj_ids=obj_ids), + default_com_pose, + atol=1e-5, + ) + def test_get_user_ids(self): """Test get_user_ids method.""" user_ids = self.obj_group.get_user_ids() @@ -171,12 +300,20 @@ def setup_method(self): self.setup_simulation("cpu") -@pytest.mark.skip(reason="Skipping CUDA tests temporarily") class TestRigidObjectGroupCUDA(BaseRigidObjectGroupTest): def setup_method(self): self.setup_simulation("cuda") +class TestRigidObjectGroupNewton(BaseRigidObjectGroupTest): + def setup_method(self): + self.setup_simulation("cuda", physics="newton") + + def teardown_method(self): + super().teardown_method() + _teardown_newton_physics() + + if __name__ == "__main__": # pytest.main(["-s", __file__]) test = TestRigidObjectGroupCPU() diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index e6e050f91..2d0cdbbe4 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -25,7 +25,9 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.objects import Robot +from embodichain.lab.sim.objects.backends.newton import _default_mujoco_mimic_solref 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,6 +52,20 @@ ], } +W1_ACTIVE_DOF = 40 # Dexforce W1 v021 scalar active-DOF count. + + +@pytest.mark.no_sim +def test_default_mujoco_mimic_solref_preserves_damping_and_timestep_floor(): + np.testing.assert_allclose( + _default_mujoco_mimic_solref(physics_dt=0.01, num_substeps=10), + [2.0e-3, 1.0e1], + ) + np.testing.assert_allclose( + _default_mujoco_mimic_solref(physics_dt=1.0e-4, num_substeps=10), + [1.0e-4, 1.0e1], + ) + def test_get_qf_selects_control_part_joint_efforts(): full_qf = torch.tensor( @@ -68,11 +84,11 @@ def test_get_qf_selects_control_part_joint_efforts(): # 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( @@ -83,10 +99,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): - cls.sim.init_gpu_physics() + cls.sim.prepare() def test_get_joint_ids(self): left_joint_ids = self.robot.get_joint_ids("left_arm") @@ -260,6 +273,52 @@ def test_mimic(self): len(right_eef_ids_without_mimic) == 6 ), f"Expected 6 right eef joint IDs without mimic, got {len(right_eef_ids_without_mimic)}" + def test_default_mimic_tracks_closed_hand_target(self): + """Keep W1 hand mimic constraints equally stiff on CPU and CUDA.""" + self.robot.reset() + open_target = torch.tensor( + [[0.0, 1.5, 0.0, 0.0, 0.0, 0.0]], + dtype=torch.float32, + device=self.sim.device, + ) + close_target = torch.tensor( + [[0.1, 1.5, 0.3, 0.2, 0.3, 0.3]], + dtype=torch.float32, + device=self.sim.device, + ) + for target in (open_target, close_target): + self.robot.set_qpos( + target.repeat(self.robot.num_instances, 1), name="right_eef" + ) + self.sim.update(step=100) + + qpos = self.robot.body_data.qpos + target_qpos = self.robot.body_data.target_qpos + right_eef_ids = self.robot.get_joint_ids("right_eef") + right_mimic_errors = [] + for mimic_id, parent_id, multiplier, offset in zip( + self.robot.mimic_ids, + self.robot.mimic_parents, + self.robot.mimic_multipliers, + self.robot.mimic_offsets, + strict=True, + ): + if not self.robot.joint_names[parent_id].startswith("RIGHT_HAND"): + continue + right_mimic_errors.append( + torch.abs( + qpos[:, mimic_id] - (qpos[:, parent_id] * multiplier + offset) + ) + ) + + assert torch.max(torch.stack(right_mimic_errors)).item() < 0.02 + assert ( + torch.max( + torch.abs(qpos[:, right_eef_ids] - target_qpos[:, right_eef_ids]) + ).item() + < 0.01 + ) + def test_setter_and_getter_with_control_part(self): left_arm_qpos = self.robot.get_qpos(name="left_arm") assert left_arm_qpos.shape == (10, 7) @@ -452,7 +511,7 @@ def test_robot_cfg_merge(self): cfg = deepcopy(self.robot.cfg) cfg_dict = { - "drive_pros": { + "joint_drive_props": { "max_effort": { "(LEFT|RIGHT)_HAND_(THUMB[12]|INDEX|MIDDLE|RING|PINKY)": 1.0, }, @@ -467,7 +526,7 @@ def test_robot_cfg_merge(self): cfg = merge_robot_cfg(cfg, cfg_dict) assert ( - cfg.drive_pros.max_effort[ + cfg.joint_drive_props.max_effort[ "(LEFT|RIGHT)_HAND_(THUMB[12]|INDEX|MIDDLE|RING|PINKY)" ] == 1.0 @@ -519,6 +578,201 @@ 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, prepare, control surface). + + A robot is a URDF articulation; the Newton ``load_urdf`` patch builds a + NewtonArticulation. This exercises the add_robot -> prepare -> 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"}) + cfg.init_qpos = [0.0001 * (index + 1) for index in range(W1_ACTIVE_DOF)] + self.robot: Robot = self.sim.add_robot(cfg=cfg) + self.sim.prepare() + + 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, prepares, and exposes a working control surface.""" + assert self.sim.is_newton_backend + assert self.robot.body_data.is_ready + assert self.robot.dof > 0 + + state_joint_names = self.robot.body_data.articulation_view.joint_names + assert self.robot.joint_names == state_joint_names + source_joint_names = self.robot._entities[0].get_actived_joint_names() + initial_qpos_by_name = dict( + zip(source_joint_names, self.robot.cfg.init_qpos, strict=True) + ) + mimic_relations = list( + zip( + self.robot.mimic_ids, + self.robot.mimic_parents, + self.robot.mimic_multipliers, + self.robot.mimic_offsets, + strict=True, + ) + ) + assert all( + state_joint_names[mimic_id].endswith("_PIP") + and "_HAND_" in state_joint_names[parent_id] + for mimic_id, parent_id, _, _ in mimic_relations + ) + initial_qpos = self.robot.body_data.qpos[0].detach().cpu().tolist() + assert dict(zip(state_joint_names, initial_qpos, strict=True)) == pytest.approx( + initial_qpos_by_name + ) + + binding = self.robot._entities[0]._physics_binding + model = binding._runtime.model + runtime_joints = {joint.name: joint for joint in binding.joints} + mimic_joint0 = np.asarray(model.constraint_mimic_joint0.numpy()).reshape(-1) + mimic_joint1 = np.asarray(model.constraint_mimic_joint1.numpy()).reshape(-1) + row_by_pair = { + (int(child), int(parent)): row + for row, (child, parent) in enumerate( + zip(mimic_joint0, mimic_joint1, strict=True) + ) + } + constraint_rows = [] + for mimic_id, parent_id, _, _ in mimic_relations: + child = runtime_joints[state_joint_names[mimic_id]] + parent = runtime_joints[state_joint_names[parent_id]] + constraint_rows.append( + row_by_pair[(int(child.joint_id), int(parent.joint_id))] + ) + + solver = binding._runtime.solver + mapping = np.asarray(solver.mjc_eq_to_newton_mimic.numpy()) + selected_eq = np.isin(mapping, np.asarray(constraint_rows, dtype=np.int32)) + assert int(selected_eq.sum()) == len(mimic_relations) + eq_solref = np.asarray(solver.mjw_model.eq_solref.numpy()) + np.testing.assert_allclose( + eq_solref[selected_eq], + np.broadcast_to([2.0e-3, 1.0e1], (len(mimic_relations), 2)), + ) + target_ke = np.asarray(model.joint_target_ke.numpy()) + target_kd = np.asarray(model.joint_target_kd.numpy()) + target_mode = np.asarray(model.joint_target_mode.numpy()) + eq_active = np.asarray(solver.mjw_data.eq_active.numpy()) + assert np.all(eq_active[selected_eq]) + for mimic_id, parent_id, _, _ in mimic_relations: + child = runtime_joints[state_joint_names[mimic_id]] + parent = runtime_joints[state_joint_names[parent_id]] + assert target_ke[child.qd_start] == pytest.approx( + target_ke[parent.qd_start] * 1.0e-2 + ) + assert target_kd[child.qd_start] == pytest.approx( + target_kd[parent.qd_start] * 1.0e-2 + ) + assert target_mode[child.qd_start] == target_mode[parent.qd_start] + + # This physical check covers both hands and keeps the native coupled + # constraints bounded under the W1's self-contacts. Default can also + # deflect these compliant joints by several tenths of a radian. + self.sim.update(step=100) + settled_qpos = self.robot.body_data.qpos + settled_errors = [] + for mimic_id, parent_id, multiplier, offset in mimic_relations: + settled_errors.append( + torch.abs( + settled_qpos[:, mimic_id] + - (settled_qpos[:, parent_id] * multiplier + offset) + ) + ) + assert torch.max(torch.stack(settled_errors)).item() < 0.5 + + 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 + assert [ + state_joint_names[index] for index in left_ids + ] == self.robot.control_parts["left_arm"] + assert [ + state_joint_names[index] for index in right_ids + ] == self.robot.control_parts["right_arm"] + right_eef_ids = self.robot.get_joint_ids("right_eef") + assert [ + state_joint_names[index] for index in right_eef_ids + ] == self.robot.control_parts["right_eef"] + + right_qpos_limits = self.robot.get_qpos_limits(name="right_arm") + requested_target = torch.full( + (1, len(right_ids)), 0.1, dtype=torch.float32, device=self.sim.device + ) + expected_target = requested_target.clamp( + right_qpos_limits[..., 0], right_qpos_limits[..., 1] + ) + self.robot.set_qpos(requested_target, name="right_arm") + torch.testing.assert_close( + self.robot.body_data.target_qpos[:, right_ids], expected_target + ) + + hand_target = torch.tensor( + [[0.1, 1.0, 0.2, 0.3, 0.4, 0.5]], + dtype=torch.float32, + device=self.sim.device, + ) + self.robot.set_qpos(hand_target, name="right_eef") + target_qpos = self.robot.body_data.target_qpos + torch.testing.assert_close(target_qpos[:, right_eef_ids], hand_target) + for mimic_id, parent_id, multiplier, offset in mimic_relations: + torch.testing.assert_close( + target_qpos[:, mimic_id], + target_qpos[:, parent_id] * multiplier + offset, + ) + hand_velocity_target = torch.tensor( + [[0.05, 0.1, 0.15, 0.2, 0.25, 0.3]], + dtype=torch.float32, + device=self.sim.device, + ) + self.robot.set_qvel(hand_velocity_target, name="right_eef") + target_qvel = self.robot.body_data.target_qvel + torch.testing.assert_close(target_qvel[:, right_eef_ids], hand_velocity_target) + for mimic_id, parent_id, multiplier, _ in mimic_relations: + torch.testing.assert_close( + target_qvel[:, mimic_id], + target_qvel[:, parent_id] * multiplier, + ) + self.robot.set_qvel(torch.zeros_like(hand_velocity_target), name="right_eef") + + # 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_robot_cfg.py b/tests/sim/objects/test_robot_cfg.py index ca025a119..16020b0f0 100644 --- a/tests/sim/objects/test_robot_cfg.py +++ b/tests/sim/objects/test_robot_cfg.py @@ -21,7 +21,10 @@ import pytest from embodichain.lab.sim.cfg import ( + ArticulationRootPropertiesCfg, + CollisionPropertiesCfg, JointDrivePropertiesCfg, + RigidBodyPhysicsCfg, RobotCfg, ) from embodichain.lab.sim.workspace import RobotWorkspaceCfg @@ -64,11 +67,15 @@ def resolve(path): def test_dexforce_w1_roundtrip(): cfg = DexforceW1Cfg.from_dict({"uid": "dexforce_w1", "version": "v021"}) + assert type(cfg.root_props) is ArticulationRootPropertiesCfg + assert cfg.root_props.min_position_iters == 32 + assert cfg.root_props.min_velocity_iters == 8 d = cfg.to_dict() assert d["uid"] == "dexforce_w1" cfg2 = DexforceW1Cfg.from_dict(d) assert cfg2.uid == "dexforce_w1" assert cfg2.version == DexforceW1Version.V021 + assert type(cfg2.root_props) is ArticulationRootPropertiesCfg def test_dexforce_w1_solver_cfg_is_srs_and_set_once(): @@ -409,7 +416,7 @@ def _build_defaults(self, init_dict=None): self.uid = "roundtrip" self.variant = _RoundTripVariant(init_dict.get("variant", "a")) self.control_parts = {"arm": ["J1", "J2"]} - self.drive_pros = JointDrivePropertiesCfg( + self.joint_drive_props = JointDrivePropertiesCfg( stiffness={"J[1-2]": 1e4}, damping={"J[1-2]": 1e3} ) @@ -426,10 +433,12 @@ def test_robotcfg_to_dict_roundtrip(): assert cfg2.uid == "roundtrip" assert cfg2.variant == _RoundTripVariant.B assert cfg2.control_parts == {"arm": ["J1", "J2"]} - assert cfg2.drive_pros.stiffness == {"J[1-2]": 1e4} + assert cfg2.joint_drive_props.stiffness == {"J[1-2]": 1e4} from embodichain.lab.sim.robots.cobotmagic import CobotMagicCfg +from embodichain.lab.sim.robots.franka_panda import FrankaPandaCfg +from embodichain.lab.sim.robots.ur_robot import URRobotCfg from embodichain.lab.sim.solvers import OPWSolverCfg @@ -444,6 +453,13 @@ def test_cobotmagic_from_dict_and_roundtrip(): } assert isinstance(cfg.solver_cfg["left_arm"], OPWSolverCfg) assert isinstance(cfg.solver_cfg["right_arm"], OPWSolverCfg) + assert isinstance(cfg.attrs, RigidBodyPhysicsCfg) + assert type(cfg.attrs.collision_props) is CollisionPropertiesCfg + assert cfg.attrs.collision_props.contact_offset == pytest.approx(0.001) + assert cfg.attrs.collision_props.rest_offset == pytest.approx(0.0) + assert type(cfg.root_props) is ArticulationRootPropertiesCfg + assert cfg.root_props.min_position_iters == 8 + assert cfg.root_props.min_velocity_iters == 2 d = cfg.to_dict() assert d["uid"] == "CobotMagic" @@ -453,6 +469,27 @@ def test_cobotmagic_from_dict_and_roundtrip(): assert isinstance(cfg2.solver_cfg["left_arm"], OPWSolverCfg) +@pytest.mark.parametrize( + ("cfg_type", "init_dict"), + [ + (CobotMagicCfg, {}), + (FrankaPandaCfg, {}), + (URRobotCfg, {}), + (DexforceW1Cfg, {}), + ], +) +def test_specified_robots_use_portable_joint_drive_semantics( + cfg_type: type[RobotCfg], + init_dict: dict, +) -> None: + cfg = cfg_type.from_dict(init_dict) + + assert type(cfg.joint_drive_props) is JointDrivePropertiesCfg + assert cfg.joint_drive_props.drive_type == "force" + assert cfg.joint_drive_props.target_mode is None + assert cfg.joint_drive_props._resolve_modes() == ("position_velocity", "force") + + def test_robotcfg_save_to_file(tmp_path): cfg = _RoundTripCfg.from_dict({"variant": "b"}) fp = tmp_path / "cfg.json" @@ -523,7 +560,6 @@ def test_cobotmagic_pk_dof_matches_control_parts(): # URRobotCfg -- UR family (ur3 / ur3e / ur5 / ur5e / ur10 / ur10e) # --------------------------------------------------------------------------- # -from embodichain.lab.sim.robots.ur_robot import URRobotCfg from embodichain.lab.sim.solvers import URSolverCfg UR_TYPES = ["ur3", "ur3e", "ur5", "ur5e", "ur10", "ur10e"] @@ -562,7 +598,7 @@ def test_ur_robot_max_effort_scales_with_size(): ur3 = URRobotCfg.from_dict({"robot_type": "ur3"}) ur5 = URRobotCfg.from_dict({"robot_type": "ur5"}) ur10 = URRobotCfg.from_dict({"robot_type": "ur10"}) - eff = lambda c: c.drive_pros.max_effort["arm"] # noqa: E731 + eff = lambda c: c.joint_drive_props.max_effort["arm"] # noqa: E731 assert eff(ur3) < eff(ur5) < eff(ur10) diff --git a/tests/sim/objects/test_soft_object.py b/tests/sim/objects/test_soft_object.py index d7334bb9e..7fe8352a9 100644 --- a/tests/sim/objects/test_soft_object.py +++ b/tests/sim/objects/test_soft_object.py @@ -26,9 +26,11 @@ ) from embodichain.lab.sim.shapes import MeshCfg from embodichain.lab.sim.objects import ( + DeformableObject, SoftBodyData, SoftObject, SoftObjectCfg, + VolumeDeformableObject, ) import pytest import torch @@ -54,7 +56,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, ) @@ -88,9 +90,9 @@ def setup_simulation(self): ), ), ) + self.sim.prepare() def test_run_simulation(self): - self.sim.init_gpu_physics() for _ in range(100): self.sim.update(step=1) self.cow.reset() @@ -99,7 +101,6 @@ def test_run_simulation(self): def test_get_deformable_mesh_geometry(self): """Test current collision vertices and matching surface triangles.""" - self.sim.init_gpu_physics() vertices = self.cow.get_current_collision_vertices() triangles = self.cow.get_collision_surface_triangles(env_ids=[0]) @@ -107,11 +108,35 @@ def test_get_deformable_mesh_geometry(self): assert triangles.ndim == 3 and triangles.shape[0] == 1 assert int(triangles.max()) < vertices.shape[1] + def test_unified_deformable_contract(self): + assert isinstance(self.cow, DeformableObject) + assert isinstance(self.cow, VolumeDeformableObject) + assert self.cow.deformable_type == "volume" + assert self.sim.get_deformable_object("cow") is self.cow + assert self.sim.get_soft_object("cow") is self.cow + assert self.sim.get_deformable_object_uid_list() == ["cow"] + + positions = self.cow.get_current_nodal_position() + velocities = self.cow.get_current_nodal_velocity() + state = self.cow.get_current_nodal_state() + default_state = self.cow.get_default_nodal_state() + assert positions.shape[-1] == 3 + assert velocities.shape == positions.shape + assert state.shape == (*positions.shape[:-1], 6) + assert default_state.shape == state.shape + torch.testing.assert_close( + self.cow.get_surface_vertices(), + self.cow.get_current_collision_vertices(), + ) + torch.testing.assert_close( + self.cow.get_surface_triangles(env_ids=[0]), + self.cow.get_collision_surface_triangles(env_ids=[0]), + ) + def test_remove(self): - self.sim.remove_asset(self.cow.uid) - assert ( - self.cow.uid not in self.sim._soft_objects - ), "Cow UID still present after removal" + with pytest.raises(NotImplementedError, match="pending removal"): + self.sim.remove_asset(self.cow.uid) + assert self.sim.get_deformable_object(self.cow.uid) is self.cow def teardown_method(self): """Clean up resources after each test method.""" diff --git a/tests/sim/objects/test_spawn_backend.py b/tests/sim/objects/test_spawn_backend.py new file mode 100644 index 000000000..f06692320 --- /dev/null +++ b/tests/sim/objects/test_spawn_backend.py @@ -0,0 +1,395 @@ +# ---------------------------------------------------------------------------- +# 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 pytest +import torch + +import embodichain.lab.sim.objects.backends.spawn as spawn_backend +from embodichain.lab.sim.objects.backends.spawn import ( + SpawnArticulationView, + SpawnRigidBodyView, + _embodichain_articulation_pose, + _embodichain_pose, + _spawn_articulation_pose, + _spawn_pose, +) + +pytestmark = pytest.mark.no_sim + + +def test_spawn_pose_adapters_preserve_embodichain_xyzw_order() -> None: + pose = torch.tensor([[1.0, 2.0, 3.0, 0.1, 0.2, 0.3, 0.9]]) + expected_spawn = torch.tensor([[0.1, 0.2, 0.3, 0.9, 1.0, 2.0, 3.0]]) + + torch.testing.assert_close(_spawn_pose(pose), expected_spawn) + torch.testing.assert_close(_spawn_articulation_pose(pose), expected_spawn) + torch.testing.assert_close(_embodichain_pose(expected_spawn), pose) + torch.testing.assert_close(_embodichain_articulation_pose(expected_spawn), pose) + + +class _SelectedRigidBatch: + def __init__(self, owner: _RigidBatch, rows: torch.Tensor) -> None: + self.owner = owner + self.rows = rows + + def apply_force(self, values: torch.Tensor) -> int: + self.owner.force[self.rows] = values + return len(self.rows) + + def apply_pose(self, values: torch.Tensor) -> int: + self.owner.pose[self.rows] = values + return len(self.rows) + + def apply_linear_velocity(self, values: torch.Tensor) -> int: + self.owner.linear_velocity[self.rows] = values + return len(self.rows) + + def apply_angular_velocity(self, values: torch.Tensor) -> int: + self.owner.angular_velocity[self.rows] = values + return len(self.rows) + + def apply_friction(self, values: torch.Tensor) -> int: + self.owner.friction[self.rows] = values + return len(self.rows) + + def fetch_friction(self, out: torch.Tensor) -> int: + out.copy_(self.owner.friction[self.rows]) + return len(self.rows) + + +class _RigidBatch: + def __init__(self) -> None: + self.force = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]) + self.pose = torch.tensor( + [ + [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0], + ] + ) + self.friction = torch.tensor([[0.1], [0.2], [0.3]]) + self.linear_velocity = torch.zeros((3, 3)) + self.angular_velocity = torch.zeros((3, 3)) + self.selections: list[tuple[int, ...]] = [] + + def __len__(self) -> int: + return len(self.force) + + def select(self, rows: torch.Tensor) -> _SelectedRigidBatch: + selected = rows.detach().cpu().to(dtype=torch.long) + self.selections.append(tuple(selected.tolist())) + return _SelectedRigidBatch(self, selected) + + +class _SelectedArticulationBatch: + def __init__(self, owner: _ArticulationBatch, rows: torch.Tensor) -> None: + self.owner = owner + self.rows = rows + + def apply_joint_force( + self, + values: torch.Tensor, + *, + dof_ids: torch.Tensor, + ) -> int: + columns = dof_ids.detach().cpu().to(dtype=torch.long) + self.owner.force[self.rows[:, None], columns] = values + self.owner.last_dof_ids = tuple(columns.tolist()) + return len(self.rows) + + def fetch_root_pose(self, out: torch.Tensor) -> int: + self.owner.root_pose_fetch_rows.append(tuple(self.rows.tolist())) + out.copy_(self.owner.root_pose[self.rows]) + return len(self.rows) + + def apply_root_pose(self, values: torch.Tensor) -> int: + self.owner.root_pose_apply_rows.append(tuple(self.rows.tolist())) + self.owner.root_pose[self.rows] = values + return len(self.rows) + + +class _ArticulationBatch: + def __init__(self) -> None: + layouts = tuple( + SimpleNamespace(name=f"joint_{index}", dof_start=index, dof_count=1) + for index in range(3) + ) + self.dof_counts = (3, 3) + self.link_counts = (1, 1) + self.joint_names_per_articulation = (("joint_0", "joint_1", "joint_2"),) * 2 + self.link_names_per_articulation = (("root",),) * 2 + self.joint_layouts_per_articulation = (layouts,) * 2 + self.dof_width = 3 + self.link_width = 1 + self.force = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + # Spawn articulation poses use xyzw + xyz layout. + self.root_pose = torch.tensor( + [ + [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 0.0, 1.0, 2.0, 0.0, 1.0], + ] + ) + self.last_dof_ids: tuple[int, ...] | None = None + self.root_pose_fetch_rows: list[tuple[int, ...]] = [] + self.root_pose_apply_rows: list[tuple[int, ...]] = [] + self.selections: list[tuple[int, ...]] = [] + + def __len__(self) -> int: + return len(self.force) + + def select(self, rows: torch.Tensor) -> _SelectedArticulationBatch: + selected = rows.detach().cpu().to(dtype=torch.long) + self.selections.append(tuple(selected.tolist())) + return _SelectedArticulationBatch(self, selected) + + +def test_rigid_partial_writes_delegate_to_selected_batch() -> None: + batch = _RigidBatch() + view = SpawnRigidBodyView( + SimpleNamespace(backend="newton"), + batch, + torch.device("cpu"), + ) + + view.apply_force(torch.tensor([[10.0, 20.0, 30.0]]), torch.tensor([1])) + view.apply_friction(torch.tensor([[0.9]]), torch.tensor([2])) + + assert torch.equal( + batch.force, + torch.tensor([[1.0, 2.0, 3.0], [10.0, 20.0, 30.0], [7.0, 8.0, 9.0]]), + ) + assert torch.equal(batch.friction, torch.tensor([[0.1], [0.2], [0.9]])) + assert batch.selections == [(1,), (2,)] + + +def test_rigid_partial_fetch_reads_only_selected_batch() -> None: + batch = _RigidBatch() + view = SpawnRigidBodyView( + SimpleNamespace(backend="dexsim"), + batch, + torch.device("cpu"), + ) + out = torch.empty((2, 1)) + + view.fetch_friction(out, torch.tensor([2, 0])) + + assert torch.equal(out, torch.tensor([[0.3], [0.1]])) + assert batch.selections == [(2, 0)] + + +def test_rigid_batch_failure_status_is_not_silently_ignored() -> None: + batch = _RigidBatch() + view = SpawnRigidBodyView( + SimpleNamespace(backend="dexsim"), + batch, + torch.device("cpu"), + ) + selected = batch.select(torch.tensor([0])) + selected.fetch_friction = lambda _out: -2 + batch.select = lambda _rows: selected + + with pytest.raises(RuntimeError, match="fetch_friction.*status -2"): + view.fetch_friction(torch.empty((1, 1)), torch.tensor([0])) + + +def test_newton_rigid_pose_write_synchronizes_free_joint_state(monkeypatch) -> None: + batch = _RigidBatch() + current_state = object() + other_state = object() + runtime = SimpleNamespace( + model=object(), + current_state=current_state, + other_state=other_state, + ) + batch._binding = SimpleNamespace( + _runtime=runtime, + _indices=torch.tensor([10, 11, 12]), + ) + synchronized_states: list[tuple[object, object]] = [] + created_body_ids: list[tuple[int, ...]] = [] + + class _StateSync: + def synchronize(self, states: tuple[object, object]) -> None: + synchronized_states.append(states) + + def _create_state_sync(_model: object, body_ids: list[int]) -> _StateSync: + created_body_ids.append(tuple(body_ids)) + return _StateSync() + + monkeypatch.setattr( + spawn_backend, + "_create_newton_standalone_state_sync", + _create_state_sync, + ) + view = SpawnRigidBodyView( + SimpleNamespace(backend="newton", topology_revision=3), + batch, + torch.device("cpu"), + ) + + view.apply_pose( + torch.tensor([[4.0, 5.0, 6.0, 0.0, 0.0, 0.0, 1.0]]), + torch.tensor([1]), + ) + view.apply_pose( + torch.tensor([[7.0, 8.0, 9.0, 0.0, 0.0, 0.0, 1.0]]), + torch.tensor([2]), + ) + + assert created_body_ids == [(10, 11, 12)] + assert synchronized_states == [ + (current_state, other_state), + (current_state, other_state), + ] + assert torch.equal( + batch.pose[1:], + torch.tensor( + [ + [0.0, 0.0, 0.0, 1.0, 4.0, 5.0, 6.0], + [0.0, 0.0, 0.0, 1.0, 7.0, 8.0, 9.0], + ] + ), + ) + + +def test_newton_rigid_velocity_writes_synchronize_free_joint_state( + monkeypatch, +) -> None: + batch = _RigidBatch() + current_state = object() + other_state = object() + runtime = SimpleNamespace( + model=object(), + current_state=current_state, + other_state=other_state, + ) + batch._binding = SimpleNamespace( + _runtime=runtime, + _indices=torch.tensor([10, 11, 12]), + ) + synchronized_states: list[tuple[object, object]] = [] + created_body_ids: list[tuple[int, ...]] = [] + + class _StateSync: + def synchronize(self, states: tuple[object, object]) -> None: + synchronized_states.append(states) + + def _create_state_sync(_model: object, body_ids: list[int]) -> _StateSync: + created_body_ids.append(tuple(body_ids)) + return _StateSync() + + monkeypatch.setattr( + spawn_backend, + "_create_newton_standalone_state_sync", + _create_state_sync, + ) + view = SpawnRigidBodyView( + SimpleNamespace(backend="newton", topology_revision=3), + batch, + torch.device("cpu"), + ) + + view.apply_linear_velocity( + torch.tensor([[1.0, 2.0, 3.0]]), + torch.tensor([1]), + ) + view.apply_angular_velocity( + torch.tensor([[4.0, 5.0, 6.0]]), + torch.tensor([2]), + ) + + assert created_body_ids == [(10, 11, 12)] + assert synchronized_states == [ + (current_state, other_state), + (current_state, other_state), + ] + assert torch.equal(batch.linear_velocity[1], torch.tensor([1.0, 2.0, 3.0])) + assert torch.equal(batch.angular_velocity[2], torch.tensor([4.0, 5.0, 6.0])) + + +def test_articulation_partial_force_preserves_other_rows_and_dofs() -> None: + batch = _ArticulationBatch() + view = SpawnArticulationView( + SimpleNamespace(backend="newton"), + batch, + torch.device("cpu"), + ) + + view.apply_qf( + torch.tensor([[50.0]]), + env_ids=torch.tensor([1]), + joint_ids=torch.tensor([1]), + ) + + assert torch.equal( + batch.force, + torch.tensor([[1.0, 2.0, 3.0], [4.0, 50.0, 6.0]]), + ) + assert batch.selections == [(1,)] + assert batch.last_dof_ids == (1,) + + +def test_newton_idempotent_root_pose_write_is_skipped() -> None: + batch = _ArticulationBatch() + view = SpawnArticulationView( + SimpleNamespace(backend="newton"), + batch, + torch.device("cpu"), + ) + current_pose = torch.tensor( + [ + [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + [2.0, 0.0, 1.0, 0.0, 0.0, 0.0, -1.0], + ] + ) + + view.apply_root_pose(current_pose, env_ids=torch.tensor([0, 1])) + + assert batch.root_pose_fetch_rows == [(0, 1)] + assert batch.root_pose_apply_rows == [] + + +def test_newton_root_pose_write_keeps_only_changed_rows() -> None: + batch = _ArticulationBatch() + view = SpawnArticulationView( + SimpleNamespace(backend="newton"), + batch, + torch.device("cpu"), + ) + target_pose = torch.tensor( + [ + [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + [3.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0], + ] + ) + + view.apply_root_pose(target_pose, env_ids=torch.tensor([0, 1])) + + assert batch.root_pose_fetch_rows == [(0, 1)] + assert batch.root_pose_apply_rows == [(1,)] + assert torch.equal( + batch.root_pose, + torch.tensor( + [ + [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0], + [0.0, 0.0, 0.0, 1.0, 3.0, 0.0, 1.0], + ] + ), + ) diff --git a/tests/sim/objects/test_usd.py b/tests/sim/objects/test_usd.py index 46e9f38ce..dd8dc0b10 100644 --- a/tests/sim/objects/test_usd.py +++ b/tests/sim/objects/test_usd.py @@ -29,41 +29,39 @@ ArticulationCfg, RigidObjectCfg, JointDrivePropertiesCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) 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): - self.sim.init_gpu_physics() - def test_import_rigid(self): - default_attr = RigidBodyAttributesCfg() + default_attr = RigidBodyPhysicsCfg() sugar_box_path = get_data_path("SugarBox/sugar_box_usd/sugar_box.usda") sugar_box: RigidObject = self.sim.add_rigid_object( cfg=RigidObjectCfg( uid="sugar_box", shape=MeshCfg(fpath=sugar_box_path), body_type="dynamic", - use_usd_properties=False, + asset_physics_mode="overlay", init_pos=[0.0, 1.0, 0.1], attrs=default_attr, ) ) + self.sim.prepare() body0 = sugar_box._entities[0].get_physical_body() print(sugar_box._entities[0].get_physical_attr()) assert pytest.approx(body0.get_mass()) == default_attr.mass @@ -75,20 +73,32 @@ 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() + default_drive = JointDrivePropertiesCfg( + drive_type="force", + stiffness=1e4, + damping=1e3, + max_effort=1e10, + max_velocity=1e10, + friction=0.0, + armature=0.0, + ) h1_path = get_data_path("UnitreeH1Usd/H1_usd/h1.usd") h1: Articulation = self.sim.add_articulation( cfg=ArticulationCfg( uid="h1", fpath=h1_path, build_pk_chain=False, - use_usd_properties=False, + asset_physics_mode="overlay", init_pos=[0.0, 0.0, 1.2], - drive_pros=default_drive, + joint_drive_props=default_drive, ) ) + self.sim.prepare() stiffness = h1.body_data.joint_stiffness damping = h1.body_data.joint_damping @@ -106,17 +116,18 @@ def test_import_articulation(self): ) def test_usd_properties(self): - """In this test, we set use_usd_properties=True to verify that the USD properties are correctly applied.""" + """Verify that preserve mode keeps physics authored in USD assets.""" h1_path = get_data_path("UnitreeH1Usd/H1_usd/h1.usd") h1: Articulation = self.sim.add_articulation( cfg=ArticulationCfg( uid="h1_beta", fpath=h1_path, build_pk_chain=False, - use_usd_properties=True, + asset_physics_mode="preserve", init_pos=[1.0, 0.0, 1.2], ) ) + self.sim.prepare() stiffness = h1.body_data.joint_stiffness damping = h1.body_data.joint_damping @@ -152,14 +163,14 @@ def test_usd_properties(self): uid="sugar_box_beta", shape=MeshCfg(fpath=sugar_box_path), body_type="dynamic", - use_usd_properties=True, + asset_physics_mode="preserve", init_pos=[1.0, 1.0, 0.1], ) ) body0 = sugar_box._entities[0].get_physical_body() print(sugar_box._entities[0].get_physical_attr()) assert pytest.approx(body0.get_mass(), 0.001) == 0.514 - # TODO: nvidia physx attrs in usd currently are not fully suported + # TODO: vendor-specific rigid-body attributes in USD are not fully supported. # assert(body0.get_linear_damping()==0) # assert(body0.get_angular_damping()==0.05) # assert(body0.get_solver_iteration_counts()==(4, 1)) @@ -180,13 +191,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_curobo_integration.py b/tests/sim/planners/test_curobo_integration.py index 5c941900c..a49c84d14 100644 --- a/tests/sim/planners/test_curobo_integration.py +++ b/tests/sim/planners/test_curobo_integration.py @@ -36,7 +36,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg # noqa: E402 from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 from embodichain.lab.sim.robots import FrankaPandaCfg # noqa: E402 -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 +from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg # noqa: E402 from embodichain.lab.sim.shapes import CubeCfg # noqa: E402 from embodichain.lab.sim.planners import ( # noqa: E402 MotionGenCfg, @@ -74,12 +74,13 @@ def _make_sim_robot(num_envs: int = 1): cfg=RigidObjectCfg( uid="demo_block", shape=CubeCfg(size=DEMO_BLOCK_DIMS), - attrs=RigidBodyAttributesCfg(), - body_type="kinematic", + attrs=RigidBodyPhysicsCfg(), + body_type="static", init_pos=DEMO_BLOCK_POS, init_rot=[0.0, 0.0, 0.0], ) ) + sim.prepare() return sim, robot, block diff --git a/tests/sim/planners/test_curobo_planner.py b/tests/sim/planners/test_curobo_planner.py index 552ef6dfc..c4acf045c 100644 --- a/tests/sim/planners/test_curobo_planner.py +++ b/tests/sim/planners/test_curobo_planner.py @@ -34,6 +34,7 @@ import yaml from embodichain.lab.sim.planners import CuroboPlannerCfg +from embodichain.utils.math import matrix_from_quat from embodichain.lab.sim.planners.curobo.curobo_planner import ( CuroboPlanOptions, CuroboPlanner, @@ -127,9 +128,14 @@ def test_public_config_imports_without_curobo(): def test_matrix_to_position_quaternion_uses_wxyz(): matrix = torch.eye(4).unsqueeze(0) + xyzw = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) / math.sqrt(30.0) + matrix[:, :3, :3] = matrix_from_quat(xyzw) position, quaternion = _matrix_to_position_quaternion(matrix) assert torch.equal(position, torch.zeros(1, 3)) - assert torch.equal(quaternion, torch.tensor([[1.0, 0.0, 0.0, 0.0]])) + torch.testing.assert_close( + quaternion, + torch.tensor([[4.0, 1.0, 2.0, 3.0]]) / math.sqrt(30.0), + ) assert position.is_contiguous() assert quaternion.is_contiguous() @@ -468,7 +474,7 @@ def _identity_pose( translation: tuple[float, float, float] = (0.45, 0.0, 0.18), ) -> torch.Tensor: return torch.tensor( - [*translation, 1.0, 0.0, 0.0, 0.0], + [*translation, 0.0, 0.0, 0.0, 1.0], dtype=torch.float32, ) @@ -530,7 +536,7 @@ def test_cuboid_entry_off_origin_mesh_offsets_center(): def test_cuboid_entry_rotated_pose_preserves_center(): quaternion = torch.tensor( - [math.cos(math.pi / 4), 0.0, 0.0, math.sin(math.pi / 4)], + [0.0, 0.0, math.sin(math.pi / 4), math.cos(math.pi / 4)], dtype=torch.float32, ) pose = torch.cat([torch.tensor([0.45, 0.0, 0.18]), quaternion]) @@ -543,7 +549,9 @@ def test_cuboid_entry_rotated_pose_preserves_center(): )[0] assert fields["pose"][:3] == pytest.approx([0.45, 0.0, 0.18]) - assert fields["pose"][3:] == pytest.approx(quaternion.tolist()) + assert fields["pose"][3:] == pytest.approx( + [math.cos(math.pi / 4), 0.0, 0.0, math.sin(math.pi / 4)] + ) def test_cuboid_entry_accepts_homogeneous_pose(): @@ -572,7 +580,7 @@ def test_mesh_entry_serializes_flat_face_buffer(): assert (top_key, name) == ("mesh", "demo_block") assert len(fields["vertices"]) == 8 assert len(fields["faces"]) == 36 - assert fields["pose"] == pytest.approx(_identity_pose().tolist()) + assert fields["pose"] == pytest.approx([0.45, 0.0, 0.18, 1.0, 0.0, 0.0, 0.0]) def test_invalid_obstacle_representation_raises(): @@ -927,7 +935,7 @@ def test_generated_mesh_yaml_loads_in_curobo_scene_cfg(tmp_path): def _build_curobo_scene(sim_device: str = "cuda") -> tuple[object, object, object]: from embodichain.lab.sim import SimulationManager, SimulationManagerCfg - from embodichain.lab.sim.cfg import RigidBodyAttributesCfg + from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg from embodichain.lab.sim.objects import RigidObjectCfg from embodichain.lab.sim.robots import FrankaPandaCfg from embodichain.lab.sim.shapes import CubeCfg @@ -948,12 +956,13 @@ def _build_curobo_scene(sim_device: str = "cuda") -> tuple[object, object, objec cfg=RigidObjectCfg( uid="block", shape=CubeCfg(size=_SIM_BLOCK_DIMS), - attrs=RigidBodyAttributesCfg(), - body_type="kinematic", + attrs=RigidBodyPhysicsCfg(), + body_type="static", init_pos=_SIM_BLOCK_POS, init_rot=(0.0, 0.0, 0.0), ) ) + sim.prepare() return sim, robot, block diff --git a/tests/sim/planners/test_motion_generator.py b/tests/sim/planners/test_motion_generator.py index d04cb42bb..628eac77c 100644 --- a/tests/sim/planners/test_motion_generator.py +++ b/tests/sim/planners/test_motion_generator.py @@ -53,7 +53,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) @@ -97,6 +97,7 @@ def setup_simulation(self): cls.robot: Robot = cls.robot_sim.add_robot( cfg=CobotMagicCfg.from_dict(cfg_dict) ) + cls.robot_sim.prepare() cls.arm_name = "left_arm" diff --git a/tests/sim/planners/test_neural_planner.py b/tests/sim/planners/test_neural_planner.py index c1e2d2786..dc3b6d6f6 100644 --- a/tests/sim/planners/test_neural_planner.py +++ b/tests/sim/planners/test_neural_planner.py @@ -99,7 +99,7 @@ def compute_fk( batch = qpos.shape[0] if qpos.dim() > 1 else 1 if to_matrix: return torch.eye(4).repeat(batch, 1, 1) - return torch.tensor([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]]).repeat(batch, 1) + return torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]).repeat(batch, 1) class FakeSimulationManager: diff --git a/tests/sim/planners/test_toppra_batched.py b/tests/sim/planners/test_toppra_batched.py index a769327ce..8b097a89d 100644 --- a/tests/sim/planners/test_toppra_batched.py +++ b/tests/sim/planners/test_toppra_batched.py @@ -140,13 +140,14 @@ def _make_planner(self): from embodichain.lab.sim.robots import CobotMagicCfg sim = SimulationManager( - SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=2) + SimulationManagerCfg(headless=True, device="cpu", num_envs=2) ) robot = sim.add_robot( cfg=CobotMagicCfg.from_dict( {"uid": "t", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} ) ) + sim.prepare() planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="t", max_workers=1)) return planner, sim @@ -244,13 +245,14 @@ def test_plan_batched_pool_path(self, mp_context): from embodichain.lab.sim.robots import CobotMagicCfg sim = SimulationManager( - SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=3) + SimulationManagerCfg(headless=True, device="cpu", num_envs=3) ) sim.add_robot( cfg=CobotMagicCfg.from_dict( {"uid": "p", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} ) ) + sim.prepare() planner = ToppraPlanner( ToppraPlannerCfg(robot_uid="p", max_workers=2, mp_context=mp_context) ) @@ -303,7 +305,7 @@ def test_workers_reaped_on_gc(self, mp_context): from embodichain.lab.sim.robots import CobotMagicCfg sim = SimulationManager( - SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=3) + SimulationManagerCfg(headless=True, device="cpu", num_envs=3) ) sim.add_robot( cfg=CobotMagicCfg.from_dict( @@ -314,6 +316,7 @@ def test_workers_reaped_on_gc(self, mp_context): } ) ) + sim.prepare() planner = ToppraPlanner( ToppraPlannerCfg( robot_uid="close_reap", max_workers=2, mp_context=mp_context @@ -369,13 +372,14 @@ def test_batched_equals_inline_single(self): from embodichain.lab.sim.robots import CobotMagicCfg sim = SimulationManager( - SimulationManagerCfg(headless=True, sim_device="cpu", num_envs=4) + SimulationManagerCfg(headless=True, device="cpu", num_envs=4) ) sim.add_robot( cfg=CobotMagicCfg.from_dict( {"uid": "r", "init_pos": [0, 0, 0.7775], "init_qpos": [0.0] * 16} ) ) + sim.prepare() planner = ToppraPlanner(ToppraPlannerCfg(robot_uid="r", max_workers=1)) try: B, dofs = 4, 6 diff --git a/tests/sim/planners/test_toppra_planner.py b/tests/sim/planners/test_toppra_planner.py index ce343cdb4..c7165f186 100644 --- a/tests/sim/planners/test_toppra_planner.py +++ b/tests/sim/planners/test_toppra_planner.py @@ -36,7 +36,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 = { @@ -45,6 +45,7 @@ def setup_simulation(self): "init_qpos": [0.0] * 16, } cls.robot = cls.sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + cls.sim.prepare() def setup_method(self): self.setup_simulation() diff --git a/tests/sim/robots/__init__.py b/tests/sim/robots/__init__.py new file mode 100644 index 000000000..dd650e902 --- /dev/null +++ b/tests/sim/robots/__init__.py @@ -0,0 +1,15 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- diff --git a/tests/sim/robots/test_entrypoints.py b/tests/sim/robots/test_entrypoints.py new file mode 100644 index 000000000..52d6a2fe3 --- /dev/null +++ b/tests/sim/robots/test_entrypoints.py @@ -0,0 +1,87 @@ +# ---------------------------------------------------------------------------- +# 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 pathlib import Path +import runpy +import sys +from types import ModuleType + +import pytest + +import embodichain.lab.sim as sim_module +from embodichain.lab.sim.cfg import NewtonPhysicsCfg + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +_ROBOT_ENTRYPOINTS = ( + "embodichain/lab/sim/robots/cobotmagic.py", + "embodichain/lab/sim/robots/franka_panda.py", + "embodichain/lab/sim/robots/ur_robot.py", + "embodichain/lab/sim/robots/dual_arm.py", + "embodichain/lab/sim/robots/dexforce_w1/cfg.py", +) + + +@pytest.mark.parametrize("relative_path", _ROBOT_ENTRYPOINTS) +def test_robot_entrypoint_selects_newton_backend( + monkeypatch: pytest.MonkeyPatch, + relative_path: str, +) -> None: + """Each robot smoke program must forward ``--physics newton``.""" + captured: dict[str, object] = {} + + class SimulationManagerCfgSpy: + def __init__(self, **kwargs: object) -> None: + captured.update(kwargs) + + class SimulationManagerSpy: + def __init__(self, _cfg: SimulationManagerCfgSpy) -> None: + pass + + def add_robot(self, *, cfg: object) -> object: + return cfg + + def prepare(self) -> None: + pass + + def update(self, *, step: int) -> None: + pass + + def open_window(self) -> None: + pass + + def destroy(self) -> None: + pass + + monkeypatch.setattr(sim_module, "SimulationManagerCfg", SimulationManagerCfgSpy) + monkeypatch.setattr(sim_module, "SimulationManager", SimulationManagerSpy) + ipython_module = ModuleType("IPython") + ipython_module.embed = lambda: None # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "IPython", ipython_module) + monkeypatch.setattr( + sys, + "argv", + [relative_path, "--physics", "newton"], + ) + + original_path_entry = sys.path[0] + try: + runpy.run_path(_REPOSITORY_ROOT / relative_path, run_name="__main__") + finally: + sys.path[0] = original_path_entry + + assert isinstance(captured["physics_cfg"], NewtonPhysicsCfg) diff --git a/tests/sim/sensors/test_camera.py b/tests/sim/sensors/test_camera.py index f9d522e4b..9e420e50c 100644 --- a/tests/sim/sensors/test_camera.py +++ b/tests/sim/sensors/test_camera.py @@ -20,6 +20,9 @@ import torch import os +from types import SimpleNamespace +from unittest.mock import MagicMock + from tensordict import TensorDict from embodichain.lab.sim import SimulationManager, SimulationManagerCfg @@ -27,6 +30,7 @@ from embodichain.lab.sim.objects import Articulation from embodichain.lab.sim.cfg import ArticulationCfg, RenderCfg from embodichain.data import get_data_path +from scripts.tutorials.sim.create_sensor import create_sensor as create_tutorial_sensor FULL_NUM_ENVS = 4 FULL_WIDTH = 640 @@ -50,7 +54,7 @@ def setup_simulation( # Setup SimulationManager config = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=sim_device, render_cfg=RenderCfg(renderer=renderer), num_envs=num_envs, ) @@ -67,6 +71,7 @@ def setup_simulation( } cfg = SensorCfg.from_dict(cfg_dict) self.camera: Camera = self.sim.add_sensor(cfg) + self.sim.prepare() def test_get_data(self): @@ -146,6 +151,7 @@ def test_attach_to_parent(self): uid="test", extrinsics=CameraCfg.ExtrinsicsCfg(parent="handle_xpos") ) ) + assert self.camera.is_attached def test_set_intrinsics(self): # Define new intrinsic parameters @@ -185,6 +191,30 @@ def setup_method(self): self.setup_simulation("cuda", renderer="hybrid") +def test_create_sensor_tutorial_preserves_attached_camera_view() -> None: + """Keep the wrist-camera view stable after the xyzw convention migration.""" + sim = MagicMock() + + create_tutorial_sensor(sim, SimpleNamespace(attach_sensor=True)) + + cfg = sim.add_sensor.call_args.kwargs["sensor_cfg"] + expected_rotation = torch.tensor( + [ + [0.579228, 0.573576, 0.579228], + [0.405580, -0.819152, 0.405580], + [0.707107, 0.0, -0.707107], + ], + dtype=torch.float32, + ) + assert cfg.extrinsics.parent == "ee_link" + torch.testing.assert_close( + cfg.extrinsics.transformation[:3, :3], + expected_rotation, + atol=1.0e-6, + rtol=1.0e-6, + ) + + @pytest.mark.parametrize( ("sim_device", "renderer"), [("cpu", "hybrid"), ("cpu", "fast-rt"), ("cuda", "fast-rt")], @@ -208,6 +238,23 @@ def test_camera_backend_smoke(sim_device, renderer): test.teardown_method() +def test_camera_parent_attachment_cpu() -> None: + """Attach a camera to a materialized articulation link on the CPU backend.""" + test = CameraTest() + test.setup_simulation( + "cpu", + renderer="hybrid", + num_envs=SMOKE_NUM_ENVS, + width=SMOKE_WIDTH, + height=SMOKE_HEIGHT, + enable_auxiliary_data=False, + ) + try: + test.test_attach_to_parent() + finally: + test.teardown_method() + + if __name__ == "__main__": test = TestCameraHybridCUDA() test.setup_method() diff --git a/tests/sim/sensors/test_contact.py b/tests/sim/sensors/test_contact.py index ef5e7f1b3..58df59f1e 100644 --- a/tests/sim/sensors/test_contact.py +++ b/tests/sim/sensors/test_contact.py @@ -26,7 +26,7 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import ( RenderCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.sim.sensors import ( ContactSensorCfg, @@ -43,14 +43,14 @@ class ContactTest: - def setup_simulation(self, sim_device, renderer="hybrid"): + def setup_simulation(self, device, renderer="hybrid"): sim_cfg = SimulationManagerCfg( width=CONTACT_TEST_WIDTH, height=CONTACT_TEST_HEIGHT, 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), ) @@ -69,8 +69,10 @@ def setup_simulation(self, sim_device, renderer="hybrid"): contact_filter_cfg.articulation_cfg_list = [contact_filter_art_cfg] contact_filter_cfg.filter_need_both_actor = True + self.sim.prepare() self.to_grasp_pose(cube2) self.contact_sensor = self.sim.add_sensor(sensor_cfg=contact_filter_cfg) + self.sim.prepare() def create_cube(self, uid: str, position: list = (0.0, 0.0, 0)) -> RigidObject: """create cube @@ -89,12 +91,16 @@ def create_cube(self, uid: str, position: list = (0.0, 0.0, 0)) -> RigidObject: uid=uid, shape=CubeCfg(size=cube_size), body_type="dynamic", - attrs=RigidBodyAttributesCfg( - mass=0.1, - dynamic_friction=0.9, - static_friction=0.95, - restitution=0.01, - sleep_threshold=0.0, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 0.1}, + "rigid_props": {"sleep_threshold": 0.0}, + "material_props": { + "dynamic_friction": 0.9, + "static_friction": 0.95, + "restitution": 0.01, + }, + } ), init_pos=position, ) @@ -125,7 +131,7 @@ def create_robot(self, uid: str, position: list = (0.0, 0.0, 0)) -> Robot: }, "init_pos": position, "init_qpos": [0.0, -1.57, 1.57, -1.57, -1.57, 0.0, 0.0, 0.0], - "drive_pros": { + "joint_drive_props": { "stiffness": {"finger[1-2]_joint": 1e2}, "damping": {"finger[1-2]_joint": 1e1}, "max_effort": {"finger[1-2]_joint": 1e3}, diff --git a/tests/sim/sensors/test_stereo.py b/tests/sim/sensors/test_stereo.py index 156393ed7..704e00120 100644 --- a/tests/sim/sensors/test_stereo.py +++ b/tests/sim/sensors/test_stereo.py @@ -42,7 +42,7 @@ def setup_simulation( # Setup SimulationManager config = SimulationManagerCfg( headless=True, - sim_device=sim_device, + device=sim_device, num_envs=num_envs, render_cfg=RenderCfg(renderer=renderer), ) @@ -61,6 +61,7 @@ def setup_simulation( } cfg = SensorCfg.from_dict(cfg_dict) self.camera: StereoCamera = self.sim.add_sensor(cfg) + self.sim.prepare() def test_get_data(self): diff --git a/tests/sim/skills/test_calls.py b/tests/sim/skills/test_calls.py index 8c0cf0311..0c38ea7d3 100644 --- a/tests/sim/skills/test_calls.py +++ b/tests/sim/skills/test_calls.py @@ -50,7 +50,7 @@ def _identity_pose() -> SemanticPose: - return SemanticPose((0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)) + return SemanticPose((0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0)) def _call_descriptor( @@ -71,35 +71,35 @@ def _call_descriptor( def test_semantic_pose_owns_inputs_and_returns_independent_tensors() -> None: position = torch.tensor([1.0, 2.0, 3.0]) - quaternion = torch.tensor([1.0, 0.0, 0.0, 0.0]) + quaternion = torch.tensor([0.0, 0.0, 0.0, 1.0]) pose = SemanticPose(position, quaternion) position.zero_() quaternion.zero_() returned_position = pose.position - returned_quaternion = pose.quaternion_wxyz + returned_quaternion = pose.quaternion_xyzw returned_position.fill_(9.0) returned_quaternion.fill_(9.0) torch.testing.assert_close(pose.position, torch.tensor([1.0, 2.0, 3.0])) torch.testing.assert_close( - pose.quaternion_wxyz, - torch.tensor([1.0, 0.0, 0.0, 0.0]), + pose.quaternion_xyzw, + torch.tensor([0.0, 0.0, 0.0, 1.0]), ) -def test_semantic_pose_normalizes_wxyz_quaternion() -> None: - pose = SemanticPose((0.0, 0.0, 0.0), (2.0, 0.0, 0.0, 2.0)) +def test_semantic_pose_normalizes_xyzw_quaternion() -> None: + pose = SemanticPose((0.0, 0.0, 0.0), (0.0, 0.0, 2.0, 2.0)) expected = torch.tensor( - [math.sqrt(0.5), 0.0, 0.0, math.sqrt(0.5)], + [0.0, 0.0, math.sqrt(0.5), math.sqrt(0.5)], dtype=torch.float32, ) - torch.testing.assert_close(pose.quaternion_wxyz, expected) + torch.testing.assert_close(pose.quaternion_xyzw, expected) def test_semantic_pose_converts_to_homogeneous_matrix() -> None: - pose = SemanticPose((1.0, 2.0, 3.0), (2.0, 0.0, 0.0, 2.0)) + pose = SemanticPose((1.0, 2.0, 3.0), (0.0, 0.0, 2.0, 2.0)) expected = torch.tensor( [ @@ -115,7 +115,7 @@ def test_semantic_pose_converts_to_homogeneous_matrix() -> None: def test_semantic_call_metadata_is_deterministic_and_json_safe() -> None: call = Place( object=SceneObjectRef("cube"), - at=SemanticPose((1.0, 2.0, 3.0), (1.0, 0.0, 0.0, 0.0)), + at=SemanticPose((1.0, 2.0, 3.0), (0.0, 0.0, 0.0, 1.0)), resources={"primary": "left_arm"}, ) diff --git a/tests/sim/skills/test_compiler.py b/tests/sim/skills/test_compiler.py index f5004d86e..d6b87ed4f 100644 --- a/tests/sim/skills/test_compiler.py +++ b/tests/sim/skills/test_compiler.py @@ -704,7 +704,7 @@ def test_curated_analysis_selects_monitors_per_semantic_call() -> None: Pick(object=SceneObjectRef("cube")), Place( object=SceneObjectRef("cube"), - at=SemanticPose((0.5, 0.0, 0.3), (1.0, 0.0, 0.0, 0.0)), + at=SemanticPose((0.5, 0.0, 0.3), (0.0, 0.0, 0.0, 1.0)), ), ) ) @@ -894,7 +894,7 @@ def test_place_effect_spec_binds_source_and_verified_detach_baseline() -> None: object=SceneObjectRef("cube"), at=SemanticPose( (0.5, -0.2, 0.4), - (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), ), ), ) @@ -1145,7 +1145,7 @@ def test_analysis_is_provider_free_and_propagates_object_target() -> None: preset=_preset("safe", action_option_templates=templates), ), ) - drop = SemanticPose((0.4, 0.2, 0.3), (1.0, 0.0, 0.0, 0.0)) + drop = SemanticPose((0.4, 0.2, 0.3), (0.0, 0.0, 0.0, 1.0)) workflow = compiler.analyze( ( @@ -1193,7 +1193,7 @@ def test_pick_lookahead_uses_downstream_place_orientation_policy() -> None: preset=_preset("safe", action_option_templates=templates), ), ) - drop = SemanticPose((0.4, 0.2, 0.3), (1.0, 0.0, 0.0, 0.0)) + drop = SemanticPose((0.4, 0.2, 0.3), (0.0, 0.0, 0.0, 1.0)) workflow = compiler.analyze( ( Pick(object=SceneObjectRef("cube")), @@ -1429,7 +1429,7 @@ def test_place_uses_verified_object_to_eef_transform() -> None: preset=_preset("safe", action_option_templates=templates), ), ) - drop = SemanticPose((0.5, -0.2, 0.4), (1.0, 0.0, 0.0, 0.0)) + drop = SemanticPose((0.5, -0.2, 0.4), (0.0, 0.0, 0.0, 1.0)) workflow = compiler.analyze((Place(object=SceneObjectRef("cube"), at=drop),)) pick_workflow = compiler.analyze((Pick(object=SceneObjectRef("cube")),)) semantics = compiler.ground( @@ -1477,7 +1477,7 @@ def test_place_can_keep_observed_object_orientation_at_target() -> None: object_to_eef = torch.eye(4).repeat(2, 1, 1) object_to_eef[:, 2, 3] = 0.12 context = _held_context(registry, semantics, object_to_eef) - drop = SemanticPose((0.5, -0.2, 0.4), (1.0, 0.0, 0.0, 0.0)) + drop = SemanticPose((0.5, -0.2, 0.4), (0.0, 0.0, 0.0, 1.0)) workflow = compiler.analyze((Place(object=SceneObjectRef("cube"), at=drop),)) grounded = compiler.ground(workflow, 0, context) @@ -1527,7 +1527,7 @@ def test_place_rejects_wrong_or_inactive_verified_holder() -> None: ( Place( object=SceneObjectRef("cube"), - at=SemanticPose((0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0)), + at=SemanticPose((0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 1.0)), ), ) ) @@ -1597,7 +1597,7 @@ def test_registered_lowerer_is_explicit_and_opaque_to_lookahead() -> None: registered, Place( object=SceneObjectRef("cube"), - at=SemanticPose((0.3, 0.0, 0.2), (1.0, 0.0, 0.0, 0.0)), + at=SemanticPose((0.3, 0.0, 0.2), (0.0, 0.0, 0.0, 1.0)), ), ) ) @@ -1614,11 +1614,11 @@ def test_registered_lowerer_can_certify_retained_object_lookahead() -> None: registry, _ = _scene_registry() registered_target = SemanticPose( (0.25, 0.1, 0.4), - (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), ) place_target = SemanticPose( (0.3, 0.0, 0.2), - (1.0, 0.0, 0.0, 0.0), + (0.0, 0.0, 0.0, 1.0), ) compiler, _ = _compiler( registry, diff --git a/tests/sim/skills/test_curobo_semantic_runtime_dynamic_recovery_gpu.py b/tests/sim/skills/test_curobo_semantic_runtime_dynamic_recovery_gpu.py index 078cbf3c7..fb8070bec 100644 --- a/tests/sim/skills/test_curobo_semantic_runtime_dynamic_recovery_gpu.py +++ b/tests/sim/skills/test_curobo_semantic_runtime_dynamic_recovery_gpu.py @@ -48,7 +48,7 @@ SkillDescriptor, ) from embodichain.lab.sim.atomic_actions.tracking import TrackingPolicy # noqa: E402 -from embodichain.lab.sim.cfg import RigidBodyAttributesCfg # noqa: E402 +from embodichain.lab.sim.cfg import RigidBodyPhysicsCfg # noqa: E402 from embodichain.lab.sim.objects import RigidObjectCfg # noqa: E402 from embodichain.lab.sim.planners import MotionGenCfg, MotionGenerator # noqa: E402 from embodichain.lab.sim.planners.curobo.curobo_planner import ( # noqa: E402 @@ -241,7 +241,7 @@ def test_semantic_runtime_replans_after_dynamic_curobo_world_change() -> None: cfg=RigidObjectCfg( uid=OBSTACLE_UID, shape=CubeCfg(size=OBSTACLE_SIZE), - attrs=RigidBodyAttributesCfg(), + attrs=RigidBodyPhysicsCfg(), body_type="kinematic", init_pos=OBSTACLE_START_POSITION, init_rot=[0.0, 0.0, 0.0], diff --git a/tests/sim/solvers/test_differential_solver.py b/tests/sim/solvers/test_differential_solver.py index c055ad617..00566eca2 100644 --- a/tests/sim/solvers/test_differential_solver.py +++ b/tests/sim/solvers/test_differential_solver.py @@ -33,7 +33,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 @@ -61,6 +61,7 @@ def setup_simulation(self, solver_type: str): } self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + self.sim.prepare() @pytest.mark.parametrize("arm_name", ["left_arm", "right_arm"]) def test_differential_solver(self, arm_name: str): diff --git a/tests/sim/solvers/test_neural_ik_solver.py b/tests/sim/solvers/test_neural_ik_solver.py index 324c329e9..7aa72b942 100644 --- a/tests/sim/solvers/test_neural_ik_solver.py +++ b/tests/sim/solvers/test_neural_ik_solver.py @@ -52,7 +52,7 @@ class TestNeuralIKSolver: def _setup(self, tmp_path): checkpoint_path = _create_fake_checkpoint(tmp_path) - config = SimulationManagerCfg(headless=True, sim_device="cpu") + config = SimulationManagerCfg(headless=True, device="cpu") self.sim = SimulationManager(config) cfg = FrankaPandaCfg.from_dict({"robot_type": "panda"}) @@ -75,6 +75,7 @@ def _setup(self, tmp_path): ) self.robot = self.sim.add_robot(cfg=cfg) + self.sim.prepare() self.sim.update(step=100) def teardown_method(self): diff --git a/tests/sim/solvers/test_opw_solver.py b/tests/sim/solvers/test_opw_solver.py index 3db7e9fb2..28c6ef37d 100644 --- a/tests/sim/solvers/test_opw_solver.py +++ b/tests/sim/solvers/test_opw_solver.py @@ -70,8 +70,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) @@ -126,6 +126,7 @@ def setup_simulation(self, sim_device): } self.robot: Robot = self.sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + self.sim.prepare() @pytest.mark.parametrize("arm_name", ["left_arm", "right_arm"]) def test_ik(self, arm_name: str): diff --git a/tests/sim/solvers/test_pink_solver.py b/tests/sim/solvers/test_pink_solver.py index 957542d40..c53d57165 100644 --- a/tests/sim/solvers/test_pink_solver.py +++ b/tests/sim/solvers/test_pink_solver.py @@ -66,7 +66,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) @@ -97,6 +97,7 @@ def setup_simulation(self, solver_type: str): } self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + self.sim.prepare() def test_differential_solver(self): # Test differential solver with a 1x4x4 homogeneous matrix pose and a joint_seed diff --git a/tests/sim/solvers/test_pinocchio_solver.py b/tests/sim/solvers/test_pinocchio_solver.py index 730d5bf59..3fd57e8b3 100644 --- a/tests/sim/solvers/test_pinocchio_solver.py +++ b/tests/sim/solvers/test_pinocchio_solver.py @@ -33,9 +33,12 @@ 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) + # Keep the scene fixed while FK/IK operate on the same robot state. + # Automatic stepping can race with the two solver calls and make the + # reconstructed pose depend on test timing. + self.sim.set_manual_update(True) # Load robot URDF file urdf = get_data_path("DexforceW1V021/DexforceW1_v02_1.urdf") @@ -62,6 +65,7 @@ def setup_simulation(self, solver_type: str): } self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + self.sim.prepare() @pytest.mark.parametrize("arm_name", ["left_arm", "right_arm"]) def test_ik(self, arm_name: str): diff --git a/tests/sim/solvers/test_pytorch_solver.py b/tests/sim/solvers/test_pytorch_solver.py index be720e7ce..a2e742a84 100644 --- a/tests/sim/solvers/test_pytorch_solver.py +++ b/tests/sim/solvers/test_pytorch_solver.py @@ -73,7 +73,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 @@ -104,6 +104,7 @@ def setup_simulation(self, solver_type: str): } self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + self.sim.prepare() # Wait for robot to stabilize. self.sim.update(step=100) diff --git a/tests/sim/solvers/test_srs_solver.py b/tests/sim/solvers/test_srs_solver.py index f760ed616..888abc3ba 100644 --- a/tests/sim/solvers/test_srs_solver.py +++ b/tests/sim/solvers/test_srs_solver.py @@ -401,7 +401,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 @@ -426,7 +426,7 @@ def setup_simulation(self, solver_type: str, device: str = "cpu"): "torso": ["ANKLE", "KNEE", "BUTTOCK", "WAIST"], "head": [f"NECK{i + 1}" for i in range(2)], }, - "drive_pros": { + "joint_drive_props": { "stiffness": { "LEFT_J[1-7]": 1e4, "RIGHT_J[1-7]": 1e4, @@ -453,14 +453,18 @@ def setup_simulation(self, solver_type: str, device: str = "cpu"): }, }, "attrs": { - "mass": 1e-1, - "static_friction": 0.95, - "dynamic_friction": 0.9, - "linear_damping": 0.7, - "angular_damping": 0.7, - "max_depenetration_velocity": 10.0, - "min_position_iters": 32, - "min_velocity_iters": 8, + "mass_props": {"mass": 1e-1}, + "rigid_props": { + "linear_damping": 0.7, + "angular_damping": 0.7, + "max_depenetration_velocity": 10.0, + "min_position_iters": 32, + "min_velocity_iters": 8, + }, + "material_props": { + "static_friction": 0.95, + "dynamic_friction": 0.9, + }, }, "solver_cfg": { "left_arm": { @@ -489,6 +493,7 @@ def setup_simulation(self, solver_type: str, device: str = "cpu"): } self.robot: Robot = self.sim.add_robot(cfg=RobotCfg.from_dict(cfg_dict)) + self.sim.prepare() # Wait for robot to stabilize. self.sim.update(step=100) diff --git a/tests/sim/solvers/test_ur_solver.py b/tests/sim/solvers/test_ur_solver.py index 522a2e5cc..69d4da5b0 100644 --- a/tests/sim/solvers/test_ur_solver.py +++ b/tests/sim/solvers/test_ur_solver.py @@ -29,7 +29,6 @@ JointDrivePropertiesCfg, RobotCfg, LightCfg, - RigidBodyAttributesCfg, RigidObjectCfg, URDFCfg, ) @@ -79,8 +78,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) @@ -95,7 +94,7 @@ def setup_simulation(self, sim_device): {"component_type": "hand", "urdf_path": gripper_urdf_path}, ] ), - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness={"Joint[0-9]": 1e4, "FINGER[1-2]": 1e3}, damping={"Joint[0-9]": 1e3, "FINGER[1-2]": 1e2}, max_effort={"Joint[0-9]": 1e5, "FINGER[1-2]": 1e4}, @@ -129,6 +128,7 @@ def setup_simulation(self, sim_device): init_pos=(0, 0, 0), ) self.robot: Robot = self.sim.add_robot(cfg=cfg) + self.sim.prepare() def test_ik(self): # Test inverse kinematics (IK) with a 1x4x4 homogeneous matrix pose and a joint_seed diff --git a/tests/sim/spawn/__init__.py b/tests/sim/spawn/__init__.py new file mode 100644 index 000000000..19567d22d --- /dev/null +++ b/tests/sim/spawn/__init__.py @@ -0,0 +1,19 @@ +# ---------------------------------------------------------------------------- +# 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 EmbodiChain Spawn descriptor translation.""" + +from __future__ import annotations diff --git a/tests/sim/spawn/test_create_robot_integration.py b/tests/sim/spawn/test_create_robot_integration.py new file mode 100644 index 000000000..f1c6028a4 --- /dev/null +++ b/tests/sim/spawn/test_create_robot_integration.py @@ -0,0 +1,155 @@ +# ---------------------------------------------------------------------------- +# 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. +# ---------------------------------------------------------------------------- + +"""Regression coverage for robots configured by simulation tutorials.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import dexsim +from embodichain.lab.sim.spawn.descriptors import ( + articulation_desc_from_cfg, + configure_articulation_desc, +) +from embodichain.lab.sim.spawn.scene import SpawnScene +from scripts.tutorials.sim.create_sensor import create_robot as create_sensor_robot +from scripts.tutorials.sim.create_robot import create_robot + +pytestmark = pytest.mark.requires_sim + +ARM_BASE_MASS = 3.167 # SR5 base_link inertial mass from the source URDF. +ARM_BASE_INERTIA = (5.677594, 30.912516, 31.167990) +ARM_STIFFNESS = 1.0e4 +ARM_DAMPING = 1.5e3 +ARM_MAX_EFFORT = 1.0e4 + + +class _ConfigCapture: + def add_robot(self, cfg): + return cfg + + +def _resolve_tutorial_properties(world, cfg): + scene = SpawnScene(world, num_envs=1) + scene.builder.prepare_arenas() + descriptor = articulation_desc_from_cfg(cfg, per_env=False) + scene.declare( + "articulation", + "robot", + descriptor, + configure_source=lambda value: configure_articulation_desc(value, cfg), + ) + result = scene.commit() + descriptor = scene.handles("robot")[0].desc + + base = descriptor.get_link_desc("arm_base_link") + joint = descriptor.get_joint_desc("joint1") + properties = ( + base.rigid_body.mass, + base.rigid_body.inertia.copy(), + joint.dexsim.stiffness, + joint.dexsim.damping, + joint.dexsim.max_force, + joint.newton.target_ke, + joint.newton.target_kd, + joint.effort_limit, + ) + result.close() + return properties + + +def test_create_robot_preserves_source_inertia_and_arm_drive() -> None: + cfg = create_robot(_ConfigCapture()) + cfg.fpath = cfg.urdf_cfg.assemble_urdf() + + config = dexsim.WorldConfig() + config.open_windows = False + config.renderer = dexsim.types.Renderer.HYBRID + config.backend = dexsim.types.Backend.VULKAN + world = dexsim.World(config) + + ( + mass, + inertia, + stiffness, + damping, + max_effort, + newton_ke, + newton_kd, + common_max_effort, + ) = _resolve_tutorial_properties(world, cfg) + + assert mass == pytest.approx(ARM_BASE_MASS) + np.testing.assert_allclose( + inertia, + ARM_BASE_INERTIA, + rtol=1.0e-5, + ) + assert stiffness == pytest.approx(ARM_STIFFNESS) + assert damping == pytest.approx(ARM_DAMPING) + assert max_effort == pytest.approx(ARM_MAX_EFFORT) + assert newton_ke == pytest.approx(ARM_STIFFNESS) + assert newton_kd == pytest.approx(ARM_DAMPING) + assert common_max_effort == pytest.approx(ARM_MAX_EFFORT) + + +def test_create_sensor_uses_the_matched_arm_drive() -> None: + """Keep the sensor tutorial's arm controller aligned across backends.""" + cfg = create_sensor_robot(_ConfigCapture()) + + assert cfg.joint_drive_props is not None + assert cfg.joint_drive_props.max_effort == { + "joint[1-6]": ARM_MAX_EFFORT, + "LEFT_.*": ARM_MAX_EFFORT, + } + cfg.fpath = cfg.urdf_cfg.assemble_urdf() + + config = dexsim.WorldConfig() + config.open_windows = False + config.renderer = dexsim.types.Renderer.HYBRID + config.backend = dexsim.types.Backend.VULKAN + world = dexsim.World(config) + + ( + _, + _, + stiffness, + damping, + max_effort, + newton_ke, + newton_kd, + common_max_effort, + ) = _resolve_tutorial_properties(world, cfg) + + assert ( + stiffness, + damping, + max_effort, + newton_ke, + newton_kd, + common_max_effort, + ) == pytest.approx( + ( + ARM_STIFFNESS, + ARM_DAMPING, + ARM_MAX_EFFORT, + ARM_STIFFNESS, + ARM_DAMPING, + ARM_MAX_EFFORT, + ) + ) diff --git a/tests/sim/spawn/test_descriptors.py b/tests/sim/spawn/test_descriptors.py new file mode 100644 index 000000000..9c6189bee --- /dev/null +++ b/tests/sim/spawn/test_descriptors.py @@ -0,0 +1,1816 @@ +# ---------------------------------------------------------------------------- +# 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 solver-aware Spawn descriptor translation.""" + +from __future__ import annotations + +import copy +import warnings +from dataclasses import fields, is_dataclass +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np +import pytest + +import dexsim +from dexsim.types import DriveType +from dexsim.spawn import ( + ArticulationDesc, + ClothDesc, + CollisionDesc, + CollisionApproximation, + DexsimCollisionDesc, + DexsimClothPhysicsDesc, + DexsimJointDesc, + DexsimPhysicsDesc, + DexsimSoftBodyPhysicsDesc, + JointDesc, + LinkDesc, + NewtonCollisionDesc, + NewtonJointDesc, + ObjectDesc, + RigidBodyPhysicsDesc, + SoftBodyDesc, +) + +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + ArticulationRootPropertiesCfg, + ClothObjectCfg, + ClothPhysicalAttributesCfg, + CollisionPropertiesCfg, + DefaultCollisionPropertiesCfg, + DefaultRigidBodyPropertiesCfg, + JointDrivePropertiesCfg, + LinkPhysicsOverrideCfg, + MassPropertiesCfg, + MeshCollisionCfg, + NewtonCollisionPropertiesCfg, + NewtonJointDrivePropertiesCfg, + NewtonRigidBodyMaterialCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + RobotCfg, + SoftbodyPhysicalAttributesCfg, + SoftbodyVoxelAttributesCfg, + SoftObjectCfg, +) +from embodichain.lab.sim.shapes import CubeCfg, LoadOption, MeshCfg +from embodichain.lab.sim.objects import Articulation +from embodichain.lab.sim.spawn.descriptors import ( + articulation_desc_from_cfg, + cloth_desc_from_cfg, + configure_articulation_desc, + rigid_desc_from_cfg, + soft_desc_from_cfg, +) +from embodichain.lab.sim.spawn.usd import ( + articulation_desc_from_usd, + rigid_desc_from_usd, +) + +pytestmark = pytest.mark.no_sim + +RESTITUTION = 0.25 +DEFORMABLE_MESH_PATH = "/assets/deformable.obj" + + +def test_soft_descriptor_projects_current_dexsim_particle_schema() -> None: + youngs = 1.0e5 + poissons = 0.4 + density = 75.0 + dynamic_friction = 0.2 + min_position_iters = 8 + simplify_target = 40 + remesh_resolution = 12 + voxel_resolution = 16 + cfg = SoftObjectCfg( + uid="soft", + shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + voxel_attr=SoftbodyVoxelAttributesCfg( + triangle_remesh_resolution=remesh_resolution, + triangle_simplify_target=simplify_target, + simulation_mesh_resolution=voxel_resolution, + ), + physical_attr=SoftbodyPhysicalAttributesCfg( + youngs=youngs, + poissons=poissons, + density=density, + dynamic_friction=dynamic_friction, + min_position_iters=min_position_iters, + ), + ) + + descriptor, materials = soft_desc_from_cfg(cfg, per_env=False) + + assert isinstance(descriptor, SoftBodyDesc) + assert descriptor.mesh.file_path == DEFORMABLE_MESH_PATH + assert descriptor.per_env is False + assert descriptor.meshing is not None + assert descriptor.meshing.proxy_simplify_target == simplify_target + assert descriptor.meshing.proxy_remesh_resolution == remesh_resolution + assert descriptor.meshing.voxel_resolution == voxel_resolution + assert descriptor.physics.volume_density == density + assert descriptor.physics.k_mu == pytest.approx(youngs / (2.0 * (1.0 + poissons))) + assert descriptor.physics.k_lambda == pytest.approx( + youngs * poissons / ((1.0 + poissons) * (1.0 - 2.0 * poissons)) + ) + assert isinstance(descriptor.physics.dexsim, DexsimSoftBodyPhysicsDesc) + assert descriptor.physics.dexsim.dynamic_friction == dynamic_friction + assert descriptor.physics.dexsim.min_position_iters == min_position_iters + assert materials == {} + + +def test_cloth_descriptor_projects_current_dexsim_particle_schema() -> None: + density = 2.5 + mass = 0.05 + thickness = 0.02 + bending_stiffness = 0.1 + cfg = ClothObjectCfg( + uid="cloth", + shape=MeshCfg(fpath=DEFORMABLE_MESH_PATH), + physical_attr=ClothPhysicalAttributesCfg( + density=density, + mass=mass, + thickness=thickness, + bending_stiffness=bending_stiffness, + ), + ) + + descriptor, materials = cloth_desc_from_cfg(cfg, per_env=False) + + assert isinstance(descriptor, ClothDesc) + assert descriptor.mesh.file_path == DEFORMABLE_MESH_PATH + assert descriptor.per_env is False + assert descriptor.physics.surface_density == density + assert isinstance(descriptor.physics.dexsim, DexsimClothPhysicsDesc) + assert descriptor.physics.dexsim.mass == mass + assert descriptor.physics.dexsim.thickness == thickness + assert descriptor.physics.dexsim.bending_stiffness == bending_stiffness + assert materials == {} + + +def _resolved_articulation_desc() -> ArticulationDesc: + source_inertia = np.ones(3, dtype=np.float32) + base = LinkDesc( + "base", + "", + np.eye(4, dtype=np.float32), + collisions=[CollisionDesc()], + rigid_body=RigidBodyPhysicsDesc.dynamic( + mass=0.5, + inertia=source_inertia, + ), + ) + finger = LinkDesc( + "finger_left", + "base", + np.eye(4, dtype=np.float32), + collisions=[CollisionDesc()], + rigid_body=RigidBodyPhysicsDesc.dynamic( + mass=0.25, + inertia=source_inertia, + ), + ) + return ArticulationDesc( + name="robot", + links=[base, finger], + joints=[ + JointDesc( + "arm_joint", + "base", + "finger_left", + dexsim.engine.JointType.REVOLUTE, + ) + ], + root_link_name="base", + ) + + +def _assert_property_tree_equal(actual: object, expected: object) -> None: + if isinstance(expected, np.ndarray): + np.testing.assert_array_equal(actual, expected) + elif is_dataclass(expected): + assert type(actual) is type(expected) + for field in fields(expected): + _assert_property_tree_equal( + getattr(actual, field.name), + getattr(expected, field.name), + ) + elif isinstance(expected, dict): + assert actual.keys() == expected.keys() + for key, value in expected.items(): + _assert_property_tree_equal(actual[key], value) + elif isinstance(expected, (list, tuple)): + assert type(actual) is type(expected) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected): + _assert_property_tree_equal(actual_item, expected_item) + else: + assert actual == expected + + +@pytest.mark.parametrize( + ("solver_type", "expected_restitution"), + [ + ("mujoco_warp", None), + ("semi_implicit", None), + ("featherstone", None), + ("xpbd", RESTITUTION), + (None, RESTITUTION), + ], +) +def test_rigid_descriptor_projects_restitution_only_to_supported_solvers( + solver_type: str | None, + expected_restitution: float | None, +) -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(restitution=RESTITUTION) + ), + ) + + descriptor, _ = rigid_desc_from_cfg( + cfg, + newton_solver_type=solver_type, + ) + + newton = descriptor.collisions[0].newton + if expected_restitution is None: + assert newton is None + else: + assert newton.restitution == expected_restitution + + +def test_rigid_descriptor_preserves_default_backend_restitution() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(restitution=RESTITUTION) + ), + ) + + descriptor, _ = rigid_desc_from_cfg( + cfg, + newton_solver_type="mujoco_warp", + ) + + assert descriptor.collisions[0].dexsim.restitution == RESTITUTION + + +def test_flat_rigid_physics_is_rejected_at_config_boundary() -> None: + with pytest.raises(ValueError, match="Removed flat rigid-body attrs fields"): + RigidObjectCfg.from_dict( + { + "uid": "cube", + "shape": {"shape_type": "Cube", "size": [0.1, 0.1, 0.1]}, + "attrs": {"mass": 2.0}, + } + ) + + +def test_rigid_descriptor_authors_mass_or_density_exclusively() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 1.0, "density": 1.0}} + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert descriptor.physics.mass == 1.0 + assert descriptor.physics.density is None + + +def test_rigid_descriptor_forwards_explicit_mass_properties() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg( + mass=2.0, + inertia=[1.0, 2.0, 3.0], + com_position=[0.1, 0.2, 0.3], + com_quaternion=[1.0, 2.0, 3.0, 4.0], + ), + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + np.testing.assert_array_equal(descriptor.physics.inertia, [1.0, 2.0, 3.0]) + np.testing.assert_allclose( + descriptor.physics.com_position, + [0.1, 0.2, 0.3], + ) + np.testing.assert_allclose( + descriptor.physics.com_quaternion, + np.array([4.0, 1.0, 2.0, 3.0]) / np.sqrt(30.0), + ) + + +@pytest.mark.parametrize( + ("attrs", "error_match"), + [ + ( + RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 0.0, "inertia": [1.0, 2.0, 3.0]}} + ), + "density is required when mass is zero", + ), + ( + RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 1.0, "inertia": [1.0, 2.0]}} + ), + "inertia must contain", + ), + ( + RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 1.0, "com_quaternion": [0.0, 0.0, 0.0, 0.0]}} + ), + "com_quaternion cannot be zero", + ), + ], + ids=["inertia-without-mass", "invalid-inertia-shape", "zero-com-quaternion"], +) +def test_rigid_descriptor_rejects_invalid_mass_properties( + attrs: RigidBodyPhysicsCfg, + error_match: str, +) -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=attrs, + ) + + with pytest.raises(ValueError, match=error_match): + rigid_desc_from_cfg(cfg) + + +def test_static_rigid_descriptor_omits_mass_properties() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + body_type="static", + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": { + "mass": 2.0, + "density": 3.0, + "inertia": [1.0, 2.0, 3.0], + "com_position": [0.1, 0.2, 0.3], + } + } + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert descriptor.physics.mass is None + assert descriptor.physics.density is None + assert descriptor.physics.inertia is None + assert descriptor.physics.com_position is None + + +def test_kinematic_rigid_descriptor_honors_mass_priority() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + body_type="kinematic", + attrs=RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 2.0, "density": 3.0}} + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert descriptor.physics.mass == 2.0 + assert descriptor.physics.density is None + + +def test_grouped_rigid_physics_routes_common_and_backend_properties() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=2.0), + rigid_props=DefaultRigidBodyPropertiesCfg(linear_damping=0.2), + collision_props=NewtonCollisionPropertiesCfg( + collision_enabled=False, + margin=0.01, + ), + material_props=NewtonRigidBodyMaterialCfg( + dynamic_friction=0.4, + ke=1000.0, + torsional_friction=0.02, + ), + ), + ) + + descriptor, _ = rigid_desc_from_cfg( + cfg, + newton_solver_type="mujoco_warp", + ) + + assert descriptor.physics.mass == 2.0 + assert descriptor.physics.dexsim.linear_damping == 0.2 + assert descriptor.physics.dexsim.angular_damping is None + collision = descriptor.collisions[0] + assert collision.enable_collision is False + assert collision.dexsim.dynamic_friction == 0.4 + assert collision.dexsim.static_friction is None + assert collision.newton.margin == 0.01 + assert collision.newton.mu == 0.4 + assert collision.newton.ke == 1000.0 + assert collision.newton.mu_torsional == 0.02 + + +def test_portable_collision_envelope_compiles_to_both_backends() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg( + contact_offset=0.015, + rest_offset=0.005, + ) + ), + ) + + descriptor, _ = rigid_desc_from_cfg( + cfg, + newton_solver_type="mujoco_warp", + ) + + collision = descriptor.collisions[0] + assert collision.dexsim.contact_offset == pytest.approx(0.015) + assert collision.dexsim.rest_offset == pytest.approx(0.005) + assert collision.newton.margin == pytest.approx(0.005) + assert collision.newton.gap == pytest.approx(0.01) + + +def test_newton_native_collision_envelope_overrides_portable_translation() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg( + contact_offset=0.015, + rest_offset=0.005, + margin=0.007, + gap=0.004, + ) + ), + ) + + descriptor, _ = rigid_desc_from_cfg( + cfg, + newton_solver_type="mujoco_warp", + ) + + collision = descriptor.collisions[0] + assert collision.dexsim.contact_offset == pytest.approx(0.015) + assert collision.dexsim.rest_offset == pytest.approx(0.005) + assert collision.newton.margin == pytest.approx(0.007) + assert collision.newton.gap == pytest.approx(0.004) + + +def test_portable_collision_envelope_rejects_invalid_ordering() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg( + contact_offset=0.001, + rest_offset=0.002, + ) + ), + ) + + with pytest.raises(ValueError, match="no smaller than rest_offset"): + rigid_desc_from_cfg(cfg, newton_solver_type="mujoco_warp") + + +def test_newton_rejects_ambiguous_portable_contact_offset() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg(contact_offset=0.001) + ), + ) + + with pytest.raises(ValueError, match="requires rest_offset"): + rigid_desc_from_cfg(cfg, newton_solver_type="mujoco_warp") + + +def test_grouped_rigid_physics_keeps_unset_backend_blocks_absent() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + collision_props=CollisionPropertiesCfg(collision_enabled=True) + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert descriptor.physics.dexsim is None + assert descriptor.physics.newton is None + assert descriptor.collisions[0].dexsim is None + assert descriptor.collisions[0].newton is None + + +def test_grouped_rigid_physics_overlays_usd_without_erasing_source( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = ObjectDesc( + name="source", + physics=RigidBodyPhysicsDesc.dynamic( + mass=7.0, + inertia=np.array([1.0, 2.0, 3.0], dtype=np.float32), + dexsim=DexsimPhysicsDesc( + linear_damping=0.6, + angular_damping=0.8, + ), + ), + collisions=[ + CollisionDesc( + enable_collision=False, + dexsim=DexsimCollisionDesc( + dynamic_friction=0.9, + contact_offset=0.05, + ), + newton=NewtonCollisionDesc(margin=0.03, gap=0.07), + ) + ], + ) + scene = SimpleNamespace(materials={}) + + def parse_singleton(path, collection, label): + return scene, source + + monkeypatch.setattr( + "embodichain.lab.sim.spawn.usd._parse_singleton", + parse_singleton, + ) + cfg = RigidObjectCfg( + uid="cube", + shape=MeshCfg(fpath="cube.usd"), + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + rigid_props=DefaultRigidBodyPropertiesCfg(linear_damping=0.2), + collision_props=NewtonCollisionPropertiesCfg(margin=0.01), + material_props=RigidBodyMaterialCfg(dynamic_friction=0.4), + ), + ) + + descriptor, _ = rigid_desc_from_usd(cfg) + + assert descriptor.physics.mass == 7.0 + np.testing.assert_array_equal(descriptor.physics.inertia, [1.0, 2.0, 3.0]) + assert descriptor.physics.dexsim.linear_damping == 0.2 + assert descriptor.physics.dexsim.angular_damping == 0.8 + collision = descriptor.collisions[0] + assert collision.enable_collision is False + assert collision.dexsim.dynamic_friction == 0.4 + assert collision.dexsim.contact_offset == 0.05 + assert collision.newton.margin == 0.01 + assert collision.newton.gap == 0.07 + + +def test_rigid_usd_can_recompute_source_inertia( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = ObjectDesc( + name="source", + physics=RigidBodyPhysicsDesc.dynamic( + mass=7.0, + inertia=np.array([1.0, 2.0, 3.0], dtype=np.float32), + ), + collisions=[CollisionDesc()], + ) + monkeypatch.setattr( + "embodichain.lab.sim.spawn.usd._parse_singleton", + lambda path, collection, label: (SimpleNamespace(materials={}), source), + ) + cfg = RigidObjectCfg( + uid="cube", + shape=MeshCfg(fpath="cube.usd"), + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg( + mass=2.0, + recompute_inertia=True, + ) + ), + ) + + descriptor, _ = rigid_desc_from_usd(cfg) + + assert descriptor.physics.mass == 2.0 + assert descriptor.physics.inertia is None + + +def test_rigid_usd_preserves_asset_physics_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + source_mass = 7.0 + source_scale = np.array([2.0, 3.0, 4.0], dtype=np.float32) + source = ObjectDesc( + name="source", + physics=RigidBodyPhysicsDesc.dynamic(mass=source_mass), + collisions=[CollisionDesc(enable_collision=False)], + body_scale=source_scale, + ) + monkeypatch.setattr( + "embodichain.lab.sim.spawn.usd._parse_singleton", + lambda path, collection, label: (SimpleNamespace(materials={}), source), + ) + cfg = RigidObjectCfg( + uid="cube", + shape=MeshCfg(fpath="cube.usd"), + body_type="static", + body_scale=(1.0, 1.0, 1.0), + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=1.0)), + ) + + descriptor, _ = rigid_desc_from_usd(cfg) + + assert descriptor.physics.mass == source_mass + assert descriptor.physics.actor_type == dexsim.types.ActorType.DYNAMIC + np.testing.assert_array_equal(descriptor.body_scale, source_scale) + assert descriptor.collisions[0].enable_collision is False + assert cfg.body_type == "dynamic" + assert cfg.body_scale == tuple(source_scale) + + +def test_rigid_descriptor_forwards_newton_sdf_options() -> None: + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.glb", + collision=MeshCollisionCfg( + approximation="sdf", + sdf_padding=0.02, + ), + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert descriptor.collisions[0].newton.force_sdf is True + assert descriptor.collisions[0].newton.sdf_padding == pytest.approx(0.02) + + +def test_mesh_collision_and_backend_property_slots_compile_independently() -> None: + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.glb", + collision=MeshCollisionCfg( + approximation="sdf", + sdf_target_voxel_size=0.005, + sdf_padding=0.02, + ), + ), + attrs=RigidBodyPhysicsCfg( + rigid_props=DefaultRigidBodyPropertiesCfg(linear_damping=0.2), + collision_props=NewtonCollisionPropertiesCfg(margin=0.04), + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + collision = descriptor.collisions[0] + + assert descriptor.physics.dexsim.linear_damping == pytest.approx(0.2) + assert collision.approximation == CollisionApproximation.SDF + assert collision.decomp_max_hulls == 1 + assert collision.newton.margin == pytest.approx(0.04) + assert collision.newton.sdf_target_voxel_size == pytest.approx(0.005) + assert collision.newton.sdf_max_resolution is None + assert collision.newton.sdf_padding == pytest.approx(0.02) + + +def test_mesh_cfg_legacy_collision_fields_normalize_before_compilation() -> None: + with pytest.warns(DeprecationWarning): + cfg = RigidObjectCfg.from_dict( + { + "uid": "mesh", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh.glb", + "max_convex_hull_num": 3, + "acd_method": "coacd", + }, + } + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert ( + descriptor.collisions[0].approximation + == CollisionApproximation.CONVEX_DECOMPOSITION + ) + assert descriptor.collisions[0].decomp_max_hulls == 3 + + +def test_static_triangle_mesh_collision_compiles_without_convex_cooking() -> None: + cfg = RigidObjectCfg( + uid="mesh", + body_type="static", + shape=MeshCfg( + fpath="mesh.glb", + collision=MeshCollisionCfg(approximation="triangle_mesh"), + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + assert descriptor.collisions[0].approximation == CollisionApproximation.NONE + + +def test_dynamic_triangle_mesh_collision_is_rejected_before_spawn() -> None: + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.glb", + collision=MeshCollisionCfg(approximation="triangle_mesh"), + ), + ) + + with pytest.raises(ValueError, match="only for static"): + rigid_desc_from_cfg(cfg) + + +def test_spawn_rejects_unsupported_convex_decomposition_method() -> None: + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.glb", + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=4, + acd_method="vhacd", + ), + ), + ) + + with pytest.raises(ValueError, match="only acd_method='coacd'"): + rigid_desc_from_cfg(cfg) + + +def test_default_collision_solver_fields_compile_from_collision_slot() -> None: + cfg = RigidObjectCfg( + uid="cube", + shape=CubeCfg(size=(0.1, 0.1, 0.1)), + attrs=RigidBodyPhysicsCfg( + collision_props=DefaultCollisionPropertiesCfg( + contact_offset=0.01, + torsional_patch_radius=0.02, + min_torsional_patch_radius=0.005, + disable_strong_friction=True, + ) + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + default_collision = descriptor.collisions[0].dexsim + assert default_collision.contact_offset == pytest.approx(0.01) + assert default_collision.torsional_patch_radius == pytest.approx(0.02) + assert default_collision.min_torsional_patch_radius == pytest.approx(0.005) + assert default_collision.disable_strong_friction is True + + +def test_mesh_descriptor_passes_load_options_to_spawn() -> None: + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.glb", + load_option=LoadOption( + rebuild_normals=True, + rebuild_tangent=True, + rebuild_3rdnormal=False, + rebuild_3rdtangent=False, + smooth=45.0, + ), + ), + ) + + descriptor, _ = rigid_desc_from_cfg(cfg) + + option = descriptor.renders[0].load_option + assert option is not None + assert option.rebuild_normals is True + assert option.rebuild_tangent is True + assert option.rebuild_3rdnormal is False + assert option.rebuild_3rdtangent is False + assert option.smooth == 45.0 + + +def test_articulation_constructor_defers_newton_properties_until_configure() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(restitution=RESTITUTION) + ), + ) + + descriptor = articulation_desc_from_cfg( + cfg, + newton_solver_type="mujoco_warp", + ) + + assert descriptor.newton_collision is None + assert descriptor.newton_drive is None + assert descriptor.urdf_read_inertia is True + + descriptor.links = _resolved_articulation_desc().links + descriptor.joints = _resolved_articulation_desc().joints + configure_articulation_desc( + descriptor, + cfg, + newton_solver_type="mujoco_warp", + ) + assert descriptor.links[0].collisions[0].newton is None + + +def test_flat_articulation_physics_is_rejected_at_config_boundary() -> None: + with pytest.raises(ValueError, match="Removed flat rigid-body attrs fields"): + ArticulationCfg.from_dict( + { + "uid": "robot", + "fpath": "robot.urdf", + "attrs": {"mass": 2.0}, + } + ) + + +def test_articulation_root_properties_compile_to_common_descriptor() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + root_props=ArticulationRootPropertiesCfg( + fixed_base=False, + self_collision_enabled=True, + ), + ) + + descriptor = articulation_desc_from_cfg(cfg) + + assert descriptor.fixed_base is False + assert descriptor.urdf_fix_root_link is False + assert descriptor.enable_self_collision is True + + +def test_articulation_root_defaults_are_resolved_at_import_boundary() -> None: + descriptor = articulation_desc_from_cfg( + ArticulationCfg(uid="robot", fpath="robot.urdf") + ) + + assert descriptor.fixed_base is True + assert descriptor.urdf_fix_root_link is True + assert descriptor.enable_self_collision is False + + +def test_explicit_root_properties_override_usd_in_preserve_mode() -> None: + source = ArticulationDesc( + name="source", + fixed_base=False, + enable_self_collision=True, + ) + cfg = ArticulationCfg( + uid="robot", + fpath="robot.usd", + asset_physics_mode="preserve", + root_props=ArticulationRootPropertiesCfg( + fixed_base=True, + self_collision_enabled=False, + ), + ) + + with patch( + "embodichain.lab.sim.spawn.usd._parse_singleton", + return_value=(SimpleNamespace(materials={}), source), + ): + descriptor, _ = articulation_desc_from_usd(cfg) + + assert descriptor.fixed_base is True + assert descriptor.enable_self_collision is False + + +def test_default_root_properties_override_usd_values() -> None: + source = ArticulationDesc( + name="source", + fixed_base=False, + enable_self_collision=True, + ) + cfg = ArticulationCfg( + uid="robot", + fpath="robot.usd", + asset_physics_mode="preserve", + ) + + with patch( + "embodichain.lab.sim.spawn.usd._parse_singleton", + return_value=(SimpleNamespace(materials={}), source), + ): + descriptor, _ = articulation_desc_from_usd(cfg) + + assert descriptor.fixed_base is True + assert descriptor.enable_self_collision is False + + +def test_explicit_none_root_properties_preserve_usd_values() -> None: + source = ArticulationDesc( + name="source", + fixed_base=False, + enable_self_collision=True, + ) + cfg = ArticulationCfg( + uid="robot", + fpath="robot.usd", + asset_physics_mode="preserve", + root_props=ArticulationRootPropertiesCfg( + fixed_base=None, + self_collision_enabled=None, + ), + ) + + with patch( + "embodichain.lab.sim.spawn.usd._parse_singleton", + return_value=(SimpleNamespace(materials={}), source), + ): + descriptor, _ = articulation_desc_from_usd(cfg) + + assert descriptor.fixed_base is False + assert descriptor.enable_self_collision is True + + +def test_articulation_descriptor_rejects_newton_acceleration_drive() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + drive_type="acceleration", + ), + ) + + descriptor = articulation_desc_from_cfg(cfg, newton_solver_type="mujoco_warp") + descriptor.links = _resolved_articulation_desc().links + descriptor.joints = _resolved_articulation_desc().joints + + with pytest.raises(NotImplementedError, match="acceleration-drive"): + configure_articulation_desc( + descriptor, + cfg, + newton_solver_type="mujoco_warp", + ) + + +@pytest.mark.parametrize( + ( + "target_mode", + "expected_default_mode", + "expected_newton_mode", + "expected_stiffness", + "expected_damping", + ), + [ + ("none", DriveType.NONE, 0, 0.0, 0.0), + ("position", DriveType.FORCE, 1, 12.0, 4.0), + ("velocity", DriveType.FORCE, 2, 0.0, 4.0), + ("position_velocity", DriveType.FORCE, 3, 12.0, 4.0), + ("effort", DriveType.NONE, 4, 0.0, 0.0), + ], +) +def test_portable_joint_target_modes_compile_for_both_backends( + target_mode: str, + expected_default_mode: DriveType, + expected_newton_mode: int, + expected_stiffness: float, + expected_damping: float, +) -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + target_mode=target_mode, # type: ignore[arg-type] + stiffness=12.0, + damping=4.0, + ), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc( + descriptor, + cfg, + newton_solver_type="mujoco_warp", + ) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.drive_mode == expected_default_mode + assert joint.newton.target_mode == expected_newton_mode + assert joint.dexsim.stiffness == pytest.approx(expected_stiffness) + assert joint.dexsim.damping == pytest.approx(expected_damping) + assert joint.newton.target_ke == pytest.approx(expected_stiffness) + assert joint.newton.target_kd == pytest.approx(expected_damping) + + +def test_force_drive_defaults_newton_target_to_position_velocity() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(drive_type="force"), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc( + descriptor, + cfg, + newton_solver_type="mujoco_warp", + ) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.drive_mode == DriveType.FORCE + assert joint.newton.target_mode == 3 + + +@pytest.mark.parametrize( + ("target_mode", "expected_ke", "expected_kd"), + [ + ("none", 0.0, 0.0), + ("velocity", 0.0, 4.0), + ("effort", 0.0, 0.0), + ], +) +def test_non_mode_aware_newton_solver_uses_gain_fallbacks( + target_mode: str, + expected_ke: float, + expected_kd: float, +) -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + target_mode=target_mode, # type: ignore[arg-type] + stiffness=12.0, + damping=4.0, + ), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg, newton_solver_type="xpbd") + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.newton.target_ke == pytest.approx(expected_ke) + assert joint.newton.target_kd == pytest.approx(expected_kd) + + +def test_non_mode_aware_newton_position_fallback_is_explicit() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + target_mode="position", + stiffness=12.0, + damping=4.0, + ), + ) + descriptor = _resolved_articulation_desc() + + with pytest.warns(UserWarning, match="POSITION is emulated"): + configure_articulation_desc(descriptor, cfg, newton_solver_type="xpbd") + + +def test_auto_solver_defers_position_mode_compatibility_warning() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + target_mode="position", + stiffness=12.0, + damping=4.0, + ), + ) + descriptor = _resolved_articulation_desc() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + configure_articulation_desc(descriptor, cfg, newton_solver_type="auto") + + assert not caught + + +def test_default_articulation_body_properties_compile_per_link() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + rigid_props=DefaultRigidBodyPropertiesCfg( + sleep_threshold=0.002, + min_position_iters=8, + min_velocity_iters=2, + ) + ), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc( + descriptor, + cfg, + newton_solver_type="mujoco_warp", + ) + + for link in descriptor.links: + assert link.rigid_body.dexsim.sleep_threshold == pytest.approx(0.002) + assert link.rigid_body.dexsim.min_position_iters == 8 + assert link.rigid_body.dexsim.min_velocity_iters == 2 + assert link.rigid_body.newton is None + + +def test_articulation_config_applies_to_exact_source_resolved_names() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg(dynamic_friction=0.4), + ), + link_attrs={ + "fingers": LinkPhysicsOverrideCfg( + link_names_expr=["finger_.*"], + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg( + mass=2.0, + recompute_inertia=True, + ), + material_props=RigidBodyMaterialCfg(dynamic_friction=0.8), + ), + ) + }, + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + stiffness={"arm_.*": 10.0}, + damping=3.0, + max_effort=20.0, + max_velocity=4.0, + friction=0.1, + armature=0.2, + ), + qpos_limits={"arm_.*": [-1.0, 1.0]}, + ) + + descriptor = articulation_desc_from_cfg(cfg) + assert descriptor.links == [] + assert descriptor.joints == [] + + resolved = _resolved_articulation_desc() + descriptor.links = resolved.links + descriptor.joints = resolved.joints + descriptor.root_link_name = resolved.root_link_name + + with ( + patch.object( + descriptor, + "set_link_properties", + wraps=descriptor.set_link_properties, + ) as set_link_properties, + patch.object( + descriptor, + "set_joint_properties", + wraps=descriptor.set_joint_properties, + ) as set_joint_properties, + ): + configure_articulation_desc(descriptor, cfg) + + assert set_link_properties.call_count == len(descriptor.links) + assert set_joint_properties.call_count == len(descriptor.joints) + + base = descriptor.get_link_desc("base") + finger = descriptor.get_link_desc("finger_left") + assert base.rigid_body.mass == 1.0 + assert base.collisions[0].dexsim.dynamic_friction == 0.4 + np.testing.assert_array_equal( + base.rigid_body.inertia, + np.ones(3, dtype=np.float32), + ) + assert finger.rigid_body.mass == 2.0 + assert finger.collisions[0].newton.mu == 0.8 + assert finger.rigid_body.inertia is None + assert finger.replace_inertial + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.damping == 3.0 + assert joint.newton.target_kd == 3.0 + assert joint.armature == 0.2 + assert joint.dexsim.stiffness == 10.0 + assert joint.newton.target_ke == 10.0 + assert joint.effort_limit == 20.0 + assert joint.velocity_limit == 4.0 + assert joint.lower_limit == -1.0 + assert joint.upper_limit == 1.0 + + +def test_joint_drive_properties_compile_joint_dynamics() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg( + stiffness=10.0, + max_effort=20.0, + max_velocity=2.0, + friction=0.4, + armature=0.7, + ), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.stiffness == pytest.approx(10.0) + assert joint.effort_limit == pytest.approx(20.0) + assert joint.velocity_limit == pytest.approx(2.0) + assert joint.dexsim.joint_friction == pytest.approx(0.4) + assert joint.newton.friction == pytest.approx(0.4) + assert joint.armature == pytest.approx(0.7) + + +def test_articulation_array_qpos_limits_compile_before_backend_build() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + qpos_limits=np.array([[-0.5, 0.75]], dtype=np.float32), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.lower_limit == pytest.approx(-0.5) + assert joint.upper_limit == pytest.approx(0.75) + + +def test_robot_control_part_drive_rule_expands_before_spawn() -> None: + cfg = RobotCfg( + uid="robot", + fpath="robot.urdf", + control_parts={"arm": ["arm_joint"]}, + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + stiffness={"arm": 10.0, "arm_joint": 20.0}, + ), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.stiffness == 20.0 + assert joint.newton.target_ke == 20.0 + + +def test_newton_joint_compatibility_subclass_uses_portable_target_mode() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=NewtonJointDrivePropertiesCfg( + drive_type="force", + stiffness={"arm_.*": 12.0}, + damping=4.0, + friction=0.5, + armature=0.7, + target_mode={"arm_.*": "velocity"}, + ), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.stiffness == 0.0 + assert joint.dexsim.damping == 4.0 + assert joint.dexsim.joint_friction == 0.5 + assert joint.armature == 0.7 + assert joint.newton.target_ke == 0.0 + assert joint.newton.target_kd == 4.0 + assert joint.newton.friction == 0.5 + assert joint.newton.armature is None + assert joint.newton.target_mode == 2 + + +def test_grouped_link_physics_overrides_compose_after_source_resolution() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0), + material_props=RigidBodyMaterialCfg(dynamic_friction=0.4), + ), + link_attrs={ + "fingers": LinkPhysicsOverrideCfg( + link_names_expr=["finger_.*"], + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg( + mass=2.0, + recompute_inertia=True, + ), + material_props=RigidBodyMaterialCfg(dynamic_friction=0.8), + ), + ) + }, + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + base = descriptor.get_link_desc("base") + finger = descriptor.get_link_desc("finger_left") + assert base.rigid_body.mass == 1.0 + assert base.collisions[0].newton.mu == 0.4 + assert finger.rigid_body.mass == 2.0 + assert finger.collisions[0].newton.mu == 0.8 + assert finger.rigid_body.inertia is None + + +def test_global_articulation_mass_properties_can_recompute_inertia() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(recompute_inertia=True)), + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + for link in descriptor.links: + assert link.rigid_body.inertia is None + assert link.replace_inertial is True + + +def test_per_link_mass_properties_can_preserve_global_source_inertia() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(recompute_inertia=True)), + link_attrs={ + "fingers": LinkPhysicsOverrideCfg( + link_names_expr=["finger_.*"], + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(recompute_inertia=False) + ), + ) + }, + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + base = descriptor.get_link_desc("base") + finger = descriptor.get_link_desc("finger_left") + assert base.rigid_body.inertia is None + assert base.replace_inertial is True + np.testing.assert_array_equal( + finger.rigid_body.inertia, + np.ones(3, dtype=np.float32), + ) + assert finger.replace_inertial is False + + +def test_explicit_and_recomputed_inertia_are_mutually_exclusive() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg( + mass=2.0, + inertia=[1.0, 2.0, 3.0], + recompute_inertia=True, + ) + ), + ) + + with pytest.raises(ValueError, match="recompute_inertia"): + configure_articulation_desc(_resolved_articulation_desc(), cfg) + + +def test_recompute_inertia_rejects_non_boolean_values() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg( + recompute_inertia="yes", # type: ignore[arg-type] + ) + ), + ) + + with pytest.raises(TypeError, match="recompute_inertia"): + configure_articulation_desc(_resolved_articulation_desc(), cfg) + + +def test_grouped_link_zero_mass_falls_back_to_inherited_density() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=1.0, density=500.0) + ), + link_attrs={ + "fingers": LinkPhysicsOverrideCfg( + link_names_expr=["finger_.*"], + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=0.0)), + ) + }, + ) + descriptor = _resolved_articulation_desc() + + configure_articulation_desc(descriptor, cfg) + + assert descriptor.get_link_desc("base").rigid_body.mass == 1.0 + finger_physics = descriptor.get_link_desc("finger_left").rigid_body + assert finger_physics.mass is None + assert finger_physics.density == 500.0 + + +@pytest.mark.parametrize("source_path", ["robot.urdf", "robot.usd"]) +def test_articulation_preserve_mode_keeps_source_physics(source_path: str) -> None: + descriptor = _resolved_articulation_desc() + source_joint = descriptor.get_joint_desc("arm_joint") + source_joint.lower_limit = -2.0 + source_joint.upper_limit = 2.0 + source_joint.effort_limit = 321.0 + source_joint.dexsim = DexsimJointDesc(stiffness=123.0, damping=456.0) + before = copy.deepcopy(descriptor) + cfg = ArticulationCfg( + uid="robot", + fpath=source_path, + asset_physics_mode="preserve", + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=9.0)), + joint_drive_props=JointDrivePropertiesCfg( + drive_type="force", + stiffness=10.0, + damping=20.0, + ), + qpos_limits={"arm_.*": [-1.0, 1.0]}, + ) + + with pytest.warns( + UserWarning, + match="preserve.*attrs, joint_drive_props, qpos_limits", + ): + configure_articulation_desc(descriptor, cfg) + + _assert_property_tree_equal(descriptor, before) + + +def test_articulation_drive_overlay_preserves_unspecified_source_fields() -> None: + source_stiffness = 123.0 + source_damping = 456.0 + configured_stiffness = 10.0 + descriptor = _resolved_articulation_desc() + joint = descriptor.get_joint_desc("arm_joint") + joint.effort_limit = 321.0 + joint.dexsim = DexsimJointDesc( + stiffness=source_stiffness, + damping=source_damping, + drive_mode=DriveType.FORCE, + ) + joint.newton = NewtonJointDesc( + target_ke=source_stiffness, + target_kd=source_damping, + target_mode=2, + ) + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + joint_drive_props=JointDrivePropertiesCfg(stiffness=configured_stiffness), + ) + + configure_articulation_desc(descriptor, cfg) + + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.stiffness == configured_stiffness + assert joint.dexsim.damping == source_damping + assert joint.dexsim.drive_mode == DriveType.FORCE + assert joint.newton.target_ke == configured_stiffness + assert joint.newton.target_kd == source_damping + assert joint.newton.target_mode == 2 + assert joint.effort_limit == 321.0 + + +def test_articulation_overlay_does_not_invent_collision_geometry() -> None: + descriptor = _resolved_articulation_desc() + collisionless_link = LinkDesc( + "imu_link", + "base", + np.eye(4, dtype=np.float32), + rigid_body=RigidBodyPhysicsDesc.dynamic(mass=0.1), + ) + descriptor.links.append(collisionless_link) + cfg = ArticulationCfg( + uid="robot", + fpath="robot.urdf", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(dynamic_friction=0.5) + ), + ) + + configure_articulation_desc(descriptor, cfg) + + assert descriptor.get_link_desc("imu_link").collisions == [] + + +@pytest.mark.parametrize( + ("cfg", "error_type"), + [ + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + link_attrs={ + "missing": LinkPhysicsOverrideCfg( + link_names_expr=["missing_.*"], + ) + }, + ), + ValueError, + ), + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + link_attrs={ + "first": LinkPhysicsOverrideCfg( + link_names_expr=["finger_.*"], + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=2.0) + ), + ), + "second": LinkPhysicsOverrideCfg( + link_names_expr=["finger_left"], + attrs=RigidBodyPhysicsCfg( + mass_props=MassPropertiesCfg(mass=3.0) + ), + ), + }, + ), + ValueError, + ), + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + joint_drive_props=JointDrivePropertiesCfg( + stiffness={"missing_.*": 10.0} + ), + ), + ValueError, + ), + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + joint_drive_props=JointDrivePropertiesCfg( + stiffness={"arm_.*": "not-a-number"} + ), + ), + TypeError, + ), + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + qpos_limits={"arm_.*": [1.0, -1.0]}, + ), + ValueError, + ), + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + qpos_limits=np.zeros((2, 2), dtype=np.float32), + ), + ValueError, + ), + ( + ArticulationCfg( + uid="robot", + fpath="robot.urdf", + joint_drive_props=NewtonJointDrivePropertiesCfg( + target_mode={"arm_.*": "servo"} + ), + ), + ValueError, + ), + ], + ids=[ + "unmatched-link", + "overlapping-link-groups", + "unmatched-joint", + "non-numeric-joint-property", + "invalid-qpos-limit", + "invalid-array-qpos-shape", + "invalid-newton-target-mode", + ], +) +def test_articulation_config_validation_failure_is_atomic( + cfg: ArticulationCfg, + error_type: type[Exception], +) -> None: + cfg.asset_physics_mode = "overlay" + descriptor = _resolved_articulation_desc() + before = copy.deepcopy(descriptor) + + with pytest.raises(error_type): + configure_articulation_desc(descriptor, cfg) + + _assert_property_tree_equal(descriptor, before) + finger = descriptor.get_link_desc("finger_left") + np.testing.assert_array_equal( + finger.rigid_body.inertia, + np.ones(3, dtype=np.float32), + ) + + +def test_usd_articulation_uses_the_same_exact_name_configuration() -> None: + cfg = ArticulationCfg( + uid="robot", + fpath="robot.usd", + asset_physics_mode="overlay", + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=1.0)), + link_attrs={ + "fingers": LinkPhysicsOverrideCfg( + link_names_expr=["finger_.*"], + attrs=RigidBodyPhysicsCfg(mass_props=MassPropertiesCfg(mass=2.0)), + ) + }, + joint_drive_props=JointDrivePropertiesCfg(stiffness={"arm_.*": 10.0}), + ) + source = ArticulationDesc( + name="source", + links=[ + LinkDesc( + "finger_left", + "", + np.eye(4, dtype=np.float32), + collisions=[CollisionDesc()], + rigid_body=RigidBodyPhysicsDesc.dynamic(mass=0.5), + ) + ], + joints=[ + JointDesc( + "arm_joint", + "finger_left", + "tip", + dexsim.engine.JointType.REVOLUTE, + ) + ], + ) + + with patch( + "embodichain.lab.sim.spawn.usd._parse_singleton", + return_value=(SimpleNamespace(materials={}), source), + ): + descriptor, _ = articulation_desc_from_usd( + cfg, + newton_solver_type="mujoco_warp", + ) + + configure_articulation_desc( + descriptor, + cfg, + newton_solver_type="mujoco_warp", + ) + + assert descriptor.get_link_desc("finger_left").rigid_body.mass == 2.0 + joint = descriptor.get_joint_desc("arm_joint") + assert joint.dexsim.stiffness == 10.0 + assert joint.newton.target_ke == 10.0 + + +def test_spawn_post_config_only_applies_render_uv() -> None: + render_body = Mock() + entity = Mock() + entity.joint_dof_layout = [] + entity.get_render_body.return_value = render_body + articulation = object.__new__(Articulation) + articulation.cfg = SimpleNamespace(compute_uv=True) + articulation._entities = [entity] + articulation.__dict__["link_names"] = ["base"] + articulation._prepared_default_root_topology_revision = -1 + articulation._mimic_info = SimpleNamespace( + mimic_id=np.array([], dtype=np.int32), + mimic_parent=np.array([], dtype=np.int32), + ) + articulation._set_default_joint_drive = Mock() + + articulation._apply_spawn_config() + + articulation._set_default_joint_drive.assert_not_called() + entity.get_render_body.assert_called_once_with("base") + render_body.set_projective_uv.assert_called_once_with() + + +def test_spawn_post_config_applies_default_only_root_properties() -> None: + native_articulation = Mock() + entity = SimpleNamespace( + _physics_binding=native_articulation, + joint_dof_layout=[], + ) + articulation = object.__new__(Articulation) + articulation.cfg = ArticulationCfg( + root_props=ArticulationRootPropertiesCfg( + sleep_threshold=0.005, + min_position_iters=8, + min_velocity_iters=2, + ) + ) + articulation._spawn_result = SimpleNamespace(backend="dexsim", topology_revision=0) + articulation._entities = [entity] + articulation._prepared_default_root_topology_revision = -1 + articulation._mimic_info = SimpleNamespace( + mimic_id=np.array([], dtype=np.int32), + mimic_parent=np.array([], dtype=np.int32), + ) + + articulation._apply_spawn_config() + + native_articulation.set_sleep_threshold.assert_called_once_with(0.005) + native_articulation.set_solver_iteration_counts.assert_called_once_with( + min_position_iters=8, + min_velocity_iters=2, + ) + + +def test_newton_skips_default_only_articulation_root_properties() -> None: + native_articulation = Mock() + entity = SimpleNamespace( + _physics_binding=native_articulation, + joint_dof_layout=[], + ) + articulation = object.__new__(Articulation) + articulation.cfg = ArticulationCfg( + root_props=ArticulationRootPropertiesCfg( + sleep_threshold=0.005, + min_position_iters=8, + min_velocity_iters=2, + ) + ) + articulation._spawn_result = SimpleNamespace(backend="newton", topology_revision=0) + articulation._entities = [entity] + articulation._prepared_default_root_topology_revision = -1 + articulation._mimic_info = SimpleNamespace( + mimic_id=np.array([], dtype=np.int32), + mimic_parent=np.array([], dtype=np.int32), + ) + + articulation._apply_spawn_config() + + native_articulation.set_sleep_threshold.assert_not_called() + native_articulation.set_solver_iteration_counts.assert_not_called() diff --git a/tests/sim/spawn/test_scene.py b/tests/sim/spawn/test_scene.py new file mode 100644 index 000000000..d1326fba2 --- /dev/null +++ b/tests/sim/spawn/test_scene.py @@ -0,0 +1,401 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from embodichain.lab.sim.cfg import ArticulationRootPropertiesCfg +from embodichain.lab.sim.objects.articulation import Articulation +from embodichain.lab.sim.spawn.scene import SpawnScene + +pytestmark = pytest.mark.no_sim + + +def _make_scene(handles: dict[str, object]) -> SpawnScene: + scene = object.__new__(SpawnScene) + scene.builder = SimpleNamespace( + is_finalized=True, + result=SimpleNamespace(handles=handles), + ) + scene._assets = {} + return scene + + +class _RetryableFacade: + def __init__(self, *, fail_first: bool = False) -> None: + self._entities: list[object] = [] + self.is_declared = True + self.fail_first = fail_first + self.bind_attempts = 0 + + def attach_spawn_handles(self, entities: tuple[object, ...]) -> None: + self._entities = list(entities) + + def bind_spawn(self, _result: object) -> None: + self.bind_attempts += 1 + if self.fail_first and self.bind_attempts == 1: + raise RuntimeError("bind failed") + self.is_declared = False + + +class _RuntimeConfigFacade(_RetryableFacade): + def __init__(self, events: list[str]) -> None: + super().__init__() + self.events = events + + def attach_spawn_handles(self, entities: tuple[object, ...]) -> None: + self.events.append("attach") + super().attach_spawn_handles(entities) + + def _prepare_spawn_runtime_config(self, _result: object) -> None: + self.events.append("runtime_config") + + +def test_bind_retries_only_incomplete_declarations() -> None: + first_handle = object() + second_handle = object() + scene = _make_scene({"first": first_handle, "second": second_handle}) + first = _RetryableFacade() + second = _RetryableFacade(fail_first=True) + + scene.track( + "rigid_object", + "first", + SimpleNamespace(name="first", per_env=False), + facade=first, + ) + scene.track( + "rigid_object", + "second", + SimpleNamespace(name="second", per_env=False), + facade=second, + ) + + with pytest.raises(RuntimeError, match="bind failed"): + scene.bind() + scene.bind() + scene.bind() + + assert first._entities == [first_handle] + assert second._entities == [second_handle] + assert first.bind_attempts == 1 + assert second.bind_attempts == 2 + + +def test_runtime_config_attaches_articulation_before_preparing_it() -> None: + scene = _make_scene({}) + events: list[str] = [] + facade = _RuntimeConfigFacade(events) + scene.track( + "articulation", + "robot", + SimpleNamespace(name="robot", per_env=False), + facade=facade, + ) + handle = object() + scene.builder.result.handles["robot"] = handle + + scene.prepare_runtime_config(scene.builder.result) + + assert facade._entities == [handle] + assert events == ["attach", "runtime_config"] + + +def test_default_root_properties_prepare_once_per_topology_revision() -> None: + native_articulation = MagicMock() + articulation = object.__new__(Articulation) + articulation.cfg = SimpleNamespace( + root_props=ArticulationRootPropertiesCfg( + min_position_iters=32, + min_velocity_iters=8, + ) + ) + articulation._entities = [SimpleNamespace(_physics_binding=native_articulation)] + articulation._prepared_default_root_topology_revision = -1 + result = SimpleNamespace(backend="dexsim", topology_revision=3) + + articulation._prepare_spawn_runtime_config(result) + articulation._prepare_spawn_runtime_config(result) + + native_articulation.set_solver_iteration_counts.assert_called_once_with( + min_position_iters=32, + min_velocity_iters=8, + ) + + result.topology_revision = 4 + articulation._prepare_spawn_runtime_config(result) + assert native_articulation.set_solver_iteration_counts.call_count == 2 + + +def test_newton_skips_default_root_runtime_properties() -> None: + native_articulation = MagicMock() + articulation = object.__new__(Articulation) + articulation.cfg = SimpleNamespace( + root_props=ArticulationRootPropertiesCfg( + min_position_iters=32, + min_velocity_iters=8, + ) + ) + articulation._entities = [SimpleNamespace(_physics_binding=native_articulation)] + articulation._prepared_default_root_topology_revision = -1 + + articulation._prepare_spawn_runtime_config( + SimpleNamespace(backend="newton", topology_revision=3) + ) + + native_articulation.set_solver_iteration_counts.assert_not_called() + + +def test_commit_resolves_and_configures_before_finalize(monkeypatch) -> None: + events: list[str] = [] + descriptor = SimpleNamespace(name="robot", per_env=True, links=[]) + result = object() + builder = SimpleNamespace( + backend="newton", + is_finalized=False, + result=None, + replicate_plan=SimpleNamespace(env_names=lambda: ["arena_0"]), + add_articulation=lambda value: value, + ) + + def resolve_source(_builder: object, value: object) -> None: + events.append("resolve") + value.links = [SimpleNamespace(name="base")] + + def finalize() -> object: + events.append("finalize") + builder.is_finalized = True + builder.result = result + return result + + builder.finalize = finalize + monkeypatch.setattr( + "embodichain.lab.sim.spawn.source.resolve_articulation_source", + resolve_source, + ) + scene = object.__new__(SpawnScene) + scene.builder = builder + scene._assets = {} + + def configure(value: object) -> None: + assert value.links[0].name == "base" + events.append("configure") + + scene.declare( + "articulation", + "robot", + descriptor, + configure_source=configure, + ) + + assert scene.commit() is result + assert events == ["resolve", "configure", "finalize"] + + +@pytest.mark.parametrize("is_finalized", [False, True]) +def test_materialized_articulation_is_configured_before_backend_add( + is_finalized: bool, +) -> None: + events: list[str] = [] + descriptor = SimpleNamespace(name="robot", per_env=True, links=[]) + result = SimpleNamespace(handles={}) + builder = SimpleNamespace( + is_finalized=is_finalized, + result=result, + replicate_plan=SimpleNamespace(env_names=lambda: ["arena_0"]), + ) + + def resolve_source(value: object) -> None: + events.append("resolve") + value.links = [SimpleNamespace(name="base")] + + def configure(value: object) -> None: + assert value.links[0].name == "base" + events.append("configure") + + def add_articulation(value: object) -> object: + assert value.links[0].name == "base" + events.append("add") + return value + + builder.resolve_articulation_source = resolve_source + builder.add_articulation = add_articulation + scene = object.__new__(SpawnScene) + scene.builder = builder + scene._assets = {} + + scene.declare( + "articulation", + "robot", + descriptor, + configure_source=configure, + ) + + assert events == ["resolve", "configure", "add"] + + +@pytest.mark.parametrize("is_finalized", [False, True]) +def test_default_eager_articulation_is_configured_after_native_add( + is_finalized: bool, +) -> None: + events: list[str] = [] + descriptor = SimpleNamespace(name="robot", per_env=True, links=[]) + result = SimpleNamespace(backend="dexsim", handles={}) + builder = SimpleNamespace( + backend="dexsim", + is_finalized=is_finalized, + result=result, + replicate_plan=SimpleNamespace(env_names=lambda: ["arena_0"]), + ) + + def configure(value: object) -> None: + assert value.links[0].name == "base" + events.append("configure") + + def add_articulation(value: object) -> object: + events.append("add") + value.links = [SimpleNamespace(name="base")] + result.handles["arena_0/robot"] = SimpleNamespace( + articulation_desc=value, + apply_dexsim_properties=lambda source: events.append("apply"), + ) + return value + + builder.add_articulation = add_articulation + scene = object.__new__(SpawnScene) + scene.builder = builder + scene._assets = {} + + scene.declare( + "articulation", + "robot", + descriptor, + configure_source=configure, + ) + + assert events == ["add", "configure", "apply"] + + +def test_source_configuration_retries_failure_then_runs_only_once() -> None: + events: list[str] = [] + descriptor = SimpleNamespace(name="robot", per_env=True, links=[]) + builder = SimpleNamespace( + is_finalized=False, + result=None, + replicate_plan=SimpleNamespace(env_names=lambda: ["arena_0"]), + add_articulation=lambda value: value, + ) + + def resolve_sources() -> None: + events.append("resolve") + descriptor.links = [SimpleNamespace(name="base")] + + attempts = 0 + + def configure(_value: object) -> None: + nonlocal attempts + attempts += 1 + events.append("configure") + if attempts == 1: + raise RuntimeError("configuration failed") + + builder.resolve_sources = resolve_sources + scene = object.__new__(SpawnScene) + scene.builder = builder + scene._assets = {} + + scene.declare( + "articulation", + "robot", + descriptor, + configure_source=configure, + ) + + with pytest.raises(RuntimeError, match="configuration failed"): + scene.resolve_sources() + scene.resolve_sources() + scene.resolve_sources() + + assert attempts == 2 + assert events == [ + "resolve", + "configure", + "resolve", + "configure", + "resolve", + ] + + +class _RetryableArticulation(Articulation): + bind_attempts = 0 + reset_attempts = 0 + + def __init__( + self, + cfg: object, + entities: list[object] | None = None, + device: object = "cpu", + *, + spawn_result: object | None = None, + declared_num_instances: int | None = None, + ) -> None: + self.cfg = cfg + self.uid = cfg.uid + self.device = device + self._entities = [] if entities is None else entities + self._spawn_result = spawn_result + self._world = None if spawn_result is None else object() + self._declared_num_instances = ( + len(entities) if entities is not None else int(declared_num_instances or 0) + ) + + def attach_spawn_handles(self, entities: list[object]) -> None: + self._entities = list(entities) + + def _apply_spawn_config(self) -> None: + type(self).bind_attempts += 1 + if type(self).bind_attempts == 1: + raise RuntimeError("configuration failed") + + def reset(self, env_ids: object | None = None) -> None: + del env_ids + type(self).reset_attempts += 1 + + +def test_articulation_binding_is_atomic_and_retryable() -> None: + _RetryableArticulation.bind_attempts = 0 + _RetryableArticulation.reset_attempts = 0 + facade = _RetryableArticulation( + SimpleNamespace(uid="robot"), + declared_num_instances=1, + ) + result = object() + handles = [object()] + facade.attach_spawn_handles(handles) + + with pytest.raises(RuntimeError, match="configuration failed"): + facade.bind_spawn(result) + + assert facade.is_declared + assert _RetryableArticulation.reset_attempts == 0 + facade.bind_spawn(result) + assert facade.is_spawn_bound + assert facade._entities == handles + assert _RetryableArticulation.reset_attempts == 1 diff --git a/tests/sim/test_backend_parity.py b/tests/sim/test_backend_parity.py new file mode 100644 index 000000000..c09f2149a --- /dev/null +++ b/tests/sim/test_backend_parity.py @@ -0,0 +1,191 @@ +# ---------------------------------------------------------------------------- +# 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}, + "volume_deformables": {"default": True, "newton": False}, + "surface_deformables": {"default": True, "newton": False}, + "soft_bodies": {"default": True, "newton": False}, + "cloth": {"default": True, "newton": False}, + "rigid_object_group": {"default": True, "newton": True}, + "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", + "volume_deformables": "add_deformable_object", + "surface_deformables": "add_deformable_object", + "soft_bodies": None, + "cloth": None, + "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._deformable_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) + + # Deformable dispatch needs its topology discriminator before the guard. + deformable_types = { + "volume_deformables": "volume", + "surface_deformables": "surface", + } + cfg = SimpleNamespace(uid=None) + if feature in deformable_types: + cfg.deformable_type = deformable_types[feature] + + 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): + 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 000000000..78bd7cd0c --- /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_cfg.py b/tests/sim/test_cfg.py index c9cfc28fe..3493fb63f 100644 --- a/tests/sim/test_cfg.py +++ b/tests/sim/test_cfg.py @@ -16,55 +16,836 @@ from __future__ import annotations +from dataclasses import fields + import dexsim import pytest +import embodichain.lab.sim.cfg as sim_cfg + +from dexsim.engine.newton_physics import ( + NewtonCollisionPipelineCfg as SpawnNewtonCollisionPipelineCfg, +) +from dexsim.spawn import DexsimCollisionDesc, DexsimPhysicsDesc, NewtonCollisionDesc from dexsim.types import DenoiserType, Renderer, ToneMappingType -from embodichain.lab.sim.cfg import ArticulationCfg, PhysicsCfg, RenderCfg, RobotCfg +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + ArticulationRootPropertiesCfg, + CollisionPropertiesCfg, + DefaultCollisionPropertiesCfg, + DefaultPhysicsCfg, + DefaultRigidBodyPropertiesCfg, + JointDrivePropertiesCfg, + LinkPhysicsOverrideCfg, + MassPropertiesCfg, + MeshCollisionCfg, + NewtonCollisionPipelineCfg, + NewtonCollisionPropertiesCfg, + NewtonJointDrivePropertiesCfg, + NewtonPhysicsCfg, + NewtonRigidBodyMaterialCfg, + PhysicsBackendCfg, + RenderCfg, + RigidBodyMaterialCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, + RobotCfg, + RobotPresetCfg, + physics_cfg_for_backend, +) +from embodichain.lab.sim.utility.cfg_utils import merge_robot_cfg +from embodichain.utils import configclass + +def test_cfg_package_preserves_the_public_facade() -> None: + from embodichain.lab.sim.cfg.rigid import ( + RigidBodyPhysicsCfg as LeafRigidBodyPhysicsCfg, + ) + from embodichain.lab.sim.cfg.robot import RobotCfg as LeafRobotCfg -def test_articulation_cfg_defaults_to_no_joint_drive() -> None: - """Generic articulations are passive unless a drive is requested.""" + assert hasattr(sim_cfg, "__path__") + assert not hasattr(sim_cfg, "PhysicsCfg") + assert sim_cfg.RigidBodyPhysicsCfg is LeafRigidBodyPhysicsCfg + assert sim_cfg.RobotCfg is LeafRobotCfg + + +def test_articulation_cfg_defaults_to_preserving_asset_physics() -> None: + """Generic articulations do not author source drive properties.""" articulation_cfg = ArticulationCfg() - assert articulation_cfg.drive_pros.drive_type == "none" + assert articulation_cfg.joint_drive_props is None + assert articulation_cfg.resolve_asset_physics_mode() == "preserve" + + +def test_articulation_cfg_uses_grouped_physics_fields_only() -> None: + field_names = {item.name for item in fields(ArticulationCfg)} + root_props = ArticulationCfg().root_props + + assert { + "fix_base", + "disable_self_collision", + "sleep_threshold", + "min_position_iters", + "min_velocity_iters", + "articulation_props", + "drive_pros", + "joint_props", + }.isdisjoint(field_names) + assert root_props == ArticulationRootPropertiesCfg() + assert root_props.fixed_base is True + assert root_props.self_collision_enabled is False + + +@pytest.mark.parametrize( + "field_name", + [ + "fix_base", + "disable_self_collision", + "sleep_threshold", + "min_position_iters", + "min_velocity_iters", + "articulation_props", + "drive_pros", + "joint_props", + ], +) +def test_removed_articulation_fields_fail_with_migration_target( + field_name: str, +) -> None: + with pytest.raises(ValueError, match=f"{field_name} ->"): + ArticulationCfg.from_dict({field_name: True}) + + with pytest.raises(ValueError, match=f"{field_name} ->"): + merge_robot_cfg(RobotCfg(), {field_name: True}) + + +def test_physics_cfg_factory_rejects_noncanonical_backend_names() -> None: + with pytest.raises(ValueError, match="expected 'default' or 'newton'"): + physics_cfg_for_backend("alternate") # type: ignore[arg-type] + +def test_articulation_cfg_parses_sparse_drive_overrides() -> None: + """Unspecified drive fields remain source-owned.""" + articulation_cfg = ArticulationCfg.from_dict( + {"joint_drive_props": {"stiffness": 0.0, "damping": 0.0}} + ) + + assert articulation_cfg.joint_drive_props.drive_type is None + assert articulation_cfg.joint_drive_props.stiffness == 0.0 + assert articulation_cfg.joint_drive_props.damping == 0.0 + assert articulation_cfg.joint_drive_props.max_effort is None + + +def test_robot_cfg_defaults_to_portable_position_velocity_drive() -> None: + """The original force drive resolves to position+velocity targets.""" + robot_cfg = RobotCfg() + + assert robot_cfg.joint_drive_props.drive_type == "force" + assert robot_cfg.joint_drive_props.target_mode is None + assert robot_cfg.joint_drive_props._resolve_modes() == ( + "position_velocity", + "force", + ) + assert robot_cfg.resolve_asset_physics_mode() == "overlay" + + +def test_robot_cfg_partial_drive_properties_preserve_portable_drive() -> None: + """Partial robot drive overrides retain the original force mode.""" + robot_cfg = RobotCfg.from_dict( + {"joint_drive_props": {"stiffness": 0.0, "damping": 0.0}} + ) -def test_articulation_cfg_partial_drive_properties_preserve_no_drive() -> None: - """Partial articulation drive overrides retain the passive default.""" + assert robot_cfg.joint_drive_props.drive_type == "force" + assert robot_cfg.joint_drive_props.target_mode is None + assert robot_cfg.joint_drive_props._resolve_modes() == ( + "position_velocity", + "force", + ) + + +def test_drive_type_override_replaces_robot_force_default() -> None: + override = {"joint_drive_props": {"drive_type": "none"}} + robot_cfg = RobotCfg.from_dict(override) + merged_cfg = merge_robot_cfg(RobotCfg(), override) + + for cfg in (robot_cfg, merged_cfg): + assert cfg.joint_drive_props.target_mode is None + assert cfg.joint_drive_props.drive_type == "none" + assert cfg.joint_drive_props._resolve_modes() == ("none", "none") + + +def test_common_target_mode_does_not_require_newton_subclass() -> None: articulation_cfg = ArticulationCfg.from_dict( - {"drive_pros": {"stiffness": 0.0, "damping": 0.0}} + { + "joint_drive_props": { + "target_mode": "effort", + "drive_type": "force", + } + } ) - assert articulation_cfg.drive_pros.drive_type == "none" + assert type(articulation_cfg.joint_drive_props) is JointDrivePropertiesCfg + assert articulation_cfg.joint_drive_props.target_mode == "effort" + assert articulation_cfg.joint_drive_props.drive_type == "force" -def test_robot_cfg_defaults_to_force_joint_drive() -> None: - """Robots retain force-based joint drives by default.""" +def test_asset_physics_policy_uses_explicit_modes() -> None: + rigid_cfg = RigidObjectCfg() + articulation_cfg = ArticulationCfg() robot_cfg = RobotCfg() + overlay_cfg = ArticulationCfg(asset_physics_mode="overlay") + + assert rigid_cfg.asset_physics_mode == "preserve" + assert rigid_cfg.resolve_asset_physics_mode() == "preserve" + assert articulation_cfg.asset_physics_mode == "preserve" + assert articulation_cfg.resolve_asset_physics_mode() == "preserve" + assert robot_cfg.asset_physics_mode == "overlay" + assert robot_cfg.resolve_asset_physics_mode() == "overlay" + assert overlay_cfg.resolve_asset_physics_mode() == "overlay" + + invalid_cfg = RigidObjectCfg(asset_physics_mode="replace") # type: ignore[arg-type] + with pytest.raises(ValueError, match="must be 'preserve' or 'overlay'"): + invalid_cfg.resolve_asset_physics_mode() + + +def test_articulation_cfg_parses_polymorphic_newton_joint_drive() -> None: + articulation_cfg = ArticulationCfg.from_dict( + { + "joint_drive_props": { + "backend": "newton", + "stiffness": {"arm_.*": 25.0}, + "target_mode": "position", + } + } + ) + + assert articulation_cfg.joint_drive_props.drive_type is None + assert isinstance(articulation_cfg.joint_drive_props, NewtonJointDrivePropertiesCfg) + assert articulation_cfg.joint_drive_props.stiffness == {"arm_.*": 25.0} + assert articulation_cfg.joint_drive_props.target_mode == "position" + + +def test_joint_drive_from_dict_preserves_newton_subclass_defaults() -> None: + defaults = NewtonJointDrivePropertiesCfg( + stiffness=10.0, + target_mode="position", + ) + + cfg = JointDrivePropertiesCfg.from_dict( + {"damping": 4.0}, + defaults=defaults, + ) + + assert isinstance(cfg, NewtonJointDrivePropertiesCfg) + assert cfg.stiffness == 10.0 + assert cfg.damping == 4.0 + assert cfg.target_mode == "position" + + +def test_robot_cfg_merge_preserves_typed_backend_property_configs() -> None: + base = RobotCfg( + joint_drive_props=NewtonJointDrivePropertiesCfg( + stiffness=10.0, + target_mode="position", + ), + attrs=RigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg(margin=0.01), + material_props=NewtonRigidBodyMaterialCfg(ke=1000.0), + ), + ) + + merged = merge_robot_cfg( + base, + { + "joint_drive_props": {"backend": "newton", "damping": 4.0}, + "attrs": {"material_props": {"backend": "newton", "kd": 50.0}}, + }, + ) + + assert isinstance(merged.joint_drive_props, NewtonJointDrivePropertiesCfg) + assert merged.joint_drive_props.stiffness == 10.0 + assert merged.joint_drive_props.damping == 4.0 + assert merged.joint_drive_props.target_mode == "position" + assert isinstance(merged.attrs, RigidBodyPhysicsCfg) + assert isinstance(merged.attrs.material_props, NewtonRigidBodyMaterialCfg) + assert merged.attrs.material_props.ke == 1000.0 + assert merged.attrs.material_props.kd == 50.0 + + +def test_rigid_physics_uses_one_slot_per_physical_concept() -> None: + """Backend blocks and geometry cooking are not parallel physics owners.""" + assert {item.name for item in fields(RigidBodyPhysicsCfg)} == { + "mass_props", + "rigid_props", + "collision_props", + "material_props", + } + assert issubclass(DefaultCollisionPropertiesCfg, CollisionPropertiesCfg) + assert issubclass(NewtonCollisionPropertiesCfg, CollisionPropertiesCfg) + assert issubclass(NewtonRigidBodyMaterialCfg, RigidBodyMaterialCfg) + assert issubclass(NewtonJointDrivePropertiesCfg, JointDrivePropertiesCfg) + for removed_name in ( + "DefaultRigidBodyPhysicsCfg", + "NewtonRigidBodyPhysicsCfg", + "MeshCollisionPropertiesCfg", + "NewtonMeshCollisionPropertiesCfg", + ): + assert not hasattr(sim_cfg, removed_name) + + +def test_backend_property_groups_track_dexsim_spawn_descriptors() -> None: + def names(config_type: type) -> set[str]: + return {item.name for item in fields(config_type)} + + assert names(DefaultRigidBodyPropertiesCfg) == names(DexsimPhysicsDesc) + default_collision_fields = ( + (names(CollisionPropertiesCfg) - {"collision_enabled"}) + | (names(DefaultCollisionPropertiesCfg) - names(CollisionPropertiesCfg)) + | names(RigidBodyMaterialCfg) + ) + assert default_collision_fields == names(DexsimCollisionDesc) + + newton_fields = ( + names(NewtonCollisionPropertiesCfg) - names(CollisionPropertiesCfg) + ) | (names(NewtonRigidBodyMaterialCfg) - names(RigidBodyMaterialCfg)) + newton_fields.remove("torsional_friction") + newton_fields.remove("rolling_friction") + newton_fields.update( + { + "mu", + "restitution", + "mu_torsional", + "mu_rolling", + "is_hydroelastic", + "sdf_narrow_band_range", + "sdf_target_voxel_size", + "sdf_max_resolution", + "sdf_texture_format", + "force_sdf", + "sdf_padding", + } + ) + intentionally_unowned_shape_fields = { + "is_solid", + "collision_group", + "collision_filter_parent", + "has_particle_collision", + "is_visible", + "is_site", + } + assert ( + newton_fields == names(NewtonCollisionDesc) - intentionally_unowned_shape_fields + ) + + assert names(NewtonCollisionPipelineCfg) == names( + SpawnNewtonCollisionPipelineCfg + ) - {"requires_grad"} + + +def test_rigid_physics_from_dict_selects_backend_subclasses() -> None: + cfg = RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 2.0}, + "rigid_props": {"backend": "default", "has_gravity": False}, + "collision_props": {"backend": "newton", "margin": 0.01}, + "material_props": { + "backend": "newton", + "dynamic_friction": 0.4, + "ke": 1000.0, + }, + } + ) + + assert isinstance(cfg.mass_props, MassPropertiesCfg) + assert isinstance(cfg.rigid_props, DefaultRigidBodyPropertiesCfg) + assert isinstance(cfg.collision_props, NewtonCollisionPropertiesCfg) + assert isinstance(cfg.material_props, NewtonRigidBodyMaterialCfg) + + +def test_recompute_inertia_is_a_mass_property() -> None: + cfg = RigidBodyPhysicsCfg.from_dict( + {"mass_props": {"mass": 2.0, "recompute_inertia": True}} + ) + + assert "replace_inertial" not in { + item.name for item in fields(LinkPhysicsOverrideCfg) + } + assert cfg.mass_props.recompute_inertia is True + assert cfg.to_dict()["mass_props"]["recompute_inertia"] is True + + link_cfg = LinkPhysicsOverrideCfg.from_dict( + { + "link_names_expr": ["finger_.*"], + "attrs": {"mass_props": {"recompute_inertia": True}}, + } + ) + assert link_cfg.attrs.mass_props.recompute_inertia is True + + with pytest.raises(ValueError, match="attrs.mass_props.recompute_inertia"): + LinkPhysicsOverrideCfg.from_dict( + { + "link_names_expr": ["finger_.*"], + "replace_inertial": True, + } + ) + + +@pytest.mark.parametrize( + ("removed_field", "replacement"), + [ + ("default_props", "polymorphic property slot"), + ("newton_props", "polymorphic property slot"), + ("mesh_collision_props", "MeshCfg.collision"), + ], +) +def test_rigid_physics_rejects_removed_parallel_owners( + removed_field: str, + replacement: str, +) -> None: + with pytest.raises(ValueError, match=replacement): + RigidBodyPhysicsCfg.from_dict({removed_field: {}}) + + +def test_articulation_cfg_parses_joint_drive_and_dynamics() -> None: + cfg = ArticulationCfg.from_dict( + { + "joint_drive_props": { + "stiffness": 12.0, + "max_effort": 20.0, + "friction": {"arm_.*": 0.2}, + }, + } + ) + + assert cfg.joint_drive_props.stiffness == pytest.approx(12.0) + assert cfg.joint_drive_props.max_effort == pytest.approx(20.0) + assert cfg.joint_drive_props.friction == {"arm_.*": 0.2} + + +def test_robot_cfg_merge_composes_single_slot_and_joint_drive_properties() -> None: + base = RobotCfg( + attrs=RigidBodyPhysicsCfg( + rigid_props=DefaultRigidBodyPropertiesCfg(linear_damping=0.1), + ), + joint_drive_props=JointDrivePropertiesCfg( + max_effort={"arm": 10.0}, + friction=0.1, + ), + ) + + merged = merge_robot_cfg( + base, + { + "attrs": { + "rigid_props": { + "backend": "default", + "angular_damping": 0.2, + }, + }, + "joint_drive_props": { + "max_effort": {"wrist": 20.0}, + "armature": 0.3, + }, + }, + ) + + assert merged.attrs.rigid_props.linear_damping == pytest.approx(0.1) + assert merged.attrs.rigid_props.angular_damping == pytest.approx(0.2) + assert merged.joint_drive_props.max_effort == {"arm": 10.0, "wrist": 20.0} + assert merged.joint_drive_props.friction == pytest.approx(0.1) + assert merged.joint_drive_props.armature == pytest.approx(0.3) + + +def test_portable_collision_envelope_round_trips_as_common_config() -> None: + cfg = RigidBodyPhysicsCfg.from_dict( + { + "collision_props": { + "collision_enabled": True, + "contact_offset": 0.01, + "rest_offset": 0.002, + } + } + ) + + assert type(cfg.collision_props) is CollisionPropertiesCfg + assert cfg.to_dict()["collision_props"] == { + "collision_enabled": True, + "contact_offset": 0.01, + "rest_offset": 0.002, + } + + +@configclass +class _RobotPhysicsPresetCfg(RobotPresetCfg): + default: RobotCfg = RobotCfg(uid="default") + newton: RobotCfg = RobotCfg(uid="newton") + newton_xpbd: RobotCfg = RobotCfg(uid="newton_xpbd") + + +def test_robot_preset_selects_complete_backend_and_solver_variants() -> None: + preset = _RobotPhysicsPresetCfg() + + default_cfg = preset.resolve(DefaultPhysicsCfg()) + newton_cfg = preset.resolve(NewtonPhysicsCfg()) + xpbd_cfg = preset.resolve(NewtonPhysicsCfg(solver_cfg={"solver_type": "xpbd"})) + + assert default_cfg.uid == "default" + assert newton_cfg.uid == "newton" + assert xpbd_cfg.uid == "newton_xpbd" + assert default_cfg is not preset.default + + +@configclass +class _CommonRobotPresetCfg(RobotPresetCfg): + default: RobotCfg = RobotCfg(uid="portable") + + +def test_robot_preset_falls_back_to_one_portable_definition() -> None: + preset = _CommonRobotPresetCfg() + + assert preset.resolve(DefaultPhysicsCfg()).uid == "portable" + assert preset.resolve(NewtonPhysicsCfg()).uid == "portable" + + +@configclass +class _NewtonSolverAliasRobotPresetCfg(RobotPresetCfg): + default: RobotCfg = RobotCfg(uid="fallback") + newton_mjwarp: RobotCfg = RobotCfg(uid="mjwarp") + + +def test_robot_preset_accepts_newton_solver_alias() -> None: + preset = _NewtonSolverAliasRobotPresetCfg() + + assert preset.resolve(DefaultPhysicsCfg()).uid == "fallback" + assert preset.resolve(NewtonPhysicsCfg()).uid == "fallback" + assert ( + preset.resolve(NewtonPhysicsCfg(solver_cfg={"solver_type": "mjwarp"})).uid + == "mjwarp" + ) + + +@configclass +class _UnsupportedRobotPresetCfg(RobotPresetCfg): + default: RobotCfg = RobotCfg(uid="default") + alternate: RobotCfg = RobotCfg(uid="alternate") + + +def test_robot_preset_rejects_noncanonical_backend_names() -> None: + with pytest.raises(TypeError, match="unsupported preset name"): + _UnsupportedRobotPresetCfg().resolve(DefaultPhysicsCfg()) + + +def test_backend_property_configs_round_trip_without_losing_subclasses() -> None: + cfg = RigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg(margin=0.01), + material_props=NewtonRigidBodyMaterialCfg(ke=1000.0), + ) + + serialized = cfg.to_dict() + restored = RigidBodyPhysicsCfg.from_dict(serialized) + + assert serialized["rigid_props"] is None + assert serialized["collision_props"]["backend"] == "newton" + assert serialized["material_props"]["backend"] == "newton" + assert isinstance(restored.collision_props, NewtonCollisionPropertiesCfg) + assert isinstance(restored.material_props, NewtonRigidBodyMaterialCfg) + + +def test_default_property_configs_use_the_default_discriminator() -> None: + cfg = RigidBodyPhysicsCfg( + rigid_props=DefaultRigidBodyPropertiesCfg(linear_damping=0.2), + collision_props=DefaultCollisionPropertiesCfg( + contact_offset=0.01, + disable_strong_friction=True, + ), + material_props=RigidBodyMaterialCfg(dynamic_friction=0.5), + ) + + serialized = cfg.to_dict() + restored = RigidBodyPhysicsCfg.from_dict(serialized) + + assert serialized["rigid_props"]["backend"] == "default" + assert serialized["collision_props"]["backend"] == "default" + assert "backend" not in serialized["material_props"] + assert isinstance(restored.rigid_props, DefaultRigidBodyPropertiesCfg) + assert isinstance(restored.collision_props, DefaultCollisionPropertiesCfg) + assert type(restored.material_props) is RigidBodyMaterialCfg + + +def test_backend_property_parser_infers_unique_fields_without_discriminator() -> None: + cfg = RigidBodyPhysicsCfg.from_dict( + { + "rigid_props": {"linear_damping": 0.2}, + "collision_props": {"margin": 0.01}, + "material_props": {"rolling_friction": 0.03}, + } + ) + + assert isinstance(cfg.rigid_props, DefaultRigidBodyPropertiesCfg) + assert isinstance(cfg.collision_props, NewtonCollisionPropertiesCfg) + assert isinstance(cfg.material_props, NewtonRigidBodyMaterialCfg) + + +def test_mesh_collision_cfg_requires_explicit_strategy_fields() -> None: + collision = MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=8, + acd_method="coacd", + ) + + assert collision.max_hulls == 8 + with pytest.raises(ValueError, match="valid only for convex_decomposition"): + MeshCollisionCfg(approximation="convex_hull", acd_method="coacd") + with pytest.raises(ValueError, match="only one"): + MeshCollisionCfg( + approximation="sdf", + sdf_resolution=64, + sdf_target_voxel_size=0.005, + ) + + +@pytest.mark.parametrize( + ("legacy_max_hulls", "expected_approximation", "expected_max_hulls"), + [ + (1, "convex_hull", None), + (4, "convex_decomposition", 4), + ], +) +def test_mesh_collision_cfg_accepts_deprecated_hull_count_alias( + legacy_max_hulls: int, + expected_approximation: str, + expected_max_hulls: int | None, +) -> None: + with pytest.warns(DeprecationWarning): + collision = MeshCollisionCfg(max_convex_hull_num=legacy_max_hulls) + + assert collision.approximation == expected_approximation + assert collision.max_hulls == expected_max_hulls + assert "max_convex_hull_num" not in collision.to_dict() + + +def test_mesh_collision_cfg_deprecated_hull_count_view_uses_canonical_value() -> None: + collision = MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=4, + ) + + assert collision.max_convex_hull_num == 4 + + +def test_mesh_collision_cfg_rejects_both_hull_count_names() -> None: + with ( + pytest.warns(DeprecationWarning), + pytest.raises(ValueError, match="cannot both be configured"), + ): + MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=4, + max_convex_hull_num=8, + ) + + +@pytest.mark.parametrize( + "collision_kwargs", + [ + {"approximation": "convex_decomposition", "max_hulls": 2.5}, + {"approximation": "sdf", "sdf_resolution": 64.5}, + {"approximation": "sdf", "sdf_padding": float("nan")}, + {"approximation": "sdf", "sdf_texture_format": "invalid"}, + ], +) +def test_mesh_collision_cfg_rejects_invalid_numeric_types_and_values( + collision_kwargs: dict[str, object], +) -> None: + with pytest.raises(ValueError): + MeshCollisionCfg(**collision_kwargs) + + +def test_mesh_cfg_legacy_collision_fields_normalize_to_nested_config() -> None: + with pytest.warns(DeprecationWarning): + cfg = RigidObjectCfg.from_dict( + { + "uid": "mesh", + "shape": { + "shape_type": "Mesh", + "fpath": "mesh.obj", + "max_convex_hull_num": 4, + "acd_method": "coacd", + }, + } + ) + + assert cfg.shape.collision == MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=4, + acd_method="coacd", + ) + serialized_shape = cfg.shape.to_dict() + assert "max_convex_hull_num" not in serialized_shape + assert serialized_shape["collision"]["approximation"] == "convex_decomposition" + + +def test_rigid_object_legacy_physics_mesh_collision_moves_to_shape() -> None: + with pytest.warns(DeprecationWarning): + cfg = RigidObjectCfg.from_dict( + { + "uid": "mesh", + "shape": {"shape_type": "Mesh", "fpath": "mesh.obj"}, + "attrs": { + "mesh_collision_props": {"max_convex_hull_num": 4}, + }, + } + ) + + assert cfg.shape.collision.approximation == "convex_decomposition" + assert cfg.shape.collision.max_hulls == 4 + assert "mesh_collision_props" not in cfg.attrs.to_dict() + + +def test_legacy_mesh_collision_physics_rejects_non_mesh_shape() -> None: + with pytest.raises(ValueError, match="only to a MeshCfg"): + RigidObjectCfg.from_dict( + { + "uid": "cube", + "shape": {"shape_type": "Cube", "size": [1.0, 1.0, 1.0]}, + "attrs": { + "mesh_collision_props": {"max_convex_hull_num": 4}, + }, + } + ) + + +def test_backend_joint_and_articulation_configs_round_trip() -> None: + drive = NewtonJointDrivePropertiesCfg(target_mode=None) + root = ArticulationRootPropertiesCfg(fixed_base=False) + + restored_drive = JointDrivePropertiesCfg.from_dict(drive.to_dict()) + restored_root = ArticulationRootPropertiesCfg.from_dict(root.to_dict()) + + assert isinstance(restored_drive, NewtonJointDrivePropertiesCfg) + assert root.to_dict() == { + "fixed_base": False, + "self_collision_enabled": False, + "sleep_threshold": None, + "min_position_iters": None, + "min_velocity_iters": None, + } + assert type(restored_root) is ArticulationRootPropertiesCfg + + +def test_articulation_root_config_rejects_backend_discriminator() -> None: + with pytest.raises(TypeError, match="backend"): + ArticulationRootPropertiesCfg.from_dict( + {"backend": "newton", "fixed_base": False} + ) + + +def test_articulation_root_config_round_trip() -> None: + root = ArticulationRootPropertiesCfg( + fixed_base=True, + sleep_threshold=0.005, + min_position_iters=8, + min_velocity_iters=2, + ) + + restored = ArticulationRootPropertiesCfg.from_dict(root.to_dict()) + + assert type(restored) is ArticulationRootPropertiesCfg + assert restored == root + assert "backend" not in root.to_dict() + + +def test_articulation_root_requires_both_solver_iteration_counts() -> None: + with pytest.raises(ValueError, match="must be configured together"): + ArticulationRootPropertiesCfg(min_position_iters=8) + + +def test_robot_cfg_round_trip_preserves_grouped_backend_types() -> None: + cfg = RobotCfg( + attrs=RigidBodyPhysicsCfg( + collision_props=NewtonCollisionPropertiesCfg(margin=0.01), + material_props=NewtonRigidBodyMaterialCfg(ke=1000.0), + ), + joint_drive_props=NewtonJointDrivePropertiesCfg(target_mode="position"), + root_props=ArticulationRootPropertiesCfg(fixed_base=False), + ) + + restored = RobotCfg.from_dict(cfg.to_dict()) + + assert isinstance(restored.attrs, RigidBodyPhysicsCfg) + assert isinstance(restored.attrs.collision_props, NewtonCollisionPropertiesCfg) + assert isinstance(restored.attrs.material_props, NewtonRigidBodyMaterialCfg) + assert isinstance(restored.joint_drive_props, NewtonJointDrivePropertiesCfg) + assert type(restored.root_props) is ArticulationRootPropertiesCfg + + +def test_rigid_physics_from_dict_rejects_unknown_fields() -> None: + with pytest.raises((KeyError, TypeError)): + RigidBodyPhysicsCfg.from_dict({"collision_props": {"margn": 0.01}}) + + +def test_robot_cfg_merge_preserves_grouped_overrides() -> None: + base = RobotCfg( + attrs=RigidBodyPhysicsCfg( + material_props=RigidBodyMaterialCfg(dynamic_friction=0.8) + ) + ) + + merged = merge_robot_cfg(base, {"attrs": {"mass_props": {"mass": 2.0}}}) + + assert isinstance(merged.attrs, RigidBodyPhysicsCfg) + assert merged.attrs.mass_props.mass == 2.0 + assert merged.attrs.material_props.dynamic_friction == 0.8 + + +def test_newton_physics_inherits_common_gravity_and_collision_config() -> None: + cfg = NewtonPhysicsCfg( + gravity=[0.0, 0.0, -1.5], + collision_cfg=NewtonCollisionPipelineCfg( + broad_phase="sap", + rigid_contact_max=1234, + ), + ) + + assert isinstance(cfg, PhysicsBackendCfg) + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + assert dexsim_cfg.gravity == [0.0, 0.0, -1.5] + assert dexsim_cfg.collision_pipeline_cfg.broad_phase == "sap" + assert dexsim_cfg.collision_pipeline_cfg.rigid_contact_max == 1234 + + +def test_newton_physics_normalizes_mapping_collision_config() -> None: + cfg = NewtonPhysicsCfg( + collision_cfg={"broad_phase": "sap", "rigid_contact_max": 12} + ) + + assert isinstance(cfg.collision_cfg, NewtonCollisionPipelineCfg) + assert cfg.collision_cfg.broad_phase == "sap" + assert cfg.collision_cfg.rigid_contact_max == 12 - assert robot_cfg.drive_pros.drive_type == "force" +def test_default_physics_accepts_the_same_gravity_input_shape() -> None: + cfg = DefaultPhysicsCfg(gravity=[0.0, 0.0, -1.5]) -def test_robot_cfg_partial_drive_properties_preserve_force_drive() -> None: - """Partial robot drive overrides retain the force-drive default.""" - robot_cfg = RobotCfg.from_dict({"drive_pros": {"stiffness": 0.0, "damping": 0.0}}) + assert cfg.to_dexsim_args()["gravity"] == [0.0, 0.0, -1.5] + assert DefaultPhysicsCfg().to_dexsim_args()["gravity"] == [0.0, 0.0, -9.81] - assert robot_cfg.drive_pros.drive_type == "force" + with pytest.raises(ValueError, match="three finite values"): + DefaultPhysicsCfg(gravity=[0.0, -9.81]).to_dexsim_args() -def test_physics_cfg_does_not_expose_fixed_solver_options() -> None: +def test_default_physics_cfg_does_not_expose_fixed_solver_options() -> None: """Fixed solver implementation details are not part of the public config.""" - physics_cfg = PhysicsCfg() + physics_cfg = DefaultPhysicsCfg() assert not hasattr(physics_cfg, "enable_enhanced_determinism") assert not hasattr(physics_cfg, "enable_friction_every_iteration") -def test_physics_cfg_applies_fixed_solver_defaults() -> None: - """Removed solver options retain their established DexSim defaults.""" - physics_args = PhysicsCfg(enable_ccd=True).to_dexsim_args() +def test_default_physics_cfg_applies_fixed_solver_defaults() -> None: + """Removed solver options retain the Default backend's established values.""" + physics_args = DefaultPhysicsCfg(enable_ccd=True).to_dexsim_args() assert physics_args["enable_ccd"] is True assert physics_args["enable_enhanced_determinism"] is False diff --git a/tests/sim/test_differentiable_stepper.py b/tests/sim/test_differentiable_stepper.py new file mode 100644 index 000000000..4006da542 --- /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.prepare() + 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.prepare() + 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.prepare() + 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_grasp_cup_to_caffe_demo.py b/tests/sim/test_grasp_cup_to_caffe_demo.py new file mode 100644 index 000000000..190dd414b --- /dev/null +++ b/tests/sim/test_grasp_cup_to_caffe_demo.py @@ -0,0 +1,143 @@ +# ---------------------------------------------------------------------------- +# 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 importlib.util +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +pytestmark = pytest.mark.no_sim + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +_DEMO_PATH = _REPOSITORY_ROOT / "examples/sim/demo/grasp_cup_to_caffe.py" +_INITIAL_PHYSICS_STEPS = 1 +_IDLE_LOOP_PHYSICS_STEPS = 10 + + +def _load_demo_module() -> ModuleType: + spec = importlib.util.spec_from_file_location("grasp_cup_to_caffe_demo", _DEMO_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_scene_perturbations_precede_first_physics_step(monkeypatch) -> None: + demo = _load_demo_module() + events: list[str] = [] + + class FakeSimulation: + def prepare(self) -> None: + events.append("prepare") + + def update(self, step: int) -> None: + events.append(f"update:{step}") + if step == _IDLE_LOOP_PHYSICS_STEPS: + raise KeyboardInterrupt + + def open_window(self) -> None: + events.append("open_window") + + sim = FakeSimulation() + robot = object() + cup = object() + caffe = object() + monkeypatch.setattr( + demo, + "parse_arguments", + lambda: SimpleNamespace(headless=True, seed=0), + ) + monkeypatch.setattr(demo, "initialize_simulation", lambda _args: sim) + monkeypatch.setattr(demo, "create_robot", lambda _sim: robot) + monkeypatch.setattr(demo, "create_table", lambda _sim: object()) + monkeypatch.setattr(demo, "create_caffe", lambda _sim: caffe) + monkeypatch.setattr(demo, "create_cup", lambda _sim: cup) + monkeypatch.setattr( + demo, + "apply_random_xy_perturbation", + lambda item, **_kwargs: events.append( + "perturb:cup" if item is cup else "perturb:caffe" + ), + ) + monkeypatch.setattr( + demo, + "run_simulation", + lambda *_args: events.append("run_simulation"), + ) + monkeypatch.setattr( + demo.np.random, + "seed", + lambda seed: events.append(f"seed:{seed}"), + ) + + demo.main() + + assert events[:5] == [ + "prepare", + "seed:0", + "perturb:cup", + "perturb:caffe", + f"update:{_INITIAL_PHYSICS_STEPS}", + ] + + +def test_trajectory_uses_authored_hold_target_as_ik_seed(monkeypatch) -> None: + demo = _load_demo_module() + target_reads: list[bool] = [] + + class FakeRobot: + def get_joint_ids(self, name: str) -> list[int]: + assert name == "right_arm" + return [0, 1] + + def get_qpos(self, target: bool = False) -> torch.Tensor: + target_reads.append(target) + return torch.tensor([[0.25, -0.5]], dtype=torch.float32) + + def compute_fk(self, **_kwargs) -> torch.Tensor: + return torch.eye(4, dtype=torch.float32).unsqueeze(0) + + def compute_ik( + self, *, joint_seed: torch.Tensor, **_kwargs + ) -> tuple[torch.Tensor, torch.Tensor]: + return torch.ones(1, dtype=torch.bool), joint_seed.clone() + + class FakeItem: + def get_local_pose(self, *, to_matrix: bool) -> torch.Tensor: + assert to_matrix + return torch.eye(4, dtype=torch.float32).unsqueeze(0) + + monkeypatch.setattr( + demo, + "interpolate_with_distance", + lambda trajectory, **_kwargs: trajectory, + ) + + trajectory = demo.create_trajectory( + SimpleNamespace( + device=torch.device("cpu"), num_envs=1, is_newton_backend=False + ), + FakeRobot(), + FakeItem(), + FakeItem(), + ) + + assert target_reads == [True] + assert trajectory.shape == (1, 10, 8) diff --git a/tests/sim/test_open_drawer_tutorial.py b/tests/sim/test_open_drawer_tutorial.py new file mode 100644 index 000000000..ad65fd9ba --- /dev/null +++ b/tests/sim/test_open_drawer_tutorial.py @@ -0,0 +1,179 @@ +# ---------------------------------------------------------------------------- +# Copyright (c) 2021-2026 DexForce Technology Co., Ltd. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ---------------------------------------------------------------------------- + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +pytestmark = pytest.mark.no_sim + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +_TUTORIAL_PATH = _REPOSITORY_ROOT / "scripts/tutorials/sim/open_drawer.py" + + +def _load_tutorial_module() -> ModuleType: + spec = importlib.util.spec_from_file_location( + "open_drawer_tutorial", _TUTORIAL_PATH + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize("is_newton_backend", [False, True]) +def test_create_scene_configures_newton_grasp_material_only_for_newton( + monkeypatch, + is_newton_backend, +) -> None: + tutorial = _load_tutorial_module() + robot = object() + drawer = object() + captured: dict[str, object] = {} + + class FakeSimulation: + def __init__(self): + self.is_newton_backend = is_newton_backend + + def add_robot(self, cfg): + captured["robot_cfg"] = cfg + return robot + + def add_articulation(self, cfg): + captured["drawer_cfg"] = cfg + return drawer + + monkeypatch.setattr( + tutorial.FrankaPandaCfg, + "from_dict", + lambda _config: SimpleNamespace( + joint_drive_props=SimpleNamespace(damping={}), + link_attrs=None, + ), + ) + monkeypatch.setattr(tutorial, "get_data_path", lambda asset: asset) + + tutorial.create_scene(FakeSimulation()) + + drawer_cfg = captured["drawer_cfg"] + robot_cfg = captured["robot_cfg"] + assert robot_cfg.joint_drive_props.damping == {} + assert drawer_cfg.root_props.fixed_base is True + if is_newton_backend: + robot_material = robot_cfg.link_attrs[ + "newton_gripper_contacts" + ].attrs.material_props + drawer_override = drawer_cfg.link_attrs["newton_handle_contacts"] + drawer_material = drawer_override.attrs.material_props + assert robot_material.ke == pytest.approx( + tutorial.NEWTON_GRASP_CONTACT_STIFFNESS + ) + assert robot_material.kd == pytest.approx(tutorial.NEWTON_GRASP_CONTACT_DAMPING) + assert drawer_override.link_names_expr == [tutorial.DRAWER_CONTACT_LINK_NAME] + assert drawer_material.ke == pytest.approx( + tutorial.NEWTON_GRASP_CONTACT_STIFFNESS + ) + assert drawer_material.kd == pytest.approx( + tutorial.NEWTON_GRASP_CONTACT_DAMPING + ) + else: + assert robot_cfg.link_attrs is None + assert drawer_cfg.link_attrs is None + + +def test_tutorial_newton_physics_cfg_enables_multiccd() -> None: + tutorial = _load_tutorial_module() + + cfg = tutorial._tutorial_physics_cfg("newton") + + assert cfg.num_substeps == 20 + assert cfg.solver_cfg == { + "solver_type": "mujoco_warp", + "njmax": 8192, + "nconmax": 8192, + "cone": "elliptic", + "enable_multiccd": True, + } + + +def test_main_opens_native_window_after_spawn_prepare(monkeypatch) -> None: + tutorial = _load_tutorial_module() + events: list[str] = [] + captured_cfg: dict[str, object] = {} + args = SimpleNamespace( + num_envs=1, + hold_steps=0, + record_fps=30, + record_save_path=None, + headless=False, + viser=False, + auto_start=True, + physics="default", + device="cpu", + arena_space=2.0, + renderer="hybrid", + ) + + class FakeSimulation: + num_envs = 1 + + def prepare(self) -> None: + events.append("prepare") + + def open_window(self) -> None: + events.append("open_window") + + def update(self, *, step: int) -> None: + pass + + def is_window_recording(self) -> bool: + return False + + def wait_window_record_saves(self) -> None: + pass + + def destroy(self) -> None: + pass + + monkeypatch.setattr( + tutorial.argparse.ArgumentParser, + "parse_args", + lambda _parser: args, + ) + monkeypatch.setattr( + tutorial, + "SimulationManagerCfg", + lambda **kwargs: captured_cfg.update(kwargs) or kwargs, + ) + monkeypatch.setattr(tutorial, "SimulationManager", lambda _cfg: FakeSimulation()) + monkeypatch.setattr( + tutorial, + "create_scene", + lambda _sim: events.append("create_scene") + or (SimpleNamespace(uid="robot"), object()), + ) + monkeypatch.setattr(tutorial, "MotionGenerator", lambda *, cfg: object()) + monkeypatch.setattr(tutorial, "open_drawer", lambda *_args, **_kwargs: None) + monkeypatch.setattr(tutorial, "visualization_cfg_from_args", lambda _args: None) + + tutorial.main() + + assert captured_cfg["headless"] is True + assert events == ["create_scene", "prepare", "open_window"] diff --git a/tests/sim/test_rigid_constraint_integration.py b/tests/sim/test_rigid_constraint_integration.py index 5e6aaf5df..baac80f33 100644 --- a/tests/sim/test_rigid_constraint_integration.py +++ b/tests/sim/test_rigid_constraint_integration.py @@ -38,7 +38,7 @@ from embodichain.lab.sim.cfg import ( RigidObjectCfg, RigidConstraintCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, ) from embodichain.lab.sim.shapes import MeshCfg @@ -64,21 +64,19 @@ def _delta_z(self) -> float: pose_b = self.duck_b.get_local_pose(to_matrix=True) return float(pose_b[0, 2, 3] - pose_a[0, 2, 3]) - def setup_simulation(self, sim_device: str) -> None: - if not _can_run_sim(sim_device): - pytest.skip( - f"Cannot run rigid-constraint integration test on {sim_device}." - ) - config = SimulationManagerCfg(headless=True, sim_device=sim_device, num_envs=1) + def setup_simulation(self, device: str) -> None: + if not _can_run_sim(device): + pytest.skip(f"Cannot run rigid-constraint integration test on {device}.") + config = SimulationManagerCfg(headless=True, device=device, num_envs=1) self.sim = SimulationManager(config) self.sim.enable_physics(False) duck_path = get_data_path(DUCK_PATH) # Two dynamic ducks at different heights; with default (None) local # frames the constraint welds them at their current relative pose. - attrs_a = RigidBodyAttributesCfg() + attrs_a = RigidBodyPhysicsCfg() attrs_a.mass = 0.2 - attrs_b = RigidBodyAttributesCfg() + attrs_b = RigidBodyPhysicsCfg() attrs_b.mass = 0.1 self.duck_a = self.sim.add_rigid_object( cfg=RigidObjectCfg( @@ -99,8 +97,7 @@ def setup_simulation(self, sim_device: str) -> None: ), ) - if sim_device == "cuda" and getattr(self.sim, "is_use_gpu_physics", False): - self.sim.init_gpu_physics() + self.sim.prepare() self.sim.enable_physics(True) def teardown_method(self): diff --git a/tests/sim/test_rigid_physics_cfg.py b/tests/sim/test_rigid_physics_cfg.py new file mode 100644 index 000000000..829a55cca --- /dev/null +++ b/tests/sim/test_rigid_physics_cfg.py @@ -0,0 +1,77 @@ +# ---------------------------------------------------------------------------- +# 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 grouped rigid-body physics configuration boundary.""" + +from __future__ import annotations + +import numpy as np +import pytest + +import embodichain.lab.sim as sim +import embodichain.lab.sim.cfg as sim_cfg +from embodichain.lab.sim.cfg import ( + ArticulationCfg, + RigidBodyPhysicsCfg, + RigidObjectCfg, +) + + +def test_public_cfg_facade_no_longer_exports_flat_rigid_attribute_types() -> None: + for facade in (sim, sim_cfg): + assert not hasattr(facade, "RigidBodyAttributesCfg") + assert not hasattr(facade, "RigidBodyAttributesOverrideCfg") + + +def test_grouped_cfg_converts_com_quaternion_only_at_dexsim_boundary() -> None: + input_quaternion_xyzw = [1.0, 2.0, 3.0, 4.0] + cfg = RigidBodyPhysicsCfg.from_dict( + { + "mass_props": { + "mass": 2.0, + "inertia": [1.0, 2.0, 3.0], + "com_position": [0.1, 0.2, 0.3], + "com_quaternion": input_quaternion_xyzw, + }, + "material_props": {"dynamic_friction": 0.4}, + } + ) + + native = cfg.to_dexsim_physical_attr() + restored = RigidBodyPhysicsCfg.from_dexsim_physical_attr(native) + + np.testing.assert_allclose(native.com_quaternion, [4.0, 1.0, 2.0, 3.0]) + np.testing.assert_allclose( + restored.mass_props.com_quaternion, + input_quaternion_xyzw, + ) + assert restored.mass_props.mass == pytest.approx(2.0) + assert restored.material_props.dynamic_friction == pytest.approx(0.4) + + +@pytest.mark.parametrize("config_type", [RigidObjectCfg, ArticulationCfg]) +def test_asset_config_rejects_removed_flat_rigid_attributes(config_type: type) -> None: + with pytest.raises(ValueError, match="Removed flat rigid-body attrs fields"): + config_type.from_dict({"attrs": {"mass": 2.0}}) + + +def test_grouped_attrs_parse_for_rigid_and_articulation_configs() -> None: + rigid = RigidObjectCfg.from_dict({"attrs": {"mass_props": {"mass": 2.0}}}) + articulation = ArticulationCfg.from_dict( + {"attrs": {"material_props": {"static_friction": 0.8}}} + ) + + assert rigid.attrs.mass_props.mass == pytest.approx(2.0) + assert articulation.attrs.material_props.static_friction == pytest.approx(0.8) diff --git a/tests/sim/test_sim_manager.py b/tests/sim/test_sim_manager.py index 0aab8443a..07a7bbcf7 100644 --- a/tests/sim/test_sim_manager.py +++ b/tests/sim/test_sim_manager.py @@ -20,25 +20,32 @@ import queue from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import numpy as np import pytest import torch import embodichain.lab.sim.sim_manager as sim_manager_module +from embodichain.lab.sim.cfg import ( + DefaultPhysicsCfg, + RobotCfg, + RobotPresetCfg, +) from embodichain.lab.sim.profiler import Profiler from embodichain.lab.sim.sim_manager import ( SimulationManager, SimulationManagerCfg, _WindowRecordState, ) +from embodichain.lab.sim.sensors import CameraCfg from embodichain.lab.visualization import ( GizmoCommand, PointCloudOverlay, SceneOverlays, VisualizationCfg, ) +from embodichain.utils import configclass DEFAULT_LOOK_AT = ( (2.6, -2.2, 1.6), @@ -211,12 +218,14 @@ def _make_visualization_sim_manager() -> ( runtime = FakeVisualizationRuntime() sim.sim_config = SimpleNamespace( physics_dt=0.01, + physics_cfg=DefaultPhysicsCfg(), visualization=SimpleNamespace(backend="viser"), ) sim.device = SimpleNamespace(type="cpu") sim.profiler = Profiler(None, torch.device("cpu")) sim._is_initialized_gpu_physics = False sim._world = FakeWorld() + sim.prepare = MagicMock() sim._window_record_state = None sim._visualization_runtime = runtime sim._visualization_overlays = None @@ -267,6 +276,54 @@ def test_flush_cleanup_queue_waits_after_running_pending_destroy( wait_scene_destruction.assert_called_once_with() +def test_deferred_destroy_prepares_backend_before_releasing_world( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Backend-owned views are released before Spawn and World resources.""" + events: list[str] = [] + sim = object.__new__(SimulationManager) + spawn_scene = MagicMock() + spawn_scene.close.side_effect = lambda: events.append("spawn_close") + sim.physics = SimpleNamespace( + prepare_for_teardown=lambda: events.append("backend_prepare") + ) + sim._gizmos = {} + sim._markers = {} + sim._rigid_objects = {} + sim._constraints = {} + sim._rigid_object_groups = {} + sim._deformable_objects = {} + sim._articulations = {} + sim._robots = {} + sim._sensors = {} + sim._lights = {} + sim._visual_materials = {} + sim._texture_cache = {} + sim._arenas = [] + sim._spawn_scene = spawn_scene + sim._default_plane = object() + sim._env = SimpleNamespace(clean=lambda: events.append("env_clean")) + sim._world = SimpleNamespace(quit=lambda: events.append("world_quit")) + sim.instance_id = 0 + sim.is_window_recording = lambda: False + sim.wait_window_record_saves = lambda: events.append("record_wait") + sim.clean_materials = lambda: events.append("material_clean") + sim.is_window_opened = False + + monkeypatch.setattr( + SimulationManager, + "reset", + lambda _instance_id: events.append("manager_reset"), + ) + monkeypatch.setattr(gc, "collect", lambda: events.append("gc_collect")) + + sim._deferred_destroy() + + assert events.index("backend_prepare") < events.index("gc_collect") + assert events.index("backend_prepare") < events.index("spawn_close") + assert events.index("backend_prepare") < events.index("world_quit") + + def test_sim_update_refreshes_dirty_visualization_and_captures_current_state() -> None: sim, runtime = _make_visualization_sim_manager() @@ -475,8 +532,9 @@ def test_start_visualization_rejects_open_native_window() -> None: sim.start_visualization() -def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> None: +def test_constructor_only_declares_spawn_scene(monkeypatch) -> None: lifecycle: list[str] = [] + spawn_scene = MagicMock() world = MagicMock() world.get_physics_scene.return_value = MagicMock() world.get_env.return_value = MagicMock() @@ -486,6 +544,11 @@ def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> No ) monkeypatch.setattr(sim_manager_module.wp, "init", lambda: None) monkeypatch.setattr(sim_manager_module.dexsim, "World", lambda _cfg: world) + monkeypatch.setattr( + sim_manager_module, + "SpawnScene", + lambda *_args, **_kwargs: spawn_scene, + ) monkeypatch.setattr( sim_manager_module.dexsim, "set_physics_config", lambda **_kwargs: None ) @@ -507,7 +570,7 @@ def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> No ) monkeypatch.setattr( SimulationManager, - "_create_default_plane", + "_declare_spawn_default_plane", lambda _self: lifecycle.append("plane"), ) monkeypatch.setattr( @@ -521,14 +584,9 @@ def test_constructor_starts_visualization_after_default_scene(monkeypatch) -> No lambda _self: lifecycle.append("lighting"), ) - def build_arenas(sim: SimulationManager, num: int) -> None: - lifecycle.append("arenas") - sim._arenas.extend([object() for _ in range(num)]) - def start_visualization(sim: SimulationManager) -> None: lifecycle.append(f"visualization:{sim.num_envs}") - monkeypatch.setattr(SimulationManager, "_build_multiple_arenas", build_arenas) monkeypatch.setattr( SimulationManager, "start_visualization", @@ -540,22 +598,362 @@ def start_visualization(sim: SimulationManager) -> None: assert lifecycle == [ "resources", - "plane", "background", + "plane", "lighting", - "arenas", - "visualization:3", ] + assert sim._spawn_scene is spawn_scene + assert sim._arenas == [] + + +def test_add_robot_resolves_backend_preset_before_declaration() -> None: + @configclass + class TestRobotPresetCfg(RobotPresetCfg): + default: RobotCfg = RobotCfg(uid="selected", fpath="selected.urdf") + + sim = object.__new__(SimulationManager) + sim.physics = SimpleNamespace(name="default", supports_robot=True) + sim.sim_config = SimpleNamespace(physics_cfg=DefaultPhysicsCfg()) + sim._robots = {} + sim._declare_spawn_articulation = MagicMock(return_value="robot-handle") + + robot = sim.add_robot(TestRobotPresetCfg()) + + assert robot == "robot-handle" + resolved_cfg = sim._declare_spawn_articulation.call_args.args[0] + assert isinstance(resolved_cfg, RobotCfg) + assert resolved_cfg.uid == "selected" + + +def test_default_plane_authors_repeated_uv_before_spawn() -> None: + sim = object.__new__(SimulationManager) + sim._spawn_scene = MagicMock() + sim._spawn_scene.handles.return_value = [] + sim._spawn_default_plane_material = object() + + sim._declare_spawn_default_plane() + + descriptor = sim._spawn_scene.declare.call_args.args[2] + expected_repeat = 500.0 # One two-metre texture tile across a 1000 m plane. + np.testing.assert_array_equal( + descriptor.renders[0].uv_coords, + np.asarray( + [ + [0.0, 0.0], + [expected_repeat, 0.0], + [expected_repeat, expected_repeat], + [0.0, expected_repeat], + ], + dtype=np.float32, + ), + ) + + +@pytest.mark.parametrize( + ("backend", "device", "initializes_direct_gpu"), + [ + pytest.param("default", torch.device("cpu"), False, id="default-host"), + pytest.param("default", torch.device("cuda"), True, id="default-accelerator"), + pytest.param("newton", torch.device("cpu"), False, id="newton-host"), + pytest.param("newton", torch.device("cuda"), False, id="newton-accelerator"), + ], +) +def test_prepare_initializes_runtime_for_backend_device_matrix( + backend: str, + device: torch.device, + initializes_direct_gpu: bool, +) -> None: + result = MagicMock() + result.topology_revision = 3 + spawn_scene = MagicMock() + spawn_scene.builder.is_finalized = False + spawn_scene.builder.result = None + spawn_scene.commit.return_value = result + spawn_scene.arena_names = ["arena_0"] + events: list[str] = [] + spawn_scene.prepare_runtime_config.side_effect = lambda _result: events.append( + "runtime_config" + ) + spawn_scene.bind.side_effect = lambda: events.append("bind") + + sim = object.__new__(SimulationManager) + sync_render_state = MagicMock() + sim.physics = SimpleNamespace( + name=backend, + sync_render_state=sync_render_state, + ) + sim.device = device + sim._world = MagicMock() + sim._world.init_gpu_physics.side_effect = lambda: events.append("gpu_init") + sim._spawn_scene = spawn_scene + sim._default_plane = object() + sim._pending_sensor_attachments = [] + sim._prepared_spawn_topology_revision = -1 + sim._synced_spawn_render_topology_revision = -1 + + sim.prepare() + + spawn_scene.prepare_runtime_config.assert_called_once_with(result) + spawn_scene.bind.assert_called_once_with() + sync_render_state.assert_called_once_with(result) + sim._world.update.assert_not_called() + if initializes_direct_gpu: + sim._world.init_gpu_physics.assert_called_once_with() + assert events == ["runtime_config", "gpu_init", "bind"] + else: + sim._world.init_gpu_physics.assert_not_called() + assert events == ["runtime_config", "bind"] + + +def test_prepare_retries_runtime_and_binding_without_recommit() -> None: + result = MagicMock() + result.needs_rebuild = False + result.topology_revision = 3 + spawn_scene = MagicMock() + spawn_scene.builder.is_finalized = True + spawn_scene.builder.result = result + spawn_scene.builder.has_pending_changes = False + + sim = object.__new__(SimulationManager) + sync_render_state = MagicMock() + sim.physics = SimpleNamespace( + name="default", + sync_render_state=sync_render_state, + ) + sim.device = torch.device("cuda") + sim._world = MagicMock() + sim._world.init_gpu_physics.side_effect = [RuntimeError("first attempt"), None] + sim._spawn_scene = spawn_scene + sim._pending_sensor_attachments = [] + sim._prepared_spawn_topology_revision = -1 + sim._synced_spawn_render_topology_revision = -1 + + with pytest.raises(RuntimeError, match="first attempt"): + sim.prepare() + sim.prepare() + + spawn_scene.commit.assert_not_called() + assert sim._world.init_gpu_physics.call_count == 2 + spawn_scene.bind.assert_called_once_with() + sync_render_state.assert_called_once_with(result) + + +def test_prepare_removes_each_sensor_after_successful_attachment() -> None: + result = MagicMock() + result.needs_rebuild = False + result.topology_revision = 3 + first_sensor = MagicMock() + second_sensor = MagicMock() + attach_camera_parent = MagicMock( + side_effect=[None, RuntimeError("attach failed"), None] + ) + spawn_scene = MagicMock() + spawn_scene.builder.is_finalized = True + spawn_scene.builder.result = result + spawn_scene.builder.has_pending_changes = False + + sim = object.__new__(SimulationManager) + sync_render_state = MagicMock() + sim.physics = SimpleNamespace( + name="default", + sync_render_state=sync_render_state, + ) + sim.device = torch.device("cpu") + sim._world = MagicMock() + sim._spawn_scene = spawn_scene + sim._attach_camera_parent = attach_camera_parent + sim._pending_sensor_attachments = [first_sensor, second_sensor] + sim._prepared_spawn_topology_revision = -1 + sim._synced_spawn_render_topology_revision = -1 + + with pytest.raises(RuntimeError, match="attach failed"): + sim.prepare() + sim.prepare() + + assert attach_camera_parent.call_args_list == [ + call(first_sensor), + call(second_sensor), + call(second_sensor), + ] + assert sim._pending_sensor_attachments == [] + sync_render_state.assert_called_once_with(result) + + +def test_prepare_syncs_render_state_once_per_topology_revision() -> None: + events: list[str] = [] + result = MagicMock() + result.needs_rebuild = False + result.topology_revision = 3 + spawn_scene = MagicMock() + spawn_scene.builder.is_finalized = True + spawn_scene.builder.result = result + spawn_scene.builder.has_pending_changes = False + spawn_scene.bind.side_effect = lambda: events.append("bind") + sync_render_state = MagicMock(side_effect=lambda _result: events.append("sync")) + + sim = object.__new__(SimulationManager) + sim.physics = SimpleNamespace( + name="newton", + sync_render_state=sync_render_state, + ) + sim.device = torch.device("cpu") + sim._world = MagicMock() + sim._spawn_scene = spawn_scene + sim._pending_sensor_attachments = [] + sim._prepared_spawn_topology_revision = -1 + sim._synced_spawn_render_topology_revision = -1 + + sim.prepare() + sim.prepare() + result.topology_revision = 4 + sim.prepare() + sim.prepare() + + spawn_scene.commit.assert_not_called() + assert spawn_scene.bind.call_count == 4 + assert sync_render_state.call_count == 2 + sync_render_state.assert_has_calls([call(result), call(result)]) + sim._world.update.assert_not_called() + assert events == ["bind", "sync", "bind", "bind", "sync", "bind"] + + +def test_prepare_retries_render_state_sync_without_recommit() -> None: + result = MagicMock() + result.needs_rebuild = False + result.topology_revision = 3 + spawn_scene = MagicMock() + spawn_scene.builder.is_finalized = True + spawn_scene.builder.result = result + spawn_scene.builder.has_pending_changes = False + sync_render_state = MagicMock( + side_effect=[RuntimeError("sync failed"), None], + ) + + sim = object.__new__(SimulationManager) + sim.physics = SimpleNamespace( + name="newton", + sync_render_state=sync_render_state, + ) + sim.device = torch.device("cpu") + sim._world = MagicMock() + sim._spawn_scene = spawn_scene + sim._pending_sensor_attachments = [] + sim._prepared_spawn_topology_revision = -1 + sim._synced_spawn_render_topology_revision = -1 + + with pytest.raises(RuntimeError, match="sync failed"): + sim.prepare() + sim.prepare() + sim.prepare() + + spawn_scene.commit.assert_not_called() + assert spawn_scene.bind.call_count == 3 + assert sync_render_state.call_count == 2 + sim._world.update.assert_not_called() + + +def test_add_camera_uses_owning_manager_render_context( + monkeypatch: pytest.MonkeyPatch, +) -> None: + world = object() + arenas = [object(), object()] + + sim = object.__new__(SimulationManager) + sim.sim_config = SimpleNamespace(num_envs=len(arenas)) + sim.device = torch.device("cpu") + sim._world = world + sim._arenas = arenas + sim._sensors = {} + sim._pending_sensor_attachments = [] + sim._visualization_topology_revision = 0 + sim.SUPPORTED_SENSOR_TYPES = {"Camera": sim_manager_module.Camera} + + monkeypatch.setattr( + sim_manager_module.Camera, + "_build_sensor_from_config", + lambda self, config, device: None, + ) + monkeypatch.setattr(sim_manager_module.Camera, "reset", lambda self: None) + + sensor = sim.add_sensor(CameraCfg(uid="owned_camera")) + + assert sensor._world is world + assert sensor._arenas == arenas + assert sensor.num_instances == len(arenas) + + +def test_camera_attachment_uses_resolved_nodes_and_tracks_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + entities = [MagicMock(), MagicMock()] + parent_nodes = [object(), object()] + owner = MagicMock() + owner.num_envs = len(entities) + owner.get_world.return_value = object() + owner.get_env.side_effect = [object(), object()] + + def build_camera(sensor, config, device) -> None: + sensor._entities[:] = entities + + monkeypatch.setattr( + sim_manager_module.Camera, + "_build_sensor_from_config", + build_camera, + ) + monkeypatch.setattr(sim_manager_module.Camera, "reset", lambda self: None) + + sensor = sim_manager_module.Camera( + CameraCfg( + uid="attached_camera", + extrinsics=CameraCfg.ExtrinsicsCfg(parent="robot/tool"), + ), + owner=owner, + ) + + assert sensor.is_attached is False + sensor.attach_to_parent_nodes(parent_nodes) + + assert sensor.is_attached is True + for entity, parent_node in zip(entities, parent_nodes, strict=True): + entity.attach_node.assert_called_once_with(parent_node) + + +def test_manager_resolves_camera_parent_before_attachment() -> None: + parent_nodes = [object(), object()] + sensor = MagicMock() + sensor.cfg.extrinsics.parent = "robot/tool" + + sim = object.__new__(SimulationManager) + sim._resolve_spawn_sensor_parent_nodes = MagicMock(return_value=parent_nodes) + + sim._attach_camera_parent(sensor) + + sim._resolve_spawn_sensor_parent_nodes.assert_called_once_with("robot/tool") + sensor.attach_to_parent_nodes.assert_called_once_with(parent_nodes) def test_remove_asset_marks_visualization_topology_dirty() -> None: sim, runtime = _make_visualization_sim_manager() rigid_object = MagicMock() + spawn_scene = MagicMock() + spawn_scene.__contains__.return_value = True + spawn_scene.result = object() + sim._spawn_scene = spawn_scene + sim.prepare = MagicMock() sim._rigid_objects = {"cube": rigid_object} + sim._rigid_object_groups = {} + sim._deformable_objects = {} + sim._articulations = {} + sim._robots = {} + sim._lights = {} + sim._sensors = {} assert sim.remove_asset("cube") - rigid_object.destroy.assert_called_once_with() + spawn_scene.remove.assert_called_once_with("cube") + sim.prepare.assert_called_once_with() + rigid_object.destroy.assert_not_called() + assert "cube" not in sim._rigid_objects assert sim._visualization_topology_revision == 3 sim.stop_visualization() assert runtime.stopped @@ -570,6 +968,7 @@ def test_add_stereo_camera_marks_visualization_topology_dirty() -> None: sim.SUPPORTED_SENSOR_TYPES = { "StereoCamera": lambda cfg, device: sensor, } + sim.prepare = MagicMock() cfg = SimpleNamespace(sensor_type="StereoCamera", uid="cam_high") assert sim.add_sensor(cfg) is sensor @@ -669,7 +1068,7 @@ def fake_save_window_record_worker( assert sim._window_record_save_threads == [] -def test_reset_objects_state_includes_soft_and_cloth_assets() -> None: +def test_reset_objects_state_includes_deformable_assets() -> None: sim = object.__new__(SimulationManager) sim._robots = {} sim._articulations = {} @@ -677,10 +1076,12 @@ def test_reset_objects_state_includes_soft_and_cloth_assets() -> None: sim._rigid_object_groups = {} sim._lights = {} sim._sensors = {} - sim._soft_objects = {"soft": MagicMock()} - sim._cloth_objects = {"cloth": MagicMock()} + sim._deformable_objects = { + "soft": MagicMock(), + "cloth": MagicMock(), + } sim.reset_objects_state(env_ids=[1]) - sim._soft_objects["soft"].reset.assert_called_once_with([1]) - sim._cloth_objects["cloth"].reset.assert_called_once_with([1]) + sim._deformable_objects["soft"].reset.assert_called_once_with([1]) + sim._deformable_objects["cloth"].reset.assert_called_once_with([1]) diff --git a/tests/sim/test_sim_manager_cfg.py b/tests/sim/test_sim_manager_cfg.py new file mode 100644 index 000000000..0bee187d7 --- /dev/null +++ b/tests/sim/test_sim_manager_cfg.py @@ -0,0 +1,377 @@ +# ---------------------------------------------------------------------------- +# 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 contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from embodichain.lab.sim import SimulationManager, SimulationManagerCfg +from embodichain.lab.sim.cfg import ( + DefaultPhysicsCfg, + NewtonPhysicsCfg, + WindowCameraPoseCfg, +) +from embodichain.lab.sim.physics import NewtonPhysicsBackend +from embodichain.lab.sim.physics import newton as newton_physics +from embodichain.lab.sim import sim_manager + + +def test_simulation_manager_cfg_uses_default_physics_cfg() -> None: + cfg = SimulationManagerCfg() + + assert type(cfg.physics_cfg) is DefaultPhysicsCfg + + +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_simulation_manager_cfg_initializes_window_camera_pose() -> None: + window_camera_pose = WindowCameraPoseCfg( + enable_hotkey=False, + convert_to_look_at=False, + ) + + cfg = SimulationManagerCfg(window_camera_pose=window_camera_pose) + + assert cfg.window_camera_pose == window_camera_pose + + +def test_simulation_manager_cfg_has_no_scene_construction_switch() -> None: + cfg = SimulationManagerCfg() + + assert "scene_construction" not in cfg.to_dict() + with pytest.raises(TypeError, match="scene_construction"): + SimulationManagerCfg(scene_construction="legacy") + + +def test_newton_physics_cfg_uses_device() -> None: + cfg = NewtonPhysicsCfg(device="cuda:1") + + serialized = cfg.to_dict() + assert serialized["device"] == "cuda:1" + assert serialized["physics_dt"] == 1.0 / 100.0 + assert "solver_type" not in serialized + + +@pytest.mark.no_sim +def test_newton_physics_cfg_preserves_dexsim_auto_solver_default() -> None: + from dexsim.engine.newton_physics import AutoSolverCfg + + cfg = NewtonPhysicsCfg() + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + + assert isinstance(dexsim_cfg.solver_cfg, AutoSolverCfg) + assert dexsim_cfg.solver_cfg.solver_type == "auto" + + +@pytest.mark.no_sim +def test_newton_physics_cfg_requires_dexsim_auto_solver_api( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from dexsim.engine import newton_physics + + monkeypatch.delattr(newton_physics, "AutoSolverCfg") + + with pytest.raises( + ImportError, + match="AutoSolverCfg.*dexsim_engine build pinned by EmbodiChain", + ): + NewtonPhysicsCfg().to_dexsim_cfg(gpu_id=0) + + +@pytest.mark.no_sim +def test_newton_gradient_mode_rejects_auto_solver() -> None: + cfg = NewtonPhysicsCfg(requires_grad=True) + + with pytest.raises(RuntimeError, match="explicit.*semi_implicit"): + cfg.to_dexsim_cfg(gpu_id=0) + + +def test_newton_physics_cfg_passes_warp_log_suppression() -> None: + cfg = NewtonPhysicsCfg(suppress_warp_kernel_logs=False) + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + + assert dexsim_cfg.suppress_warp_kernel_logs is False + + +@pytest.mark.parametrize( + ("physics_cfg", "expect_suppressed"), + [ + (NewtonPhysicsCfg(), True), + (NewtonPhysicsCfg(suppress_warp_kernel_logs=False), False), + (DefaultPhysicsCfg(), False), + ], +) +def test_warp_runtime_init_honors_newton_log_suppression( + monkeypatch: pytest.MonkeyPatch, + physics_cfg, + expect_suppressed: bool, +) -> None: + previous_log_level = sim_manager.wp.config.log_level + observed_log_levels = [] + + def fake_init() -> None: + observed_log_levels.append(sim_manager.wp.config.log_level) + + monkeypatch.setattr(sim_manager.wp, "init", fake_init) + try: + sim_manager._initialize_warp_runtime(physics_cfg) + expected_log_level = ( + sim_manager.wp.LOG_WARNING if expect_suppressed else previous_log_level + ) + assert observed_log_levels == [expected_log_level] + assert sim_manager.wp.config.log_level == previous_log_level + finally: + sim_manager.wp.config.log_level = previous_log_level + + +def test_newton_warp_log_suppression_covers_world_update() -> None: + previous_log_level = sim_manager.wp.config.log_level + observed_log_levels = [] + + class NoopProfiler: + def section(self, *_args, **_kwargs): + return nullcontext() + + class World: + def update(self, _physics_dt: float) -> None: + observed_log_levels.append(sim_manager.wp.config.log_level) + + manager = SimpleNamespace( + profiler=NoopProfiler(), + prepare=lambda: None, + is_physics_manually_update=True, + sim_config=SimpleNamespace( + physics_dt=0.01, + physics_cfg=NewtonPhysicsCfg(), + visualization=SimpleNamespace(backend="none"), + ), + update_gizmos=lambda: None, + _world=World(), + _visualization_sim_step=0, + _visualization_sim_time=0.0, + _window_record_state=None, + ) + try: + SimulationManager.update(manager, physics_dt=0.01) + assert observed_log_levels == [sim_manager.wp.LOG_WARNING] + assert sim_manager.wp.config.log_level == previous_log_level + finally: + sim_manager.wp.config.log_level = previous_log_level + + +@pytest.mark.no_sim +def test_newton_backend_exposes_resolved_solver_type() -> None: + backend = NewtonPhysicsBackend(SimpleNamespace()) + world_config = SimpleNamespace(newton_cfg=None) + sim_config = SimulationManagerCfg( + physics_cfg=NewtonPhysicsCfg( + device="cpu", + solver_cfg={"solver_type": "xpbd"}, + ), + ) + + backend.configure_world(world_config, sim_config) + + assert backend.solver_type == "xpbd" + assert world_config.newton_cfg.solver_cfg.solver_type == "xpbd" + + +@pytest.mark.no_sim +def test_newton_backend_reports_scene_resolved_auto_solver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + world = object() + native_backend = SimpleNamespace(solver_type="mujoco_warp") + manager = SimpleNamespace(_world=world) + backend = NewtonPhysicsBackend(manager) + world_config = SimpleNamespace(newton_cfg=None) + sim_config = SimulationManagerCfg(physics_cfg=NewtonPhysicsCfg()) + monkeypatch.setattr( + "dexsim.engine.newton_physics.backend_registry.get_newton_backend", + lambda candidate: native_backend if candidate is world else None, + ) + + backend.configure_world(world_config, sim_config) + + assert world_config.newton_cfg.solver_cfg.solver_type == "auto" + assert backend.solver_type == "mujoco_warp" + + +def test_newton_teardown_releases_render_views_on_the_resolved_device( + monkeypatch: pytest.MonkeyPatch, +) -> None: + render_sync = MagicMock() + newton_backend = SimpleNamespace(render_sync=render_sync) + manager = SimpleNamespace(_world=object()) + backend = NewtonPhysicsBackend(manager) + world_config = SimpleNamespace(newton_cfg=None) + synchronize_device = MagicMock() + sim_config = SimulationManagerCfg( + gpu_id=2, + physics_cfg=NewtonPhysicsCfg(device="cuda"), + ) + monkeypatch.setattr( + newton_physics.wp, + "synchronize_device", + synchronize_device, + ) + from dexsim.engine.newton_physics import backend_registry + + monkeypatch.setattr( + backend_registry, + "get_newton_backend", + lambda world: newton_backend if world is manager._world else None, + ) + + backend.configure_world(world_config, sim_config) + backend.prepare_for_teardown() + + synchronize_device.assert_called_once_with("cuda:2") + render_sync.clear.assert_called_once_with() + + +def test_newton_teardown_skips_cpu_devices( + monkeypatch: pytest.MonkeyPatch, +) -> None: + render_sync = MagicMock() + newton_backend = SimpleNamespace(render_sync=render_sync) + manager = SimpleNamespace(_world=object()) + backend = NewtonPhysicsBackend(manager) + world_config = SimpleNamespace(newton_cfg=None) + synchronize_device = MagicMock() + sim_config = SimulationManagerCfg(physics_cfg=NewtonPhysicsCfg(device="cpu")) + monkeypatch.setattr( + newton_physics.wp, + "synchronize_device", + synchronize_device, + ) + from dexsim.engine.newton_physics import backend_registry + + monkeypatch.setattr( + backend_registry, + "get_newton_backend", + lambda world: newton_backend if world is manager._world else None, + ) + + backend.configure_world(world_config, sim_config) + backend.prepare_for_teardown() + + synchronize_device.assert_not_called() + render_sync.clear.assert_called_once_with() + + +def test_newton_backend_syncs_render_state_without_physics_step( + monkeypatch: pytest.MonkeyPatch, +) -> None: + world = object() + native_backend = SimpleNamespace( + sync_to_dexsim=MagicMock(), + sync_particle_fluids=MagicMock(), + ) + monkeypatch.setattr( + "dexsim.engine.newton_physics.backend_registry.get_newton_backend", + lambda candidate: native_backend if candidate is world else None, + ) + backend = NewtonPhysicsBackend(SimpleNamespace()) + + backend.sync_render_state(SimpleNamespace(world=world)) + + native_backend.sync_to_dexsim.assert_called_once_with(world) + native_backend.sync_particle_fluids.assert_called_once_with(world) + + +@pytest.mark.no_sim +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, + "enable_multiccd": True, + }, + ) + + 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 + assert dexsim_cfg.solver_cfg.enable_multiccd is True + + +@pytest.mark.no_sim +def test_newton_physics_cfg_accepts_explicit_auto_solver_mapping() -> None: + from dexsim.engine.newton_physics import AutoSolverCfg + + cfg = NewtonPhysicsCfg(solver_cfg={"class_type": "AutoSolverCfg"}) + + dexsim_cfg = cfg.to_dexsim_cfg(gpu_id=0) + + assert isinstance(dexsim_cfg.solver_cfg, AutoSolverCfg) + + +@pytest.mark.no_sim +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/test_sim_profiler.py b/tests/sim/test_sim_profiler.py index bcc46a165..d197435ac 100644 --- a/tests/sim/test_sim_profiler.py +++ b/tests/sim/test_sim_profiler.py @@ -22,6 +22,7 @@ import torch from embodichain.lab.sim import Profiler, ProfilerCfg, SimulationManager +from embodichain.lab.sim.cfg import DefaultPhysicsCfg pytestmark = pytest.mark.no_sim @@ -52,8 +53,10 @@ def _make_sim_update_probe(profiler: Profiler) -> SimulationManager: sim._visualization_runtime = None sim._visualization_sim_step = 0 sim._visualization_sim_time = 0.0 + sim.prepare = lambda: None sim.sim_config = types.SimpleNamespace( physics_dt=0.01, + physics_cfg=DefaultPhysicsCfg(), visualization=types.SimpleNamespace(backend="none"), ) return sim diff --git a/tests/sim/workspace/test_analyzer.py b/tests/sim/workspace/test_analyzer.py index 57e94b61b..170d19c9e 100644 --- a/tests/sim/workspace/test_analyzer.py +++ b/tests/sim/workspace/test_analyzer.py @@ -35,7 +35,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) @@ -77,6 +77,7 @@ def setup_simulation(self): } self.robot: Robot = self.sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + self.sim.prepare() def teardown_method(self): """Clean up resources after each test method.""" diff --git a/tests/sim/workspace/test_cache.py b/tests/sim/workspace/test_cache.py index 688f33712..29f22b34b 100644 --- a/tests/sim/workspace/test_cache.py +++ b/tests/sim/workspace/test_cache.py @@ -508,7 +508,7 @@ def _robot_ns(**overrides) -> argparse.Namespace: init_pos=[0.0, 0.0, 0.0], init_rot=[0.0, 0.0, 0.0], fix_base=True, - use_usd_properties=False, + asset_physics_mode="overlay", ) defaults.update(overrides) return argparse.Namespace(**defaults) @@ -551,6 +551,7 @@ def test_build_robot_cfg_urdf_defaults_solver_urdf(): assert cfg.control_parts == {"arm": ["fr3_joint[1-7]"]} assert cfg.solver_cfg["arm"].end_link_name == "fr3_hand_tcp" assert cfg.solver_cfg["arm"].urdf_path == "/tmp/panda.urdf" + assert cfg.asset_physics_mode == "overlay" def test_build_robot_cfg_usd_requires_urdf(): @@ -573,6 +574,15 @@ def test_build_robot_cfg_usd_with_urdf(): assert cfg.solver_cfg["arm"].urdf_path == "/tmp/robot.urdf" +def test_build_robot_cfg_accepts_source_independent_preserve_mode(): + """The asset physics policy applies to either USD or URDF sources.""" + from embodichain.lab.scripts.analyze_workspace import build_robot_cfg + + cfg, _part, _urdf = build_robot_cfg(_robot_ns(asset_physics_mode="preserve")) + + assert cfg.asset_physics_mode == "preserve" + + def test_build_robot_cfg_asset_requires_ee_link(): """--asset without --ee-link raises a clear error.""" from embodichain.lab.scripts.analyze_workspace import build_robot_cfg @@ -722,6 +732,7 @@ def _make_cobotmagic_sim(tmp_path): }, } robot = sim.add_robot(cfg=CobotMagicCfg.from_dict(cfg_dict)) + sim.prepare() return sim, robot diff --git a/tests/sim/workspace/test_sim_utils.py b/tests/sim/workspace/test_sim_utils.py new file mode 100644 index 000000000..9d266993f --- /dev/null +++ b/tests/sim/workspace/test_sim_utils.py @@ -0,0 +1,101 @@ +# ---------------------------------------------------------------------------- +# 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 pytest + +from embodichain.lab.sim.cfg import RigidObjectCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg +from embodichain.lab.sim.utility.sim_utils import _load_rigid_mesh_prototype + + +class _FakeActor: + def add_rigidbody(self, *args, **kwargs) -> None: + pass + + +class _FakeArena: + def __init__(self) -> None: + self.load_actor_called = False + self.acd_method: str | None = None + + def load_actor(self, *args, **kwargs) -> _FakeActor: + self.load_actor_called = True + return _FakeActor() + + def load_actor_with_acd(self, *args, method: str, **kwargs) -> _FakeActor: + self.acd_method = method + return _FakeActor() + + +def test_load_rigid_mesh_uses_shape_collision_defaults() -> None: + arena = _FakeArena() + cfg = RigidObjectCfg(uid="mesh", shape=MeshCfg(fpath="mesh.obj")) + + _load_rigid_mesh_prototype( + arena, + cfg, + cache_dir=None, + body_type=None, + is_newton_backend=False, + ) + + assert arena.load_actor_called + + +def test_load_rigid_mesh_forwards_shape_acd_method() -> None: + arena = _FakeArena() + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.obj", + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=2, + acd_method="vhacd", + ), + ), + ) + + _load_rigid_mesh_prototype( + arena, + cfg, + cache_dir=None, + body_type=None, + is_newton_backend=False, + ) + + assert arena.acd_method == "vhacd" + + +def test_load_rigid_mesh_rejects_dynamic_triangle_mesh_collision() -> None: + cfg = RigidObjectCfg( + uid="mesh", + shape=MeshCfg( + fpath="mesh.obj", + collision=MeshCollisionCfg(approximation="triangle_mesh"), + ), + ) + + with pytest.raises(ValueError, match="only for static"): + _load_rigid_mesh_prototype( + _FakeArena(), + cfg, + cache_dir=None, + body_type=None, + is_newton_backend=False, + ) diff --git a/tests/test_release_metadata.py b/tests/test_release_metadata.py index 53600115b..71187fb3a 100644 --- a/tests/test_release_metadata.py +++ b/tests/test_release_metadata.py @@ -16,13 +16,17 @@ from __future__ import annotations -import tomllib from pathlib import Path from zipfile import ZipFile import pytest from packaging.requirements import Requirement +try: + import tomllib +except ModuleNotFoundError: # Python 3.10 + import tomli as tomllib + from scripts.validate_wheel_metadata import WheelMetadataError, validate_wheel from setup import get_package_dir, get_packages diff --git a/tests/toolkits/test_grasp_pose_generator.py b/tests/toolkits/test_grasp_pose_generator.py index ed09a6714..2799e631c 100644 --- a/tests/toolkits/test_grasp_pose_generator.py +++ b/tests/toolkits/test_grasp_pose_generator.py @@ -31,7 +31,7 @@ from embodichain.lab.sim.objects import Robot, RigidObject from embodichain.lab.sim.utility.action_utils import interpolate_with_distance from embodichain.toolkits.graspkit import ParallelJawGripperModelCfg -from embodichain.lab.sim.shapes import MeshCfg +from embodichain.lab.sim.shapes import MeshCfg, MeshCollisionCfg from embodichain.lab.sim.solvers import PytorchSolverCfg from embodichain.data import get_data_path from embodichain.lab.gym.utils.gym_utils import add_env_launcher_args_to_parser @@ -41,7 +41,7 @@ JointDrivePropertiesCfg, RobotCfg, LightCfg, - RigidBodyAttributesCfg, + RigidBodyPhysicsCfg, RigidObjectCfg, URDFCfg, ) @@ -62,7 +62,7 @@ def initialize_simulation() -> SimulationManager: """ config = SimulationManagerCfg( headless=True, - sim_device=torch.device("cuda"), + device=torch.device("cuda"), render_cfg=RenderCfg(renderer="auto"), physics_dt=1.0 / 100.0, arena_space=2.5, @@ -103,7 +103,7 @@ def create_robot(sim: SimulationManager, position=[0.0, 0.0, 0.0]) -> Robot: {"component_type": "hand", "urdf_path": gripper_urdf_path}, ] ), - drive_pros=JointDrivePropertiesCfg( + joint_drive_props=JointDrivePropertiesCfg( stiffness={"Joint[0-9]": 1e4, "FINGER[1-2]": 1e3}, damping={"Joint[0-9]": 1e3, "FINGER[1-2]": 1e2}, max_effort={"Joint[0-9]": 1e5, "FINGER[1-2]": 1e4}, @@ -136,13 +136,17 @@ def create_mug(sim: SimulationManager): uid="table", shape=MeshCfg( fpath=get_data_path("CoffeeCup/cup.ply"), + collision=MeshCollisionCfg( + approximation="convex_decomposition", + max_hulls=16, + ), ), - attrs=RigidBodyAttributesCfg( - mass=0.01, - dynamic_friction=0.97, - static_friction=0.99, + attrs=RigidBodyPhysicsCfg.from_dict( + { + "mass_props": {"mass": 0.01}, + "material_props": {"dynamic_friction": 0.97, "static_friction": 0.99}, + } ), - max_convex_hull_num=16, init_pos=[0.55, 0.0, 0.01], init_rot=[0.0, 0.0, -90], body_scale=(4, 4, 4), @@ -207,6 +211,7 @@ def test_grasp_pose_generator(): try: robot = create_robot(sim, position=[0.0, 0.0, 0.0]) mug = create_mug(sim) + sim.prepare() # get mug grasp pose grasp_generator = AntipodalGraspPoseGenerator( diff --git a/tests/utils/test_configclass.py b/tests/utils/test_configclass.py new file mode 100644 index 000000000..d0a5b8749 --- /dev/null +++ b/tests/utils/test_configclass.py @@ -0,0 +1,41 @@ +# ---------------------------------------------------------------------------- +# 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 configclass decorator.""" + +from __future__ import annotations + +from dataclasses import fields +from typing import ClassVar + +from embodichain.utils import configclass + + +@configclass +class _DeferredClassVarCfg: + values: list[int] = [] + label: ClassVar[str] = "shared" + + +def test_deferred_classvar_is_not_converted_to_a_dataclass_field() -> None: + first = _DeferredClassVarCfg() + second = _DeferredClassVarCfg() + first.values.append(1) + + assert [item.name for item in fields(_DeferredClassVarCfg)] == ["values"] + assert first.to_dict() == {"values": [1]} + assert second.values == [] + assert _DeferredClassVarCfg.label == "shared" diff --git a/tests/utils/test_math.py b/tests/utils/test_math.py new file mode 100644 index 000000000..fab263b1c --- /dev/null +++ b/tests/utils/test_math.py @@ -0,0 +1,94 @@ +# ---------------------------------------------------------------------------- +# 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.utils.math import ( + convert_quat, + default_orientation, + matrix_from_quat, + quat_apply, + quat_conjugate, + quat_from_matrix, + quat_mul, + trans_matrix_to_xyz_quat, + xyz_quat_to_4x4_matrix, +) + + +def _distinct_xyzw() -> torch.Tensor: + """Return a normalized quaternion whose components expose order mistakes.""" + quaternion = torch.tensor([[1.0, 2.0, 3.0, 4.0]], dtype=torch.float32) + return quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True) + + +def test_quaternion_matrix_round_trip_uses_xyzw() -> None: + quaternion = _distinct_xyzw() + + rotation = matrix_from_quat(quaternion) + restored = quat_from_matrix(rotation) + + torch.testing.assert_close(restored, quaternion, atol=1.0e-6, rtol=1.0e-6) + + +def test_quaternion_product_and_conjugate_return_xyzw_identity() -> None: + quaternion = _distinct_xyzw() + + product = quat_mul(quaternion, quat_conjugate(quaternion)) + + torch.testing.assert_close( + product, + torch.tensor([[0.0, 0.0, 0.0, 1.0]]), + atol=1.0e-6, + rtol=1.0e-6, + ) + + +def test_quaternion_application_reads_scalar_from_last_component() -> None: + half_sqrt_two = 2.0**-0.5 + z_quarter_turn_xyzw = torch.tensor( + [[0.0, 0.0, half_sqrt_two, half_sqrt_two]], dtype=torch.float32 + ) + + rotated = quat_apply(z_quarter_turn_xyzw, torch.tensor([[1.0, 0.0, 0.0]])) + + torch.testing.assert_close( + rotated, + torch.tensor([[0.0, 1.0, 0.0]]), + atol=1.0e-6, + rtol=1.0e-6, + ) + + +def test_pose_vector_round_trip_uses_xyz_plus_xyzw() -> None: + pose = torch.cat((torch.tensor([[0.25, -0.5, 0.75]]), _distinct_xyzw()), dim=-1) + + restored = trans_matrix_to_xyz_quat(xyz_quat_to_4x4_matrix(pose)) + + torch.testing.assert_close(restored, pose, atol=1.0e-6, rtol=1.0e-6) + + +def test_identity_and_boundary_conversion_orders_are_explicit() -> None: + xyzw = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + + torch.testing.assert_close( + default_orientation(1, "cpu"), torch.tensor([[0.0, 0.0, 0.0, 1.0]]) + ) + torch.testing.assert_close( + convert_quat(xyzw, to="wxyz"), torch.tensor([[4.0, 1.0, 2.0, 3.0]]) + ) diff --git a/tests/visualization/test_protocol.py b/tests/visualization/test_protocol.py index 5c5ca4d3c..19a161d90 100644 --- a/tests/visualization/test_protocol.py +++ b/tests/visualization/test_protocol.py @@ -35,18 +35,21 @@ ) -def test_pose_conversion_preserves_embodichain_wxyz_order() -> None: - pose = np.array([1.0, 2.0, 3.0, 2.0, 0.0, 0.0, 0.0], dtype=np.float32) +def test_pose_conversion_converts_embodichain_xyzw_to_protocol_wxyz() -> None: + pose = np.array([1.0, 2.0, 3.0, 1.0, 2.0, 3.0, 4.0], dtype=np.float32) position, wxyz = pose_to_position_wxyz(pose) np.testing.assert_allclose(position, [1.0, 2.0, 3.0]) - np.testing.assert_allclose(wxyz, [1.0, 0.0, 0.0, 0.0]) + np.testing.assert_allclose( + wxyz, + np.array([4.0, 1.0, 2.0, 3.0]) / np.sqrt(30.0), + ) def test_pose_conversion_accepts_batch_of_four_pose_vectors() -> None: poses = np.tile( - np.array([[0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]], dtype=np.float32), + np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]], dtype=np.float32), (4, 1), ) diff --git a/tests/visualization/test_scene_exporter.py b/tests/visualization/test_scene_exporter.py index ef94f03d5..dbf7945c5 100644 --- a/tests/visualization/test_scene_exporter.py +++ b/tests/visualization/test_scene_exporter.py @@ -164,7 +164,8 @@ def get_local_pose(self, to_matrix: bool = False) -> np.ndarray: class _DeformableObject: - def __init__(self) -> None: + def __init__(self, deformable_type: str) -> None: + self.deformable_type = deformable_type local_vertices = np.array( [[0.0, 0.0, 0.0], [0.15, 0.0, 0.0], [0.0, 0.15, 0.0]], dtype=np.float32, @@ -177,16 +178,10 @@ def __init__(self) -> None: ) self._faces = np.array([[0, 1, 2]], dtype=np.int32) - def get_current_collision_vertices(self) -> np.ndarray: - return self.vertices - - def get_current_vertex_position(self) -> np.ndarray: + def get_surface_vertices(self) -> np.ndarray: return self.vertices - def get_collision_surface_triangles(self, env_ids: list[int]) -> np.ndarray: - return self.get_triangles(env_ids) - - def get_triangles(self, env_ids: list[int]) -> np.ndarray: + def get_surface_triangles(self, env_ids: list[int]) -> np.ndarray: return np.stack([self._faces for _ in env_ids]) @@ -224,17 +219,11 @@ def get_rigid_object_group_uid_list(self) -> list[str]: def get_rigid_object_group(self, uid: str) -> None: raise AssertionError(f"Unexpected rigid-object-group lookup: {uid}") - def get_soft_object_uid_list(self) -> list[str]: + def get_deformable_object_uid_list(self) -> list[str]: return [] - def get_soft_object(self, uid: str) -> None: - raise AssertionError(f"Unexpected soft-object lookup: {uid}") - - def get_cloth_object_uid_list(self) -> list[str]: - return [] - - def get_cloth_object(self, uid: str) -> None: - raise AssertionError(f"Unexpected cloth-object lookup: {uid}") + def get_deformable_object(self, uid: str) -> None: + raise AssertionError(f"Unexpected deformable-object lookup: {uid}") def get_sensor_uid_list(self) -> list[str]: return [] @@ -265,17 +254,11 @@ def get_articulation_uid_list(self) -> list[str]: def get_articulation(self, uid: str) -> None: raise AssertionError(f"Unexpected articulation lookup: {uid}") - def get_soft_object_uid_list(self) -> list[str]: - return [] - - def get_soft_object(self, uid: str) -> None: - raise AssertionError(f"Unexpected soft-object lookup: {uid}") - - def get_cloth_object_uid_list(self) -> list[str]: + def get_deformable_object_uid_list(self) -> list[str]: return [] - def get_cloth_object(self, uid: str) -> None: - raise AssertionError(f"Unexpected cloth-object lookup: {uid}") + def get_deformable_object(self, uid: str) -> None: + raise AssertionError(f"Unexpected deformable-object lookup: {uid}") def get_sensor_uid_list(self) -> list[str]: return [] @@ -453,8 +436,8 @@ class _CompleteSimulation(_Simulation): def __init__(self) -> None: super().__init__() self.rigid_group = _RigidObjectGroup() - self.soft = _DeformableObject() - self.cloth = _DeformableObject() + self.soft = _DeformableObject("volume") + self.cloth = _DeformableObject("surface") def get_rigid_object_group_uid_list(self) -> list[str]: return ["pair"] @@ -463,19 +446,11 @@ def get_rigid_object_group(self, uid: str) -> _RigidObjectGroup: assert uid == "pair" return self.rigid_group - def get_soft_object_uid_list(self) -> list[str]: - return ["jelly"] - - def get_soft_object(self, uid: str) -> _DeformableObject: - assert uid == "jelly" - return self.soft - - def get_cloth_object_uid_list(self) -> list[str]: - return ["flag"] + def get_deformable_object_uid_list(self) -> list[str]: + return ["jelly", "flag"] - def get_cloth_object(self, uid: str) -> _DeformableObject: - assert uid == "flag" - return self.cloth + def get_deformable_object(self, uid: str) -> _DeformableObject: + return {"jelly": self.soft, "flag": self.cloth}[uid] def test_manifest_deduplicates_geometry_and_escapes_paths() -> None: